ykxmz/_build/wire_update_cmd.py

# -*- coding: utf-8 -*-
r'''把「静默自动更新」接进协议和服务端。

== 加两条命令 ==

    {"cmd":"update","action":"check"}   -> 查最新版，不改动任何东西
    {"cmd":"update","action":"start"}   -> 真的要更新（下载+静默安装+重启）

== 为什么"查"和"做"分成两步 ==

手表点一下按钮，先看到"有新版本 2.1.5（当前 2.1.4，约 9 MB）"，
再决定要不要升。直接一步到位的话，用户不知道自己触发了什么。

== start 为什么要开线程 ==

安装包 9 MB，在学校的网上下载可能要几十秒。
如果同步下载，手表那边的请求会一直挂着；而挂太久用户以为卡死了。

所以：立刻回一句"开始更新了"，然后在后台线程里下载 + 校验 + 应用。
进度通过 server 自己的日志广播给正在看的手表。

== 应用之后进程必须退出 ==

Windows 上正在跑的 exe 是锁着的，安装器覆盖不了。
所以 updater 会先起一个**分离的 .cmd 帮手**（不受本进程退出影响），
帮手等本进程退出 -> 静默装 -> 重新拉起。
这里的顺序不能反。
'''
from __future__ import annotations

import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
PROTO = ROOT / "windows" / "lanremote" / "protocol.py"
SERVER = ROOT / "windows" / "lanremote" / "server.py"


def patch_protocol() -> bool:
    raw = PROTO.read_bytes()
    bom = raw.startswith(b"\xef\xbb\xbf")
    text = raw.decode("utf-8-sig")
    if "CMD_UPDATE" in text:
        print("  protocol.py 已经加过了")
        return True
    anchor = 'CMD_INFO = "info"'
    if anchor not in text:
        print("  ✗ protocol.py 找不到 CMD_INFO")
        return False
    text = text.replace(
        anchor,
        anchor + "\n"
        '# 静默自动更新：从**用户自己的中转站**取新版本包，'
        '静默装好并重启。\n'
        '# action=check 只看不动；action=start 真的更新。\n'
        'CMD_UPDATE = "update"', 1)
    PROTO.write_bytes((b"\xef\xbb\xbf" if bom else b"") + text.encode("utf-8"))
    print("  ✓ protocol.py 加了 CMD_UPDATE")
    return True


