install_monitor.py

# -*- coding: utf-8 -*-
"""「ZCode Token 用量状态栏」安装监控窗口（安装完成后自动弹出的常驻命令行）。

每 10 秒检测一次，全程只读（唯一的写操作是 TMP 替换，见下），不联网：
1. 存在待替换 TMP（客户端运行中导致原子替换失败的场景）：等 ZCode 退出后自动完成
   asar 替换 —— 取代旧的 schtasks 计划任务收尾方案（finalize_if_closed.py 已删）。
2. asar 已替换：持续提醒重启 ZCode；检测到 ZCode 换了新进程（进程集合与启动采集的
   不相交）后，用 diag-<n>.json 的更新时间确认注入链路已加载，显示成功并自动退出。
   diag 是启动后第一次数据拉取就回写的（pushCount==1），空闲也会写，不依赖用户发消息。

手动运行：python install_monitor.py [install_epoch 秒] [app.asar 路径] [运行时目录]
（缺省用当前时间/默认 asar 位置；diag 在运行时目录——标准安装是数据目录副本所在，
 --dev 安装传仓库目录，不传则用本脚本所在目录）
退出：生效确认后自动退出；或 Ctrl+C / 直接关闭窗口。

输出语言：运行时目录 config.json 的 "lang" 字段（zh/en，缺省 zh）——与安装器/CLI 同一来源。
"""
import csv
import json
import os
import subprocess
import sys
import time
from pathlib import Path

HERE = Path(__file__).parent.resolve()
IS_MAC = sys.platform == "darwin"
IS_LINUX = sys.platform.startswith("linux")
# 缺省仅兜底（patch_install.py 拉起时必传 argv[2]）：按平台取常见安装位置
ASAR = Path(sys.argv[2]) if len(sys.argv) > 2 else (
    Path("/Applications/ZCode.app/Contents/Resources/app.asar") if IS_MAC
    else Path("/opt/ZCode/resources/app.asar") if IS_LINUX
    else Path(r"D:\ZCode\resources\app.asar"))
RUNTIME_DIR = Path(sys.argv[3]) if len(sys.argv) > 3 else HERE
TMP = ASAR.with_name("app.asar.zusage.tmp")
POLL = 10            # 检测间隔（秒）
CONFIRM_WAIT = 120   # 重启后等 diag 确认的超时（秒）
BYE = 8              # 成功提示停留秒数
QUIET = os.environ.get("ZUSAGE_MONITOR_QUIET") == "1"


def load_lang():
    try:
        lang = json.loads((RUNTIME_DIR / "config.json").read_text(encoding="utf-8")).get("lang")
        return lang if lang in ("zh", "en") else "zh"
    except (OSError, ValueError):
        return "zh"


LANG = load_lang()


def L(zh, en):
    """双语输出：按 LANG 返回对应文案（与 install.py 同款约定）。"""
    return en if LANG == "en" else zh


def zcode_pids():
    """ZCode 主进程 PID 集合（Windows=tasklist，macOS/Linux=pgrep）；探测命令不可用时返回 None。"""
    if os.name == "nt":
        try:
            r = subprocess.run(
                ["tasklist", "/FI", "IMAGENAME eq ZCode.exe", "/FO", "CSV", "/NH"],
                capture_output=True,
            )
        except OSError:
            return None
        pids = set()
        for row in csv.reader(r.stdout.decode("gbk", "replace").splitlines()):
            if len(row) >= 2 and row[0].strip().lower() == "zcode.exe":
                pids.add(row[1].strip())
        return pids
    try:
        r = subprocess.run(["pgrep", "-x", "ZCode"], capture_output=True)
    except OSError:
        return None
    if r.returncode == 1:
        return set()   # pgrep：无匹配进程
    if r.returncode != 0:
        return None    # pgrep 自身出错
    return {t.strip() for t in r.stdout.decode().split() if t.strip()}


