# -*- coding: utf-8 -*-
"""修防火墙判定里的两个漏报，以及 Store 写盘失败会崩的问题。

1. rule_exists：netsh 返回空输出但退出码为 0 时会返回 True ——
   明明没放行却显示"已放行"。应该要求输出里真的出现规则名。

2. list_block_rules：只认"有 Program 字段且 basename 是 lanremote-agent.exe
   /python.exe"的阻止规则。**按端口封的阻止规则（没有 Program 字段）
   完全看不见** —— 而"UDP 能搜到、TCP 连不上"正是这种规则造成的。
   这个模块存在的意义就是解决那个场景，结果它自己漏掉了。

3. Store：装到 Program Files 这类只读目录时，mkdir/写文件抛 PermissionError
   会让程序**启动就崩**，而且没有任何提示。改成失败时回退到
   %LOCALAPPDATA%，并记住实际可写的位置。
"""
from __future__ import annotations

import sys
from pathlib import Path

ROOT = Path(r"C:\Users\shu\Desktop\deepseekagent\LanRemote\windows"
            r"\lanremote")
FIREWALL = ROOT / "firewall.py"
CONFIG = ROOT / "config.py"

# ---------------- firewall ----------------

FIXES = [
    (FIREWALL,
     '''    if proc.returncode != 0:
        return False
    text = ((proc.stdout or "") + (proc.stderr or "")).lower()
    for marker in _NO_MATCH_MARKERS:
        if marker in text:
            return False
    return True''',
     '''    if proc.returncode != 0:
        return False
    text = ((proc.stdout or "") + (proc.stderr or "")).lower()
    for marker in _NO_MATCH_MARKERS:
        if marker in text:
            return False
    # 输出为空（或明显不是规则详情）时**不能**当成"规则存在" ——
    # 本地化提示语没收录、netsh 被拦、超时都可能给出空输出，
    # 那时候报"已放行"就是在骗用户。
    if len(text.strip()) < 10:
        return False
    if rule_name.lower() not in text:
        return False
    return True''',
     "rule_exists 空输出不再误判"),

    (FIREWALL,
     '''def list_block_rules() -> list[dict]:
    """列出针对本程序（及 python）的入站阻止规则。"""''',
     '''def _port_hits(info: dict, ports: tuple) -> bool:
    """这条规则的 LocalPort 是否命中我们关心的端口。"""
    raw = (info.get("LocalPort") or "").strip()
    if not raw or raw.lower() in ("any", "任意"):
        return False
    pieces = [p.strip() for p in raw.replace("，", ",").split(",")]
    for piece in pieces:
        if not piece:
            continue
        for port in ports:
            if not port:
                continue
            if piece == str(port):
                return True
            # 端口段，如 8000-9000
            if "-" in piece:
                try:
                    low, high = piece.split("-", 1)
                    if int(low) <= int(port) <= int(high):
                        return True
                except ValueError:
                    continue
    return False


def list_block_rules(tcp_port: int = 0, udp_port: int = 0) -> list[dict]:
    """列出可能挡住我们的入站阻止规则。

    两类都算：
      * 程序级：Program 指向 lanremote-agent.exe / python.exe
      * 端口级：LocalPort 命中我们的 TCP/UDP 端口

    第二类原来完全没覆盖。而"UDP 能搜到、TCP 连不上"往往就是
    一条按端口封的规则 —— 组策略、第三方安全软件、手工建规则都会这样。
    """''',
     "list_block_rules 覆盖端口级规则"),

    (FIREWALL,
     '''        program = (info.get("Program") or "").strip('"')
        base = os.path.basename(program).lower() if program else ""
        if base in PROGRAM_RULE_NAMES:
            info["ProgramBase"] = base
            blocks.append(info)
    return blocks''',
     '''        program = (info.get("Program") or "").strip('"')
        base = os.path.basename(program).lower() if program else ""
        if base in PROGRAM_RULE_NAMES:
            info["ProgramBase"] = base
            info["Kind"] = "program"
            blocks.append(info)
            continue
        if _port_hits(info, (tcp_port, udp_port)):
            info["Kind"] = "port"
            blocks.append(info)
    return blocks''',
     "list_block_rules 收集端口级"),

    (FIREWALL,
     "    all_blocked = list_block_rules()",
     "    all_blocked = list_block_rules(tcp_port, udp_port)",
     "check 传端口"),
]

# ---------------- config ----------------

CONFIG_FIX = (CONFIG,
              '''    def __init__(self, directory: Path):
        self.dir = Path(directory)
        self.dir.mkdir(parents=True, exist_ok=True)''',
              '''    def __init__(self, directory: Path):
        self.dir = Path(directory)
        try:
            self.dir.mkdir(parents=True, exist_ok=True)
            probe = self.dir / ".write_test"
            probe.write_text("ok", encoding="utf-8")
            probe.unlink()
        except OSError:
            # 装到 Program Files 这类只读目录时，原来会直接崩在启动。
            # 退到用户自己的 AppData —— 那里一定可写。
            fallback = Path(os.environ.get("LOCALAPPDATA",
                                           str(Path.home()))) / "LanRemote"
            try:
                fallback.mkdir(parents=True, exist_ok=True)
            except OSError:
                pass
            self.dir = fallback''',
              "Store 只读目录时回退")


def main() -> int:
    applied, skipped = 0, []
    for path, old, new, tag in FIXES + [CONFIG_FIX]:
        text = path.read_text(encoding="utf-8")
        if old not in text:
            skipped.append(tag)
            print(f"  [跳过] {tag} —— 锚点没找到")
            continue
        path.write_text(text.replace(old, new, 1), encoding="utf-8")
        applied += 1
        print(f"  {tag}: OK")

    # config 需要 os
    text = CONFIG.read_text(encoding="utf-8")
    if "\nimport os" not in text and "import os" not in text.split("\n\n")[0]:
        if "from __future__" in text:
            text = text.replace("from __future__ import annotations",
                                "from __future__ import annotations\n\nimport os", 1)
        else:
            text = "import os\n" + text
        CONFIG.write_text(text, encoding="utf-8")
        print("  config.py / import os: OK")

    print(f"\n  应用 {applied} / {len(FIXES) + 1}")
    if skipped:
        print(f"  跳过: {', '.join(skipped)}")

    import ast
    ok = True
    for label, path in (("firewall.py", FIREWALL), ("config.py", CONFIG)):
        try:
            ast.parse(path.read_text(encoding="utf-8"))
            print(f"  {label}: 语法 OK")
        except SyntaxError as exc:
            print(f"  {label}: 语法错误 {exc}")
            ok = False
    return 0 if ok else 1


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