def patch_server() -> bool:
    raw = SERVER.read_bytes()
    bom = raw.startswith(b"\xef\xbb\xbf")
    text = raw.decode("utf-8-sig")
    if "_do_update" in text:
        print("  server.py 已经加过了")
        return True

    anchors = {
        "info": 'if cmd == P.CMD_INFO:\n            return {"ok": True, "info": self.actions.info()}',
        "unknown": '        return {"ok": False, "error": f"未知命令: {cmd}"}',
        "imports": "from . import protocol as P",
    }
    missing = [k for k, v in anchors.items() if v not in text]
    if missing:
        print("  ✗ server.py 这些锚点找不到：%s" % ", ".join(missing))
        return False
    print("  ✓ server.py 锚点全部就位")

    # ---- import updater ----
    text = text.replace(
        anchors["imports"],
        anchors["imports"] + "\nfrom . import updater as U", 1)

    # ---- 命令分发 ----
    text = text.replace(
        anchors["info"],
        anchors["info"] + '''

        if cmd == P.CMD_UPDATE:
            return self._do_update(msg)''', 1)

    # ---- 实现 ----
    handler = '''

    # ---------------- 静默自动更新 ----------------

    def _do_update(self, msg: dict) -> dict:
        """处理 {"cmd":"update","action":...}。

        取包的地方是**用户自己配的中转站**（store 里的 relay_url），
        不是任何写死的公网地址 —— 这一点是刻意的，见 updater.py 的说明。
        """
        from . import __version__ as my_version

        action = str(msg.get("action") or "check").lower()
        relay = (self.store.relay_url() or "").strip()

        if action == "check":
            try:
                info = U.check(relay, my_version)
            except U.UpdateError as exc:
                return {"ok": False, "error": str(exc)}
            if not info.get("supported"):
                return {"ok": False,
                        "error": info.get("error") or "不支持自动更新"}
            if "error" in info:
                return {"ok": False, "error": info["error"]}
            return {
                "ok": True,
                "current": my_version,
                "latest": info["latest"],
                "needUpdate": bool(info["needUpdate"]),
                "size": info.get("size", 0),
                "name": info.get("name", ""),
            }

        if action != "start":
            return {"ok": False, "error": "action 只能是 check 或 start"}

        # ---- 真的更新 ----
        # 先查一遍（在调用方线程里查，失败能立刻报回去）
        try:
            info = U.check(relay, my_version)
        except U.UpdateError as exc:
            return {"ok": False, "error": str(exc)}
        if not info.get("supported") or "error" in info:
            return {"ok": False,
                    "error": info.get("error") or "不支持自动更新"}
        if not info["needUpdate"]:
            return {"ok": True, "alreadyLatest": True, "current": my_version,
                    "message": "电脑端已经是最新版本 %s" % my_version}

        # 下载和应用放到后台线程，先回一句让手表知道开始了
        threading.Thread(
            target=self._run_update,
            args=(info, my_version),
            name="ykxmz-update", daemon=True).start()

        return {
            "ok": True,
            "started": True,
            "from": my_version,
            "to": info["latest"],
            "size": info.get("size", 0),
            "message": "正在下载 %s（%.1f MB），下完会自动装好并重启，"
                       "设置和配对码都不变"
                       % (info["latest"], (info.get("size") or 0) / 1048576.0),
        }

    def _run_update(self, info: dict, my_version: str) -> None:
        """后台线程：下载 -> 校验 -> 起帮手 -> 自己退出。"""
        def log(text: str) -> None:
            try:
                self._log(text)
            except Exception:                             # noqa: BLE001
                pass

        try:
            log("[更新] 开始下载 %s（%.1f MB）"
                % (info["latest"], (info.get("size") or 0) / 1048576.0))
            work = Path(tempfile.gettempdir()) / "ykxmz-update"
            last = [0]

            def progress(got: int, total: int) -> None:
                # 每 25% 报一次，别把日志刷爆
                if total:
                    pct = int(got * 100 / total)
                    if pct - last[0] >= 25:
                        last[0] = pct
                        log("[更新] 已下载 %d%%（%.1f/%.1f MB）"
                            % (pct, got / 1048576.0, total / 1048576.0))

            pkg = U.download(info["url"], work,
                             info.get("sha256", ""), info.get("size", 0),
                             progress=progress)
            log("[更新] 下载完成并校验通过：%s" % pkg.name)

            # 应用：起分离的 .cmd 帮手，然后**本进程必须退出**
            app_exe = Path(sys.executable if getattr(sys, "frozen", False)
                           else sys.argv[0]).resolve()
            logfile = U.apply_update(pkg, app_exe)
            log("[更新] 已交给更新帮手（%s），被控端即将退出并重启" % logfile)
            U.cleanup_old_downloads()

            # 给日志和网络一点时间落完，再退出
            time.sleep(2)
            os._exit(0)
        except U.UpdateError as exc:
            log("[更新] 失败：%s（保持当前版本 %s）" % (exc, my_version))
        except Exception as exc:                          # noqa: BLE001
            log("[更新] 意外失败：%s: %s" % (type(exc).__name__, exc))

'''
    text = text.replace(anchors["unknown"],
                        handler + "\n" + anchors["unknown"], 1)

    # ---- 补 import ----
    need = []
    head = text.split("class ", 1)[0]
    for mod, line in (("tempfile", "import tempfile"),
                      ("sys", "import sys"),
                      ("threading", "import threading"),
                      ("time", "import time"),
                      ("os", "import os"),
                      ("Path", "from pathlib import Path")):
        if line not in head:
            need.append(line)
    if need:
        anchor = "from . import protocol as P"
        text = text.replace(anchor, "\n".join(need) + "\n" + anchor, 1)
        print("  ✓ 补了 import: %s" % ", ".join(need))

    SERVER.write_bytes((b"\xef\xbb\xbf" if bom else b"") + text.encode("utf-8"))
    print("  ✓ server.py 加了 _do_update / _run_update")
    return True


def main() -> int:
    print("=" * 58)
    print("  接「静默自动更新」进协议和服务端")
    print("=" * 58)
    print()
    ok = patch_protocol()
    print()
    ok = patch_server() and ok
    return 0 if ok else 1


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