# -*- coding: utf-8 -*-
"""数据目录（配对码 / 设备令牌）永远不给拉。

端到端测试发现：EXE 装在 Desktop 下时，`data/settings.json` 确实落在
"允许目录"范围内 —— 不是绕过白名单，但里面有配对码和全部设备令牌，
被拉走等于把控制权交出去。所以加一条硬性排除，跟白名单并行生效。
"""
from __future__ import annotations

import sys
from pathlib import Path

ROOT = Path(r"C:\Users\shu\Desktop\deepseekagent\LanRemote\windows"
            r"\lanremote")
SERVER = ROOT / "server.py"
FEATURES = ROOT / "features.py"

# ---------------- features.py：加敏感路径判断 ----------------
FEAT_OLD = '''    def assert_not_root(self, path: Path) -> None:'''

FEAT_NEW = '''    def sensitive_dirs(self) -> list:
        """程序自己的数据目录 —— 里面有配对码和设备令牌，任何情况下都不给碰。"""
        out = []
        try:
            out.append(Path(self.store.dir))
        except Exception:                                 # noqa: BLE001
            pass
        try:
            from .config import app_dir
            out.append(Path(app_dir()))
        except Exception:                                 # noqa: BLE001
            pass
        return out

    def is_sensitive(self, path: Path) -> bool:
        try:
            resolved = path.resolve()
        except OSError:
            return False
        for folder in self.sensitive_dirs():
            try:
                folder = folder.resolve()
            except OSError:
                continue
            if resolved == folder or folder in resolved.parents:
                return True
        return False

    def assert_not_root(self, path: Path) -> None:'''

# 在 resolve_in_roots 里也挡一道
RESOLVE_OLD = '''        for root in self.allowed_roots():
            try:
                resolved.relative_to(root.resolve())
                return resolved
            except ValueError:
                continue
        raise ValueError(f"只允许访问自己的目录（{resolved} 不在允许范围内）")'''

RESOLVE_NEW = '''        if self.is_sensitive(resolved):
            raise ValueError("这是程序自己的数据目录，不允许访问")

        for root in self.allowed_roots():
            try:
                resolved.relative_to(root.resolve())
                return resolved
            except ValueError:
                continue
        raise ValueError(f"只允许访问自己的目录（{resolved} 不在允许范围内）")'''

# ---------------- server.py：拉文件时也挡 ----------------
SRV_OLD = '''            for root in allowed:
                try:
                    resolved.relative_to(root.resolve())
                    return resolved
                except (ValueError, OSError):
                    continue
            self._log(f"[文件] 拒绝拉取允许范围外的路径: {raw}")
            return None'''

SRV_NEW = '''            if self._is_sensitive(resolved):
                self._log(f"[文件] 拒绝拉取数据目录: {raw}")
                return None
            for root in allowed:
                try:
                    resolved.relative_to(root.resolve())
                    return resolved
                except (ValueError, OSError):
                    continue
            self._log(f"[文件] 拒绝拉取允许范围外的路径: {raw}")
            return None'''

SRV_REL_OLD = '''        for root in allowed_roots:
            try:
                relative.relative_to(root.resolve())
                return relative if relative.is_file() else None
            except (ValueError, OSError):
                continue
        self._log(f"[文件] 拒绝拉取相对路径: {raw}")
        return None'''

SRV_REL_NEW = '''        if self._is_sensitive(relative):
            self._log(f"[文件] 拒绝拉取数据目录: {raw}")
            return None
        for root in allowed_roots:
            try:
                relative.relative_to(root.resolve())
                return relative if relative.is_file() else None
            except (ValueError, OSError):
                continue
        self._log(f"[文件] 拒绝拉取相对路径: {raw}")
        return None'''

# server 的辅助方法
SRV_HELPER_OLD = '''    def _resolve_received(self, path: str) -> Path | None:'''

SRV_HELPER_NEW = '''    def _is_sensitive(self, path: Path) -> bool:
        """程序自己的数据目录：配对码、设备令牌都在这儿。"""
        if self.features is not None:
            try:
                return self.features.is_sensitive(path)
            except Exception:                             # noqa: BLE001
                pass
        try:
            folder = Path(self.store.dir).resolve()
            resolved = path.resolve()
            return resolved == folder or folder in resolved.parents
        except OSError:
            return False

    def _resolve_received(self, path: str) -> Path | None:'''

# 去掉末尾那个没有白名单的兜底
FALLBACK_OLD = '''        # 退一步：允许传相对路径直接命中
        if raw.is_file():
            return raw
        return None'''
FALLBACK_NEW = '''        return None'''

FIXES = [
    (FEATURES, FEAT_OLD, FEAT_NEW, "features: 敏感目录判断"),
    (FEATURES, RESOLVE_OLD, RESOLVE_NEW, "features: resolve_in_roots 挡数据目录"),
    (SERVER, SRV_HELPER_OLD, SRV_HELPER_NEW, "server: _is_sensitive"),
    (SERVER, SRV_OLD, SRV_NEW, "server: 绝对路径挡数据目录"),
    (SERVER, SRV_REL_OLD, SRV_REL_NEW, "server: 相对路径挡数据目录"),
    (SERVER, FALLBACK_OLD, FALLBACK_NEW, "server: 删掉无白名单的兜底"),
]


def main() -> int:
    applied, skipped = 0, []
    for path, old, new, tag in FIXES:
        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")

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

    import ast
    ok = True
    for label, path in (("server.py", SERVER), ("features.py", FEATURES)):
        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())