ykxmz/_build/add_update_endpoint.py

# -*- coding: utf-8 -*-
r'''给中转加「最新版本查询」接口 —— 被控端静默自动更新的第一步。

== 为什么由中转来回答 ==

被控端已经知道中转地址（那是它的配置），而中转的 downloads/ 目录里
就放着各版本的成品。让中转直接回答"最新是哪版、在哪下、校验和多少"，
就不用把任何地址写死在客户端里 —— 换个域名/换台服务器都不用改代码。

== 新增 ==

    GET /v1/update/latest?platform=windows

    200 {
      "ok": true,
      "platform": "windows",
      "version": "2.1.5",
      "name": "ykxmz-setup-2.1.5.exe",
      "url": "/download/ykxmz-setup-2.1.5.exe",
      "size": 9395078,
      "sha256": "....",
      "notes": "..."          <- 可选，取 CHANGELOG 里那一节
    }

    没找到任何包时：
    200 {"ok": false, "error": "还没有可用的安装包"}

== url 给相对路径 ==

给相对路径（`/download/xxx.exe`），让被控端用**它自己配置的中转地址**
拼成完整 URL。这样：
  * 中转换域名/换 IP 不用改任何东西
  * 走 Cloudflare 还是直连由被控端那边的配置决定

== sha256 要缓存 ==

安装包 9MB，每次请求都算一遍 sha256 大概 30 毫秒 —— 不算慢，
但这个接口会被所有被控端周期性调用（每台每小时一次的话，
一个学校几十台就是几十次），没必要重复算。
按「文件名 + 大小 + 修改时间」做键缓存。
'''
from __future__ import annotations

import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
RELAY = ROOT / "relay" / "relay.py"


def main() -> int:
    raw = RELAY.read_bytes()
    bom = raw.startswith(b"\xef\xbb\xbf")
    text = raw.decode("utf-8-sig")

    if "_latest_package" in text:
        print("  已经加过了")
        return 0

    # ⚠ **先把所有锚点都验一遍，再动手改。**
    #
    # 第一版是把替换都做在内存里、最后才 write_bytes ——
    # 结果中途某个锚点没匹配上就 `return 1`，
    # 前面几处改动**全丢了**（脚本还打印了"✓ 已加"），
    # 而文件其实一个字没变。特别容易误判成"加好了"。
    #
    # 教训：**批处理脚本要幺先验后改，幺写一段落一段。**
    ANCHORS = {
        "import_json": "\nimport json",
        "list_downloads": '    items.sort(key=lambda item: item["mtime"], reverse=True)\n    return items\n',
        "downloads_route": '        if path == "/v1/downloads":',
    }
    missing = [k for k, v in ANCHORS.items() if v not in text]
    if missing:
        print("  ✗ 这些锚点找不到，一个字都不改：%s" % ", ".join(missing))
        return 1
    print("  ✓ 锚点全部就位")

    # ---- 1. import hashlib ----
    if "\nimport hashlib" not in text:
        text = text.replace(ANCHORS["import_json"],
                            "\nimport hashlib\nimport json", 1)
        print("  ✓ 加了 import hashlib")

    # ---- 2. 插到 _list_downloads 之后 ----
    anchor = ANCHORS["list_downloads"]

    helper = anchor + '''

# 最新安装包的信息缓存。键 = (文件名, 大小, 修改时间)。
_PKG_CACHE: dict = {}


def _sha256_of(path: Path) -> str:
    """算文件的 sha256（分块读，别把 9MB 全塞内存）。"""
    h = hashlib.sha256()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(256 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def _latest_package(platform: str = "windows") -> dict:
    """找 downloads/ 里最新的安装包，给被控端自更新用。

    `platform` 现在只有 windows；留着参数是为了以后加别的平台
    （手表/服务端不通过这个接口自更新，它们各有各的更新方式）。

    返回 {} 表示"还没有可用的包"。
    """
    # 各平台的成品名前缀
    prefix = {
        "windows": "ykxmz-setup-",
    }.get(platform)
    if not prefix:
        return {}

    best = None      # (版本元组, 路径)
    try:
        if not DOWNLOADS_DIR.is_dir():
            return {}
        for entry in DOWNLOADS_DIR.iterdir():
            if not entry.is_file():
                continue
            name = entry.name
            if not name.startswith(prefix) or not name.endswith(".exe"):
                continue
            if not _safe_download_name(name):
                continue
            # 从文件名里抠版本号
            mid = name[len(prefix):-len(".exe")]
            parts = mid.split(".")
            if len(parts) != 3 or not all(p.isdigit() for p in parts):
                continue
            ver = tuple(int(p) for p in parts)
            # ⚠ 按**数字段**比大小，不能按字符串 ——
            #   否则 "1.9.0" > "2.1.4"（字符 '9' > '1'）。
            if best is None or ver > best[0]:
                best = (ver, entry)
    except OSError:
        return {}

    if best is None:
        return {}

    ver, path = best
    try:
        st = path.stat()
    except OSError:
        return {}

    key = (path.name, st.st_size, int(st.st_mtime))
    cached = _PKG_CACHE.get(key)
    if cached:
        digest = cached
    else:
        try:
            digest = _sha256_of(path)
        except OSError:
            return {}
        # 只留最近几个，别让字典涨
        _PKG_CACHE.clear()
        _PKG_CACHE[key] = digest

    return {
        "platform": platform,
        "version": ".".join(str(p) for p in ver),
        "name": path.name,
        "url": "/download/" + quote(path.name),
        "size": st.st_size,
        "sha256": digest,
        "mtime": st.st_mtime,
    }
'''
    text = text.replace(anchor, helper, 1)
    print("  ✓ 加了 _latest_package / _sha256_of")

    # ---- 3. 路由 ----
    route_anchor = ANCHORS["downloads_route"]

    route = '''        if path == "/v1/update/latest":
            # 被控端自更新用：最新版是哪版、在哪下、校验和多少。
            # 公开接口 —— 文件名和哈希本来就在下载页上人人可见，
            # 不含任何敏感信息。
            platform = (urllib.parse.parse_qs(
                urllib.parse.urlparse(self.path).query)
                .get("platform", ["windows"])[0] or "windows")
            pkg = _latest_package(platform)
            if not pkg:
                self._json(200, {"ok": False,
                                 "error": "还没有可用的安装包"})
            else:
                self._json(200, dict(pkg, ok=True))
            return

''' + route_anchor
    text = text.replace(route_anchor, route, 1)
    print("  ✓ 加了 /v1/update/latest 路由")

    RELAY.write_bytes((b"\xef\xbb\xbf" if bom else b"") + text.encode("utf-8"))
    print("  ✓ 写回 relay.py")
    return 0


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