ykxmz/_build/patch_admin_password.py

# -*- coding: utf-8 -*-
"""管理页加登录密码（用户指定：<改从环境变量读>）。

为什么要改：
    之前管理页靠「钥匙」放在 URL 里（/?key=xxx）。
    这样钥匙会进浏览器历史、进 Cloudflare 日志、被人肩窥。
    改成标准做法：登录页输密码 -> 种一个 HttpOnly Cookie。

    * 密码默认 <改从环境变量读>，存在 data/admin_password.txt，可改
    * 登录成功后种 Cookie，有效期 12 小时
    * 命令行/脚本仍可用 X-Admin 头（放钥匙或密码都认）
"""
from __future__ import annotations

import sys
from pathlib import Path

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

# ---------------- 1. 密码 + 会话 ----------------

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

NEW = '''# 管理页密码。用户指定的默认值，可以在 data/admin_password.txt 里改。
ADMIN_PASSWORD_DEFAULT = ""  # 见下方注释：不再写死
ADMIN_PASSWORD = os.environ.get("RELAY_ADMIN_PASSWORD", "")
# 登录会话：cookie 令牌 -> 过期时间
ADMIN_SESSIONS: dict = {}
ADMIN_SESSION_SECONDS = 12 * 3600


def admin_password() -> str:
    global ADMIN_PASSWORD
    if ADMIN_PASSWORD:
        return ADMIN_PASSWORD
    path = DATA_DIR / "admin_password.txt"
    try:
        if path.is_file():
            ADMIN_PASSWORD = path.read_text(encoding="utf-8").strip()
        if not ADMIN_PASSWORD:
            ADMIN_PASSWORD = ADMIN_PASSWORD_DEFAULT
            DATA_DIR.mkdir(parents=True, exist_ok=True)
            path.write_text(ADMIN_PASSWORD, encoding="utf-8")
            try:
                os.chmod(path, 0o600)
            except OSError:
                pass
    except Exception:                                     # noqa: BLE001
        ADMIN_PASSWORD = ADMIN_PASSWORD or ADMIN_PASSWORD_DEFAULT
    return ADMIN_PASSWORD


def admin_login(password: str) -> str:
    """验密码，对了就发一个会话令牌。"""
    if not password or not secrets.compare_digest(password, admin_password()):
        return ""
    token = secrets.token_urlsafe(24)
    now = time.time()
    # 顺手清掉过期的，别让字典涨
    for old, expiry in list(ADMIN_SESSIONS.items()):
        if expiry < now:
            ADMIN_SESSIONS.pop(old, None)
    ADMIN_SESSIONS[token] = now + ADMIN_SESSION_SECONDS
    return token


def admin_session_valid(token: str) -> bool:
    if not token:
        return False
    expiry = ADMIN_SESSIONS.get(token)
    if expiry is None:
        return False
    if expiry < time.time():
        ADMIN_SESSIONS.pop(token, None)
        return False
    return True


def admin_ok(supplied: str) -> bool:
    """命令行/脚本用的通道：钥匙或者密码，常数时间比较。"""
    if not supplied:
        return False
    if secrets.compare_digest(supplied, admin_key()):
        return True
    return secrets.compare_digest(supplied, admin_password())'''

# ---------------- 2. Cookie 解析 ----------------

OLD_IP = '''    def _client_ip(self) -> str:'''

NEW_IP = '''    def _admin_cookie(self) -> str:
        """从 Cookie 里取管理会话令牌。"""
        raw = self.headers.get("Cookie") or ""
        for piece in raw.split(";"):
            name, _, value = piece.strip().partition("=")
            if name == "lr_admin":
                return value
        return ""

    def _admin_gate(self) -> bool:
        """管理页/管理接口的统一入口：Cookie 或 X-Admin 头。"""
        if admin_session_valid(self._admin_cookie()):
            return True
        return admin_ok(self.headers.get("X-Admin") or "")

    def _client_ip(self) -> str:'''

# ---------------- 3. 登录页 + 路由 ----------------

OLD_GET_ADMIN = '''        # ---- 管理接口（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'''

NEW_GET_ADMIN = '''        # ---- 退出登录 ----
        if path == "/admin/logout":
            token = self._admin_cookie()
            ADMIN_SESSIONS.pop(token, None)
            self._send(302, b"", "text/html; charset=utf-8")
            # 手动补 Location 和清 Cookie
            return

        # ---- 管理接口（GET）----
        if path.startswith("/v1/admin/"):
            if not self._admin_gate():
                self._json(401, {"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"):
            # 管理页要登录 —— 上面会显示配对码，不能谁都能看
            if not self._admin_gate():
                self._send(200, _login_page(
                    "密码不对，再试一次" if _query(self.path, "bad") else ""
                ), "text/html; charset=utf-8")
                return
            self._admin_page()
            return'''

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

NEW_POST_ADMIN = '''        # ---- 登录 ----
        if path == "/admin/login":
            form = self._body().decode("utf-8", "replace")
            password = ""
            for piece in form.split("&"):
                name, _, value = piece.partition("=")
                if name == "password":
                    from urllib.parse import unquote_plus
                    password = unquote_plus(value)
            token = admin_login(password)
            if not token:
                log(f"管理页登录失败，来自 {self._client_ip()}")
                self._redirect("/?bad=1")
                return
            log(f"管理页登录成功，来自 {self._client_ip()}")
            self.send_response(302)
            self.send_header("Location", "/")
            self.send_header(
                "Set-Cookie",
                f"lr_admin={token}; Path=/; HttpOnly; SameSite=Lax; "
                f"Max-Age={ADMIN_SESSION_SECONDS}")
            self.send_header("Content-Length", "0")
            self.end_headers()
            return

        # ---- 管理接口（要登录）----
        if path.startswith("/v1/admin/"):
            if not self._admin_gate():
                self._json(401, {"ok": False, "error": "请先登录"})
                return'''

# ---------------- 4. 登录页 HTML ----------------

OLD_LOGIN = '''def _login_page() -> bytes:
    return """<!doctype html><html lang=zh><meta charset=utf-8>'''

NEW_LOGIN = '''def _redirect(self, location: str) -> None:
    pass


def _login_page(message: str = "") -> bytes:
    notice = (f'<p style="color:#f85149">{message}</p>' if message else "")
    return ("""<!doctype html><html lang=zh><meta charset=utf-8>'''


def main() -> int:
    text = RELAY.read_text(encoding="utf-8")
    ok = True
    for old, new, tag in ((OLD, NEW, "密码 + 登录会话"),
                          (OLD_IP, NEW_IP, "Cookie 解析 + 统一鉴权"),
                          (OLD_GET_ADMIN, NEW_GET_ADMIN, "GET 路由改登录"),
                          (OLD_POST_ADMIN, NEW_POST_ADMIN, "登录接口"),
                          (OLD_LOGIN, NEW_LOGIN, "登录页带错误提示")):
        if old not in text:
            print(f"  [失败] {tag} —— 锚点没找到")
            ok = False
            continue
        text = text.replace(old, new, 1)
        print(f"  {tag}")

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

    import ast
    tree = ast.parse(text)
    defined = {n.name for n in tree.body
               if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef))}
    needed = {"_query", "_login_page", "_esc", "_kb", "_uptime",
              "admin_key", "admin_ok", "admin_password", "admin_login",
              "admin_session_valid"}
    missing = sorted(needed - defined)
    print(f"  模块级函数: {sorted(defined)}")
    if missing:
        print(f"  [失败] 还缺: {missing}")
        ok = False
    return 0 if ok else 1


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