ykxmz/_build/fix_relay_helpers.py

# -*- coding: utf-8 -*-
"""修两个漏掉的地方。

1. _login_page / _query 没插进去
   上一版补丁的判断条件写错了：我用
       if "_login_page" not in text.split("def _esc")[0]
   来判断"是不是已经定义过"，但 _login_page() 在 do_GET 里被**调用**，
   那段就在 def _esc 前面 —— 所以判断成了"已存在"，函数压根没插。
   结果 do_GET 一跑就 NameError，relay 每次请求都崩（Cloudflare 报 502）。

2. /v1/admin/devices 只有 POST 路由
   测试和管理脚本习惯用 GET，补上。

教训：插入前应该检查的是「有没有**定义**」，不是「有没有出现」。
"""
from __future__ import annotations

import sys
from pathlib import Path

RELAY = Path(r"C:\Users\shu\Desktop\deepseekagent\LanRemote\relay\relay.py")

HELPERS = '''def _query(url: str, name: str) -> str:
    """从 URL 里取一个查询参数（不引 urllib，省内存）。"""
    if "?" not in url:
        return ""
    try:
        from urllib.parse import parse_qs
        values = parse_qs(url.split("?", 1)[1]).get(name) or []
        return values[0] if values else ""
    except Exception:                                     # noqa: BLE001
        return ""


def _login_page() -> bytes:
    return """<!doctype html><html lang=zh><meta charset=utf-8>
<meta name=viewport content="width=device-width,initial-scale=1">
<title>中转服务器 · 需要钥匙</title>
<style>
 body{font:14px/1.6 -apple-system,'Microsoft YaHei',sans-serif;
      background:#0f1117;color:#e6edf3;display:flex;align-items:center;
      justify-content:center;height:100vh;margin:0}
 .box{background:#161b22;border:1px solid #30363d;border-radius:12px;
      padding:26px 30px;max-width:420px}
 h1{font-size:17px;margin:0 0 10px} p{color:#8b949e;font-size:13px}
 code{background:#21262d;padding:2px 7px;border-radius:4px;
      display:inline-block;margin-top:6px;word-break:break-all}
</style>
<div class=box>
<h1>需要管理钥匙</h1>
<p>钥匙在随身 WiFi 上：<br><code>/root/lanremote-relay/data/admin_key.txt</code></p>
<p>打开方式：<br><code>/?key=你的钥匙</code></p>
</div></html>""".encode("utf-8")


'''

ANCHOR = "class Handler(BaseHTTPRequestHandler):"

# do_GET 里补管理接口
OLD_GET = '''        if path in ("/", "/admin"):
            # 管理页要钥匙 —— 上面会显示配对码，不能谁都能看
            supplied = _query(self.path, "key") or ""
            if not admin_ok(supplied):
                self._send(403, _login_page(), "text/html; charset=utf-8")
                return
            self._admin_page()
            return'''

NEW_GET = '''        # ---- 管理接口（GET，也要钥匙）----
        if path.startswith("/v1/admin/"):
            supplied = (self.headers.get("X-Admin")
                        or _query(self.path, "key") or "")
            if not admin_ok(supplied):
                self._json(403, {"ok": False, "error": "管理钥匙不对"})
                return
            if path == "/v1/admin/devices":
                self._json(200, RELAY.snapshot())
                return
            self._json(404, {"ok": False, "error": "没有这个管理接口"})
            return

        if path in ("/", "/admin"):
            # 管理页要钥匙 —— 上面会显示配对码，不能谁都能看
            supplied = _query(self.path, "key") or ""
            if not admin_ok(supplied):
                self._send(403, _login_page(), "text/html; charset=utf-8")
                return
            self._admin_page()
            return'''


def main() -> int:
    text = RELAY.read_text(encoding="utf-8")

    # ---- 1. 插函数（这次按「有没有定义」判断）----
    if "\ndef _query(" in text:
        print("  _query 已经定义过")
    elif ANCHOR not in text:
        print("  [失败] 找不到 class Handler")
        return 1
    else:
        text = text.replace(ANCHOR, HELPERS + ANCHOR, 1)
        print("  _query / _login_page 已插入")

    # ---- 1b. 如果之前误插在别处，清掉重复 ----
    occurrences = text.count("def _login_page(")
    if occurrences > 1:
        print(f"  [注意] _login_page 出现了 {occurrences} 次，清理重复")
        # 保留第一份（在 class Handler 之前的那份）
        head, sep, tail = text.partition(ANCHOR)
        # 从头里删掉第二份（如果有）
        first = head.find("def _login_page(")
        second = head.find("def _login_page(", first + 1)
        if second > 0:
            # 找到这个函数块的结果
            end = head.find("\n\n\n", second)
            if end < 0:
                end = len(head)
            head = head[:second] + head[end + 3:] if end < len(head) else head[:second]
            text = head + sep + tail
            print("  已删掉重复的一份")

    # ---- 2. GET 管理接口 ----
    if OLD_GET not in text:
        print("  [失败] 找不到 do_GET 的管理页分支")
        return 1
    text = text.replace(OLD_GET, NEW_GET, 1)
    print("  do_GET 已支持 /v1/admin/devices")

    RELAY.write_text(text, encoding="utf-8")

    # ---- 3. 静态检查：用到的模块级函数都得有定义 ----
    import ast
    tree = ast.parse(text)
    defined = set()
    for node in tree.body:
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            defined.add(node.name)
    needed = {"_query", "_login_page", "_esc", "_kb", "_uptime"}
    missing = sorted(needed - defined)
    print(f"  模块级函数: {sorted(defined)}")
    if missing:
        print(f"  [失败] 还缺: {missing}")
        return 1
    print("  所有用到的模块级函数都有定义")

    for occurrence in ("def _login_page(", "def _query("):
        print(f"  {occurrence.rstrip('(')} 出现 {text.count(occurrence)} 次")
    return 0


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