ykxmz/_build/patch_relay_admin.py

# -*- coding: utf-8 -*-
"""给中转服务器加管理接口。

用户要「收发服务器端可以管理」，所以要能：
    * 看所有设备 / 在线状态 / 流量
    * 踢掉某个设备的会话
    * 抹掉一个配对码绑定（让那台设备下次能设新码）
    * 管理页要凭钥匙访问（默认谁都能看配对码不合适）

钥匙没配就自动生成一个，存在 data/admin_key.txt（权限 600）。
"""
from __future__ import annotations

import sys
from pathlib import Path

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

# ---------------- 1. 钥匙 ----------------

OLD_KEY = 'ADMIN_KEY = os.environ.get("RELAY_ADMIN_KEY", "")\n'

NEW_KEY = '''ADMIN_KEY = os.environ.get("RELAY_ADMIN_KEY", "")


def admin_key() -> str:
    """管理页 / 管理接口的钥匙。

    没配置就自动生成一个并记到 data/admin_key.txt。
    以前默认是空字符串，等于管理页谁都能打开 ——
    那上面能看到配对码，太不合适了。
    """
    global ADMIN_KEY
    if ADMIN_KEY:
        return ADMIN_KEY
    path = DATA_DIR / "admin_key.txt"
    try:
        if path.is_file():
            ADMIN_KEY = path.read_text(encoding="utf-8").strip()
        if not ADMIN_KEY:
            ADMIN_KEY = secrets.token_urlsafe(18)
            DATA_DIR.mkdir(parents=True, exist_ok=True)
            path.write_text(ADMIN_KEY, encoding="utf-8")
            try:
                os.chmod(path, 0o600)
            except OSError:
                pass
    except Exception:                                     # noqa: BLE001
        ADMIN_KEY = ADMIN_KEY or secrets.token_urlsafe(18)
    return ADMIN_KEY


def admin_ok(supplied: str) -> bool:
    """常数时间比较，避免时序侧信道。"""
    if not supplied:
        return False
    return secrets.compare_digest(supplied, admin_key())
'''

# ---------------- 2. 管理接口 ----------------

OLD_ROUTE = '''        if path == "/v1/client/detach":'''

NEW_ROUTE = '''        # ---- 管理接口（要钥匙）----
        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

            if path == "/v1/admin/kick":
                device_id = str(self._json_body().get("deviceId") or "")
                if not device_id:
                    self._json(400, {"ok": False, "error": "要给 deviceId"})
                    return
                gone = RELAY.kick(device_id)
                log(f"管理操作：踢掉 {device_id[:12]}… 的 {gone} 个会话")
                self._json(200, {"ok": True, "kicked": gone})
                return

            if path == "/v1/admin/forget":
                code = str(self._json_body().get("code") or "").strip()
                if not CODE_RE.match(code):
                    self._json(400, {"ok": False, "error": "配对码要是 6 位数字"})
                    return
                done = RELAY.forget(code)
                log(f"管理操作：抹掉配对码 {code} -> {done}")
                self._json(200, {"ok": True, "forgotten": done})
                return

            self._json(404, {"ok": False, "error": "没有这个管理接口"})
            return

        if path == "/v1/client/detach":'''

# ---------------- 3. 管理页要钥匙 ----------------

OLD_PAGE = '''        if path in ("/", "/admin"):
            self._admin_page()
            return'''

NEW_PAGE = '''        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'''

LOGIN_PAGE = '''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:400px}}
 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")


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


'''


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

    for old, new, tag in ((OLD_KEY, NEW_KEY, "管理钥匙（自动生成+常数时间比较）"),
                          (OLD_ROUTE, NEW_ROUTE, "管理接口 kick/forget/devices"),
                          (OLD_PAGE, NEW_PAGE, "管理页加钥匙校验")):
        if old not in text:
            print(f"  [失败] {tag} —— 锚点没找到")
            ok = False
            continue
        text = text.replace(old, new, 1)
        print(f"  {tag}")

    # _login_page / _query 放在 _esc 前面
    anchor = "def _esc(text) -> str:"
    if "_login_page" not in text.split("def _esc")[0] and anchor in text:
        text = text.replace(anchor, LOGIN_PAGE + anchor, 1)
        print("  _login_page / _query 已插入")

    # 启动时把钥匙打出来，方便用户找
    old_boot = '''    if not ADMIN_KEY:
        log("提示：未设 RELAY_ADMIN_KEY，管理页任何人都能看（只有配对码和流量）")'''
    new_boot = '''    key = admin_key()
    log(f"管理页：http://<本机>:{PORT}/?key={key}")
    log(f"（钥匙也在 {DATA_DIR / 'admin_key.txt'}）")'''
    if old_boot in text:
        text = text.replace(old_boot, new_boot, 1)
        print("  启动日志会打印管理页地址")

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

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


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