#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""E11：CH395F 服务器侧 KeepAlive —— "链路在、对端静默"判死验证（PC 侧注入）。

背景（E11）：手册 §5.56 规定 TCP 服务器应在 SINT_STAT_CONNECT 后开启
KeepAlive；本仓该调用 2026-10-12 被撤回，全仓仅剩客户端分支启用，而
App/ 无任何 net_connect() → 产品上 keepalive 实际关闭。本脚本用于验证
"开启后能否在有界时间内判死静默对端，且不误杀正常空闲会话"。

用法：
  # 场景「正常空闲存活」：连接后静默 hold 秒，末尾发 KEEPLIVE 验证会话仍在
  python e11_keepalive_probe.py --mode idle --hold-s 90

  # 场景「静默对端」：连接后立即装 Windows 防火墙入站 DROP（仅本连接本地端口 +
  #    设备 IP）→ 芯片 keepalive 探测得不到 ACK → 期望设备在
  #    IDLE + INTVL×CNT 内判死（本仓 60 + 5x3 = 75s）
  python e11_keepalive_probe.py --mode block --hold-s 120

判据（设备侧事件由串口日志核对，本脚本打印 PC 侧时间线）：
  期望：串口出现 `[ERR][NET] sockN TIMEOUT`（耗时≈75s），随后
           `sockN pump/relisten` 恢复；本脚本移除规则后新连接必须成功。
  期望：hold 期间**无** TIMEOUT；末尾 KEEPLIVE 必须收到 0xBB 回显。

需要管理员权限（仅 block 模式装/卸防火墙规则；脚本自行清理）。
"""

import argparse
import socket
import subprocess
import sys
import time

TAG_C, TAG_D = 0xAA, 0xBB
CMD_KEEPLIVE = 0
RULE_NAME = "SELFHEAL-E11-DROP"


def _crc(data: bytes) -> int:
    c = 0
    for b in data:
        c ^= b
    return c


def build_frame(cmd: int, data: bytes = b"") -> bytes:
    n = len(data)
    return bytes([TAG_C, cmd, (n >> 8) & 0xFF, n & 0xFF]) + data + bytes([_crc(data)])


def recv_frame(sock: socket.socket, timeout: float = 3.0):
    sock.settimeout(timeout)
    hdr = b""
    while len(hdr) < 4:
        chunk = sock.recv(4 - len(hdr))
        if not chunk:
            return None
        hdr += chunk
    tag, cmd, n = hdr[0], hdr[1], (hdr[2] << 8) | hdr[3]
    body = b""
    while len(body) < n + 1:
        chunk = sock.recv(n + 1 - len(body))
        if not chunk:
            return None
        body += chunk
    data, crcb = body[:n], body[n]
    return cmd, data, (tag == TAG_D) and (crcb == _crc(data))


def firewall(action: str, local_port: int, remote_ip: str) -> bool:
    verb = "add" if action == "add" else "delete"
    cmd = ["netsh", "advfirewall", "firewall", verb, "rule",
           f"name={RULE_NAME}", "dir=in", "action=block", "protocol=TCP",
           f"localport={local_port}", f"remoteip={remote_ip}"]
    if verb == "delete":
        cmd = ["netsh", "advfirewall", "firewall", "delete", "rule",
               f"name={RULE_NAME}"]
    r = subprocess.run(cmd, capture_output=True, text=True)
    if r.returncode != 0:
        print(f"[WARN] 防火墙 {verb} 失败：{r.stdout.strip()} {r.stderr.strip()}")
        return False
    return True


def connect(host: str, port: int, timeout: float = 5.0):
    s = socket.create_connection((host, port), timeout)
    return s


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--host", default="192.168.1.100")
    ap.add_argument("--port", type=int, default=7003)
    ap.add_argument("--mode", choices=["idle", "block"], required=True)
    ap.add_argument("--hold-s", type=float, default=90.0)
    ap.add_argument("--expect-detect-s", type=float, default=75.0,
                    help="判死上界（IDLE + INTVL*CNT），仅用于提示")
    args = ap.parse_args()

    t0 = time.time()

    def log(msg: str) -> None:
        print(f"[{time.strftime('%H:%M:%S')} +{time.time() - t0:6.1f}s] {msg}",
              flush=True)

    log(f"模式={args.mode} 目标={args.host}:{args.port} hold={args.hold_s}s")
    log(f"（设备侧判死上界参考 = {args.expect_detect_s}s）")

    try:
        s = connect(args.host, args.port)
    except OSError as e:
        print(f"[FAIL] 连接失败：{e}")
        return 2
    s.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 0)  # PC 侧不参与
    lport = s.getsockname()[1]
    log(f"已连接（本地端口 {lport}）→ 进入静默")

    rc = 0
    if args.mode == "block":
        # 入站 DROP：设备发来的 keepalive 探测包在 PC 侧被丢弃 → 永不回 ACK
        if not firewall("add", lport, args.host):
            s.close()
            return 2
        log(f"已装防火墙入站 DROP（本地端口 {lport} <- {args.host}）→ 对端静默开始")

    try:
        time.sleep(args.hold_s)
    except KeyboardInterrupt:
        pass

    if args.mode == "idle":
        # 存活探针：KEEPLIVE 必须被回显
        log("静默结束 → 发 KEEPLIVE 验证会话存活")
        try:
            s.sendall(build_frame(CMD_KEEPLIVE))
            r = recv_frame(s, 3.0)
        except OSError as e:
            r = None
            log(f"KEEPLIVE 发送/接收异常：{e}")
        if r and r[2]:
            log("PASS：收到 0xBB KEEPLIVE 回显 → 正常空闲会话未被误杀")
        else:
            log(f"FAIL：未收到有效回显（{r}）→ 空闲置的会话被断开（误杀）")
            rc = 1
        s.close()
        return rc

    # block 模式：卸规则 → 验证旧连接已死 + 监听器已恢复
    log("静默结束 → 卸载防火墙规则")
    firewall("delete", lport, args.host)
    time.sleep(1.0)

    # 旧连接是否已被设备判死（应读到 EOF/RST；读不到也算进程侧正常）
    try:
        s.settimeout(2.0)
        d = s.recv(16)
        log(f"旧连接 recv → {d!r}（b'' = 已收到 FIN，即设备已判死并释放）")
    except OSError as e:
        log(f"旧连接 recv 异常（设备已判死/连接被拆）：{e}")
    s.close()

    # 监听器恢复验证：新连接 + KEEPLIVE 回显
    log("验证监听器恢复：新建连接 + KEEPLIVE")
    ok = False
    for attempt in range(1, 6):
        try:
            s2 = connect(args.host, args.port)
            s2.sendall(build_frame(CMD_KEEPLIVE))
            r = recv_frame(s2, 3.0)
            s2.close()
            if r and r[2]:
                log(f"PASS：第 {attempt} 次新连接建立并收到回显 → 监听器已恢复")
                ok = True
                break
        except OSError as e:
            log(f"第 {attempt} 次新连接失败：{e}")
        time.sleep(2.0)
    if not ok:
        log("FAIL：监听器未恢复（端口仍不接客）")
        rc = 1
    return rc


if __name__ == "__main__":
    sys.exit(main())