# -*- coding: utf-8 -*-
"""配对限速改成按来源 IP。

之前是全局计数：局域网里任何人发 10 个错码，就会把**合法用户**也锁在门外
5 分钟 —— 这本身是个拒绝服务的口子。

改成按来源 IP 分别计数：攻击者只能锁住自己。
"""
from __future__ import annotations

import sys
from pathlib import Path

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

FIXES = [
    ('''        # 防暴力枚举：6 位数字只有 100 万种，不限速的话局域网里慢慢试就行。
        now = time.time()
        self._pair_failures = [t for t in getattr(self, "_pair_failures", [])
                               if now - t < 300]
        if len(self._pair_failures) >= 10:
            wait = max(1, int(300 - (now - self._pair_failures[0])))
            self._log(f"[配对] 失败次数过多，锁定中（{wait} 秒）")
            return {"ok": False,
                    "error": f"配对码错误次数过多，请等 {wait} 秒后再试"}
        # 用 compare_digest，避免按字符比较带来的时间侧信道
        if not secrets.compare_digest(code, expected):
            self._pair_failures.append(now)
            self._log(f"[配对失败] 来自 {session.addr[0]}，配对码不匹配"
                      f"（{len(self._pair_failures)}/10）")
            return {"ok": False, "error": "配对码错误"}
        self._pair_failures = []''',
     '''        # 防暴力枚举：6 位数字只有 100 万种，不限速的话局域网里慢慢试就行。
        # 按**来源 IP**分别计数 —— 全局计数的话，别人发 10 个错码
        # 就能把你这个合法用户也锁在门外，那本身就是个拒绝服务的口子。
        now = time.time()
        ip = session.addr[0]
        buckets = getattr(self, "_pair_failures", None)
        if not isinstance(buckets, dict):
            buckets = {}
        mine = [t for t in buckets.get(ip, []) if now - t < 300]
        buckets = {k: [t for t in v if now - t < 300]
                   for k, v in buckets.items() if any(now - t < 300 for t in v)}
        buckets[ip] = mine
        self._pair_failures = buckets

        if len(mine) >= 10:
            wait = max(1, int(300 - (now - mine[0])))
            self._log(f"[配对] {ip} 失败次数过多，锁定中（{wait} 秒）")
            return {"ok": False,
                    "error": f"配对码错误次数过多，请等 {wait} 秒后再试"}
        # 用 compare_digest，避免按字符比较带来的时间侧信道
        if not secrets.compare_digest(code, expected):
            mine.append(now)
            self._log(f"[配对失败] 来自 {ip}，配对码不匹配（{len(mine)}/10）")
            return {"ok": False, "error": "配对码错误"}
        self._pair_failures[ip] = []'''),

    ("        # 最近的配对失败时间戳，用于限速防枚举\n"
     "        self._pair_failures: list[float] = []",
     "        # 配对失败记录：{来源 IP: [时间戳]}，用于限速防枚举\n"
     "        self._pair_failures: dict = {}"),
]


def main() -> int:
    text = SERVER.read_text(encoding="utf-8")
    ok = True
    for old, new in FIXES:
        if old not in text:
            print(f"  [跳过] 锚点没找到: {old.strip().splitlines()[0][:50]}")
            ok = False
            continue
        text = text.replace(old, new, 1)
        print("  一处已改")
    if ok:
        SERVER.write_text(text, encoding="utf-8")

    import ast
    try:
        ast.parse(SERVER.read_text(encoding="utf-8"))
        print("  server.py: 语法 OK")
    except SyntaxError as exc:
        print(f"  server.py: 语法错误 {exc}")
        ok = False
    return 0 if ok else 1


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