def wait_any_exit(pids, timeout_s):
    """阻塞到任一 ZCode 进程退出或超时。Windows 用进程句柄事件驱动等待，消灭固定
    轮询间隔可能错过的「退出→立即重启」空档（等错过后监控只会看到新进程，永远收不了尾）；
    采集后已退出的 PID 视为退出已发生，立即返回；句柄不可得（权限/句柄数超限/非 Windows）
    时退化为普通睡眠。"""
    if os.name != "nt" or not pids:
        time.sleep(timeout_s)
        return
    import ctypes
    from ctypes import wintypes
    SYNCHRONIZE = 0x00100000
    ERROR_INVALID_PARAMETER = 87   # OpenProcess：该 PID 已没有进程
    MAX_WAIT = 60   # WaitForMultipleObjects 上限 64，留余量；进程数更多时退化为睡眠
    try:
        pids = [int(pid) for pid in pids]
    except ValueError:
        time.sleep(timeout_s)
        return
    if len(pids) > MAX_WAIT:
        time.sleep(timeout_s)
        return
    kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
    kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
    kernel32.OpenProcess.restype = wintypes.HANDLE
    kernel32.WaitForMultipleObjects.argtypes = [
        wintypes.DWORD, ctypes.POINTER(wintypes.HANDLE), wintypes.BOOL, wintypes.DWORD]
    kernel32.WaitForMultipleObjects.restype = wintypes.DWORD
    kernel32.CloseHandle.argtypes = [wintypes.HANDLE]
    handles = []
    try:
        for pid in pids:
            handle = kernel32.OpenProcess(SYNCHRONIZE, False, pid)
            if handle:
                handles.append(handle)
            elif ctypes.get_last_error() == ERROR_INVALID_PARAMETER:
                return   # 采集后已退出：「任一退出」已经发生，立即复查
        if not handles:
            # 进程都在但全部无权等待（如其它用户会话的 ZCode）：按轮询间隔睡，
            # 立即返回会让监控忙转。
            time.sleep(timeout_s)
            return
        array = (wintypes.HANDLE * len(handles))(*handles)
        kernel32.WaitForMultipleObjects(len(handles), array, False, int(timeout_s * 1000))
    finally:
        for handle in handles:
            kernel32.CloseHandle(handle)


def diag_fresh(since):
    """返回安装时刻之后更新过的 diag 文件名（注入链路已加载的证据），没有则 None。"""
    for f in sorted(RUNTIME_DIR.glob("diag-*.json")):
        try:
            if f.stat().st_mtime >= since:
                return f.name
        except OSError:
            pass
    return None


def patch_module():
    import patch_install as pi
    pi.set_target(ASAR)
    pi.set_runtime(RUNTIME_DIR)
    return pi


def finalize_tmp(verified=None):
    """监控按 sidecar 的 install/remove 收尾，不吞掉不可重试的完整性错误。"""
    pi = patch_module()
    return pi.finalize_tmp(operation=None, client_closed=lambda: zcode_pids() == set(),
                           verified=verified)


def draw(lines):
    os.system("cls" if os.name == "nt" else "clear")
    print("\n".join(lines), flush=True)


def quiet_finalize():
    """hook 无窗口收尾：来源由 repack 前 sidecar 决定，不用监控启动时间推断。"""
    pi = patch_module()
    verified = {}
    deadline = time.monotonic() + 86400
    while TMP.exists() and time.monotonic() < deadline:
        try:
            # 预校验并焐热 verified 缓存（issue #72）：每个 sidecar 世代只 hash 一次大包，
            # 此后轮询 stat 命中缓存；退出窗口内不再 hash 大包（剩余为 stat、client_closed
            # 进程探测、可选 publish_install 与 rename 等原有毫秒~亚秒级小步骤），
            # 不再输给新实例首进程亚秒级重开 asar 的竞速。哈希失败（OSError）下轮重试。
            pi.pending_tmp(operation=None, verify_hash=True, verified=verified)
            pids = zcode_pids()
            if pids == set() and finalize_tmp(verified=verified):
                return
            if pids:   # 客户端在跑：事件驱动等任一进程退出，立即复查（见 wait_any_exit）
                wait_any_exit(pids, POLL)
                continue
        except pi.PendingTmpError:
            return  # 完整性/来源失效立刻取消；不删旧包，更不能回滚客户端新版。
        except OSError:
            pass  # 文件占用等系统错误才重试，最多一天。
        time.sleep(POLL)


def main():
    if QUIET:
        quiet_finalize()
        return
    epoch = float(sys.argv[1]) if len(sys.argv) > 1 else time.time()
    if os.name == "nt":
        try:
            subprocess.run(["cmd", "/c", "title", L("ZCode Token 用量状态栏 - 安装监控", "ZCode Token Usage Status Bar - Install Monitor")])
        except OSError:
            pass
    pids0 = zcode_pids() or set()   # 启动采集 ≈ 安装时刻的进程集合（install 完成即拉起本窗口）
    t0 = time.time()
    new_since = None
    pi = patch_module()
    verified = {}
    deadline = time.monotonic() + 86400
    while True:
        pids = zcode_pids()
        now = time.time()
        done = False
        pending = None
        if TMP.exists():
            try:
                # 同 quiet_finalize：等待期预校验焐热缓存，退出窗口内不再 hash 大包（issue #72）。
                pending = pi.pending_tmp(operation=None, verify_hash=True, verified=verified)
                if time.monotonic() >= deadline:
                    draw([L("[!] 待替换包收尾已超过一天，停止自动重试。文件保留，请重新安装。",
                            "[!] Pending finalization timed out after one day; files kept. Re-run install.")])
                    return
            except pi.PendingTmpError as e:
                draw([L(f"[!] 待替换包来源/身份失效，取消收尾，文件保留：{e}",
                        f"[!] Pending source/identity invalid; finalization cancelled, files kept: {e}"),
                      L("请重新运行 install/remove，不要手动覆盖新版 asar。",
                        "Re-run install/remove; do not manually overwrite a newer asar.")])
                return
            except OSError as e:
                draw([L(f"[!] 暂时无法检查待替换包，稍后重试：{e}",
                        f"[!] Cannot check pending package yet; will retry: {e}")])
                if time.monotonic() >= deadline:
                    return
                time.sleep(POLL)
                continue
        if pids is None:
            lines = [L("[!] 无法检测 ZCode 进程（探测命令不可用）。",
                       "[!] Cannot detect ZCode processes (detection command unavailable).")]
        elif pending is not None:
            # 和 quiet/--finalize 共用来源、TMP hash、自检及退出校验。
            if not pids:
                try:
                    finished = finalize_tmp(verified=verified)
                except pi.PendingTmpError as e:
                    draw([L(f"[!] 待替换包已失效，取消收尾，文件保留：{e}",
                            f"[!] Pending package invalid; finalization cancelled, files kept: {e}")])
                    return
                except OSError:
                    finished = False
                if finished:
                    lines = [L("[OK] ZCode 已退出，asar 替换完成。", "[OK] ZCode has exited; app.asar replacement finished."),
                             "",
                             L("请启动 ZCode，悬浮条即可显示。", "Start ZCode and the floating bar will appear.")]
                    if pending["operation"] == "remove":
                        lines[-1] = L("注入已移除；重新启动 ZCode 后生效。",
                                      "Injection removed; takes effect when ZCode starts again.")
                        done = True
                else:
                    lines = [L("[!] ZCode 已退出但替换失败（文件仍被占用？）",
                               "[!] ZCode has exited but the replacement failed (file still locked?)"),
                             L("稍后自动重试；也可手动执行：", "Will retry automatically; you can also run manually:"),
                             "  python patch_install.py "
                             + (f'remove --asar "{ASAR}"' if pending["operation"] == "remove"
                                else f'install --finalize --asar "{ASAR}" --runtime "{RUNTIME_DIR}"')]
            else:
                lines = [L("[等待] 有待替换文件，请完全退出 ZCode。",
                           "[waiting] A pending replacement file exists; please exit ZCode completely."),
                         L("退出后本窗口将自动完成 asar 替换。",
                           "This window will finish the asar replacement automatically once you exit."),
                         "", L("当前 ZCode 进程数：%d" % len(pids), "Current ZCode process count: %d" % len(pids))]
        elif not pids:
            lines = [L("ZCode 未运行。启动 ZCode 后悬浮条即生效。",
                       "ZCode is not running. The floating bar activates once you start ZCode.")]
        elif pids & pids0:
            lines = [L(">>>>>  请 重 启 Z Code  <<<<<", ">>>>>  PLEASE RESTART ZCODE  <<<<<"),
                     "",
                     L("安装已完成，但当前还是旧实例（%d 个进程），" % len(pids),
                       "Installation finished, but the old instance is still running (%d processes)," % len(pids)),
                     L("请完全退出 ZCode（所有窗口）后重新打开。",
                       "please quit ZCode completely (all windows) and reopen it.")]
        else:
            if new_since is None:
                new_since = now
            d = diag_fresh(epoch)
            if d:
                lines = [L("[OK] ZCode 已重启，注入已加载（%s 已更新）。" % d,
                           "[OK] ZCode restarted, injection loaded (%s updated)." % d),
                         L("悬浮条应已出现在输入框下方。", "The floating bar should now be visible."),
                         "", L("本窗口即将自动关闭。", "This window will close automatically.")]
                done = True
            elif now - new_since > CONFIRM_WAIT:
                lines = [L("[?] ZCode 已重启，但 %d 秒内未检测到注入加载。" % CONFIRM_WAIT,
                           "[?] ZCode restarted, but no injection load detected within %ds." % CONFIRM_WAIT),
                         L("若 ZCode 刚升级过，asar 会被官方版本覆盖，",
                           "If ZCode was just upgraded, app.asar was overwritten by the official build —"),
                         L(f'请重新运行：python patch_install.py install --asar "{ASAR}"',
                           f'please re-run: python patch_install.py install --asar "{ASAR}"'),
                         L("（继续检测中，diag 一更新即确认）", "(still watching; will confirm as soon as diag updates)")]
            else:
                lines = [L("ZCode 已重启（新实例），正在确认注入加载……",
                           "ZCode restarted (new instance); confirming injection load…")]
        el = int(now - t0)
        lines += ["", "-" * 52,
                  L("已等待 %02d:%02d | 每 %d 秒检测一次 | 生效后自动退出" % (el // 60, el % 60, POLL),
                    "waited %02d:%02d | checking every %ds | exits automatically once confirmed" % (el // 60, el % 60, POLL)),
                  L("本窗口可随时关闭，不影响安装结果。", "You can close this window at any time; the installation result is unaffected.")]
        draw(lines)
        if done:
            time.sleep(BYE)
            return
        if pending is not None and pids:
            wait_any_exit(pids, POLL)   # 有待替换且客户端在跑：任一进程退出立即醒来复查
        else:
            time.sleep(POLL)


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        pass