382 lines
15 KiB
Python
382 lines
15 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
外部独立更新器 —— 主进程退出后完成真正的替换/重启。
|
||
|
||
不随主进程死、不被主进程文件锁影响、失败可回滚。
|
||
|
||
用法（主程序 Popen 启动）:
|
||
update_helper.exe --src <下载包路径> --dst <目标exe路径> --pid <当前PID> [--restart] [--mode auto|onefile|zip] [--keepbackup]
|
||
|
||
不依赖计划任务，普通用户权限即可覆盖自己的安装目录。
|
||
如果用户装在 C:\\Program Files 等受限目录，需配合计划任务提权（保留 update_helper.spec 打包为 exe 的能力）。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import ctypes
|
||
import ctypes.wintypes as wt
|
||
import json
|
||
import os
|
||
import shutil
|
||
import subprocess
|
||
import sys
|
||
import tempfile
|
||
import time
|
||
import traceback
|
||
import zipfile
|
||
|
||
|
||
LOG_FILE = "update.log"
|
||
|
||
|
||
# ────────── 日志 ────────────────────────────────────────────
|
||
def _log(msg: str) -> None:
|
||
try:
|
||
p = os.path.join(tempfile.gettempdir(), LOG_FILE)
|
||
with open(p, "a", encoding="utf-8") as f:
|
||
f.write(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] {msg}\n")
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
# ────────── 等待进程退出（Win32 精准版） ───────────────────
|
||
def _wait_process_exit(pid: int, timeout_sec: float = 90.0) -> bool:
|
||
"""用 OpenProcess + WaitForSingleObject 等进程退出。
|
||
比 tasklist 轮询精准：能检测 zombie（已退出但句柄未释放）、
|
||
能被进程句柄自动释放（不依赖命令行解析）。"""
|
||
if pid <= 0:
|
||
_log("wait: no pid, skip")
|
||
return True
|
||
user32 = ctypes.windll.user32
|
||
kernel32 = ctypes.windll.kernel32
|
||
SYNCHRONIZE = 0x00100000
|
||
PROCESS_QUERY_LIMITED_INFORMATION = 0x1000
|
||
handle = kernel32.OpenProcess(SYNCHRONIZE | PROCESS_QUERY_LIMITED_INFORMATION, False, pid)
|
||
if not handle:
|
||
# OpenProcess 失败通常意味着进程已不存在（访问拒绝）
|
||
_log(f"wait: OpenProcess pid={pid} failed, assume gone")
|
||
return True
|
||
try:
|
||
waited = 0.0
|
||
step = 0.5
|
||
while True:
|
||
rc = kernel32.WaitForSingleObject(handle, int(step * 1000))
|
||
waited += step
|
||
if rc == 0:
|
||
_log(f"wait: pid={pid} exited after {waited:.1f}s")
|
||
return True
|
||
if waited >= timeout_sec:
|
||
_log(f"wait: pid={pid} TIMEOUT after {timeout_sec}s")
|
||
return False
|
||
finally:
|
||
kernel32.CloseHandle(handle)
|
||
|
||
|
||
# ────────── 文件操作（带重试 + 原子替换） ────────────────────
|
||
def _atomic_replace(src: str, dst: str, backup: str | None = None, retries: int = 12) -> bool:
|
||
"""原子替换：备份旧文件 → os.replace 覆盖 → 失败则恢复备份。"""
|
||
# 1. 备份旧文件
|
||
old_exists = os.path.exists(dst)
|
||
if old_exists and backup:
|
||
try:
|
||
os.replace(dst, backup)
|
||
_log(f"backup: {dst} -> {backup}")
|
||
except Exception as e:
|
||
_log(f"backup failed: {e}")
|
||
return False # 连备份都做不了 = 没权限
|
||
|
||
# 2. 尝试替换（os.replace 在 Windows 上对目标进程已退出的文件是原子的）
|
||
for i in range(max(1, retries)):
|
||
try:
|
||
os.replace(src, dst)
|
||
_log(f"replace: ok ({i} attempts)")
|
||
return True
|
||
except Exception as e:
|
||
_log(f"replace attempt {i + 1}/{retries} failed: {e}")
|
||
time.sleep(0.7)
|
||
|
||
# 3. 全部失败 → 恢复备份
|
||
if backup and os.path.exists(backup) and not os.path.exists(dst):
|
||
try:
|
||
os.replace(backup, dst)
|
||
_log("rollback: restored backup")
|
||
except Exception as e:
|
||
_log(f"rollback failed: {e}")
|
||
return False
|
||
|
||
|
||
def _clean_onedir_dead_sibling(target_exe: str) -> None:
|
||
"""onefile 覆盖 onedir 启动器后，同级遗留的 _internal/ 等死目录清理。"""
|
||
exe_dir = os.path.dirname(os.path.abspath(target_exe))
|
||
if not exe_dir:
|
||
return
|
||
for name in ("_internal", "__pycache__"):
|
||
dead = os.path.join(exe_dir, name)
|
||
if os.path.isdir(dead):
|
||
try:
|
||
shutil.rmtree(dead, ignore_errors=True)
|
||
_log(f"cleaned dead dir: {dead}")
|
||
except Exception as e:
|
||
_log(f"clean dead dir failed: {e}")
|
||
|
||
|
||
# ────────── zip 解压 + 路径归一 ─────────────────────────────
|
||
def _extract_zip_normalized(zip_path: str, dest_dir: str) -> str:
|
||
"""解压 zip 到 dest_dir，自动 strip 单一顶层目录（zip 内嵌套一层目录）。
|
||
返回实际 payload 根目录（可能是 dest_dir 本身，也可能是 strip 后的子目录）。"""
|
||
os.makedirs(dest_dir, exist_ok=True)
|
||
with zipfile.ZipFile(zip_path, "r") as zf:
|
||
zf.extractall(dest_dir)
|
||
# 检测是否所有文件都在同一顶层目录下
|
||
entries = set()
|
||
with zipfile.ZipFile(zip_path, "r") as zf:
|
||
for name in zf.namelist():
|
||
# strip 开头的 ./
|
||
parts = name.replace("\\", "/").lstrip("./").split("/")
|
||
if parts[0]:
|
||
entries.add(parts[0])
|
||
if len(entries) == 1:
|
||
top = next(iter(entries))
|
||
candidate = os.path.join(dest_dir, top)
|
||
if os.path.isdir(candidate):
|
||
_log(f"zip: strip top folder '{top}', payload={candidate}")
|
||
return candidate
|
||
_log(f"zip: no single top folder, payload={dest_dir}")
|
||
return dest_dir
|
||
|
||
|
||
# ────────── 目录整体替换（zip 模式） ──────────────────────────
|
||
def _replace_directory(old_dir: str, new_payload: str, keep_backup: bool = False) -> bool:
|
||
"""备份旧目录为 .bak.时间戳 → 移动新 payload 到旧目录位置 → 清理备份（默认保留最近 3 个）。"""
|
||
old_dir = os.path.abspath(old_dir)
|
||
if not os.path.isdir(old_dir):
|
||
_log(f"replace dir: old_dir {old_dir} not found, treat as fresh install")
|
||
os.makedirs(old_dir, exist_ok=True)
|
||
_copy_tree(new_payload, old_dir)
|
||
return True
|
||
|
||
parent = os.path.dirname(old_dir)
|
||
base = os.path.basename(old_dir)
|
||
ts = time.strftime("%Y%m%d-%H%M%S")
|
||
backup_dir = os.path.join(parent, f"{base}.bak.{ts}")
|
||
|
||
# 1. 备份旧目录
|
||
try:
|
||
os.rename(old_dir, backup_dir)
|
||
_log(f"backup dir: {old_dir} -> {backup_dir}")
|
||
except Exception as e:
|
||
_log(f"backup dir failed: {e}")
|
||
return False
|
||
|
||
# 2. 移动新 payload 到目标位置
|
||
try:
|
||
if os.path.isdir(new_payload):
|
||
# 直接 rename 最快，但可能跨盘符 → 用 shutil.move
|
||
moved = shutil.move(new_payload, old_dir)
|
||
_log(f"move: {new_payload} -> {moved}")
|
||
else:
|
||
_log(f"new_payload {new_payload} is not a dir")
|
||
os.rename(backup_dir, old_dir) # 恢复
|
||
return False
|
||
except Exception as e:
|
||
_log(f"move failed: {e}")
|
||
try:
|
||
os.rename(backup_dir, old_dir)
|
||
except Exception:
|
||
pass
|
||
return False
|
||
|
||
# 3. 清理旧备份（保留最近 3 个）
|
||
if not keep_backup:
|
||
_prune_old_backups(parent, base, keep=3)
|
||
return True
|
||
|
||
|
||
def _copy_tree(src: str, dst: str) -> None:
|
||
"""递归复制目录树，覆盖目标文件。"""
|
||
os.makedirs(dst, exist_ok=True)
|
||
for item in os.listdir(src):
|
||
s = os.path.join(src, item)
|
||
d = os.path.join(dst, item)
|
||
if os.path.isdir(s):
|
||
_copy_tree(s, d)
|
||
else:
|
||
shutil.copy2(s, d)
|
||
|
||
|
||
def _prune_old_backups(parent: str, base: str, keep: int = 3) -> None:
|
||
"""清理旧的 .bak.* 目录，只保留最近 keep 个。"""
|
||
try:
|
||
backups = sorted(
|
||
[d for d in os.listdir(parent) if d.startswith(f"{base}.bak.")],
|
||
reverse=True,
|
||
)
|
||
for old in backups[keep:]:
|
||
old_path = os.path.join(parent, old)
|
||
try:
|
||
shutil.rmtree(old_path, ignore_errors=True)
|
||
_log(f"prune backup: {old_path}")
|
||
except Exception as e:
|
||
_log(f"prune failed: {e}")
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
# ────────── 启动新进程 ──────────────────────────────────────
|
||
def _launch_new_exe(exe_path: str) -> bool:
|
||
try:
|
||
# --from-update: 让新进程跳过单实例检测（旧 mutex 可能还没完全释放）
|
||
# CREATE_NEW_CONSOLE：新进程创建独立控制台（真正脱离 update_helper）
|
||
# CREATE_NO_WINDOW： 控制台窗口不显示
|
||
subprocess.Popen(
|
||
[exe_path, "--from-update"],
|
||
creationflags=subprocess.CREATE_NEW_CONSOLE | subprocess.CREATE_NO_WINDOW,
|
||
)
|
||
_log(f"launched: {exe_path} --from-update")
|
||
return True
|
||
except Exception as e:
|
||
_log(f"launch failed: {e}")
|
||
return False
|
||
|
||
|
||
# ────────── 包类型判断 ──────────────────────────────────────
|
||
def _detect_mode(src: str) -> str:
|
||
ext = os.path.splitext(src)[1].lower()
|
||
if ext == ".zip":
|
||
return "zip"
|
||
if ext in (".exe",):
|
||
return "onefile"
|
||
# 尝试打开 zip 看看
|
||
try:
|
||
with zipfile.ZipFile(src, "r"):
|
||
return "zip"
|
||
except Exception:
|
||
pass
|
||
return "onefile"
|
||
|
||
|
||
# ────────── 主流程 ───────────────────────────────────────────
|
||
def main(argv: list[str] | None = None) -> int:
|
||
parser = argparse.ArgumentParser(description="外部独立更新器")
|
||
parser.add_argument("--src", required=True, help="下载包路径（onefile exe 或 zip）")
|
||
parser.add_argument("--dst", required=True, help="目标 exe 路径")
|
||
parser.add_argument("--pid", type=int, default=0, help="当前运行进程 PID，用于等待退出")
|
||
parser.add_argument("--restart", action="store_true", help="替换完成后启动新进程")
|
||
parser.add_argument("--mode", choices=["auto", "onefile", "zip"], default="auto")
|
||
parser.add_argument("--keepbackup", action="store_true", help="保留备份不清理")
|
||
parser.add_argument("--zip-root", default=None,
|
||
help="zip 包内的 exe 相对路径（如 niumasoftware/niumasoftware.exe），"
|
||
"不传则自动在解压结果中查找同名 exe")
|
||
args = parser.parse_args(argv)
|
||
|
||
_log(f"=== START args={json.dumps(vars(args))}")
|
||
|
||
src = os.path.abspath(args.src)
|
||
dst = os.path.abspath(args.dst)
|
||
mode = args.mode if args.mode != "auto" else _detect_mode(src)
|
||
|
||
if not os.path.exists(src):
|
||
_log(f"ERROR: src not found: {src}")
|
||
return 2
|
||
|
||
# 1. 等待主进程退出
|
||
if args.pid > 0:
|
||
if not _wait_process_exit(args.pid, timeout_sec=90):
|
||
_log("WARN: main process still alive after timeout, attempt replacement anyway")
|
||
|
||
if mode == "onefile":
|
||
# ─── onefile 模式 ────────────────────────────────────
|
||
_log("mode: onefile exe")
|
||
dst_dir = os.path.dirname(dst)
|
||
backup_exe = os.path.join(dst_dir, os.path.basename(dst) + ".bak")
|
||
|
||
# 原子替换
|
||
if not _atomic_replace(src, dst, backup=backup_exe):
|
||
_log("FATAL: onefile replace failed, cannot continue")
|
||
return 3
|
||
|
||
# 清理遗留的 onedir 死目录
|
||
_clean_onedir_dead_sibling(dst)
|
||
|
||
# 删备份（保留 1 个版本足够回滚）
|
||
if os.path.exists(backup_exe):
|
||
try:
|
||
# 新 exe 启动正常后用户可以手动删，这里只删旧的 .bak
|
||
os.remove(backup_exe)
|
||
_log("backup cleaned")
|
||
except Exception:
|
||
pass
|
||
|
||
if args.restart:
|
||
if not _launch_new_exe(dst):
|
||
return 4
|
||
|
||
elif mode == "zip":
|
||
# ─── zip 整包模式 ───────────────────────────────────
|
||
_log("mode: zip full-package")
|
||
tmp_extract = tempfile.mkdtemp(prefix="niumasoft_update_")
|
||
try:
|
||
payload_root = _extract_zip_normalized(src, tmp_extract)
|
||
|
||
# 确定旧目录和目标 exe 在旧目录里
|
||
old_dir = os.path.dirname(dst)
|
||
new_exe_in_payload = dst # 默认同名 exe 相对路径
|
||
if args.zip_root:
|
||
new_exe_in_payload = os.path.join(payload_root, args.zip_root)
|
||
if not os.path.exists(new_exe_in_payload):
|
||
_log(f"ERROR: --zip-root exe not found: {new_exe_in_payload}")
|
||
return 5
|
||
else:
|
||
# 自动在 payload 里找同名 exe
|
||
target_name = os.path.basename(dst)
|
||
found = None
|
||
for root, _, files in os.walk(payload_root):
|
||
if target_name in files:
|
||
found = os.path.join(root, target_name)
|
||
break
|
||
if not found:
|
||
_log(f"ERROR: cannot find {target_name!r} in payload {payload_root}")
|
||
return 5
|
||
new_exe_in_payload = found
|
||
_log(f"auto-found exe in payload: {new_exe_in_payload}")
|
||
|
||
# payload 根目录 = exe 所在目录的父级（可能就是 payload_root，也可能子级）
|
||
payload_exe_dir = os.path.dirname(new_exe_in_payload)
|
||
if os.path.abspath(payload_exe_dir) != os.path.abspath(payload_root):
|
||
_log(f"payload exe dir != payload root, using exe dir as new dir root")
|
||
new_payload = payload_exe_dir
|
||
else:
|
||
new_payload = payload_root
|
||
|
||
# 替换整个旧目录
|
||
ok = _replace_directory(old_dir, new_payload, keep_backup=args.keepbackup)
|
||
if not ok:
|
||
_log("FATAL: directory replace failed")
|
||
return 3
|
||
|
||
new_dst = os.path.join(old_dir, os.path.basename(dst))
|
||
if args.restart:
|
||
if not _launch_new_exe(new_dst):
|
||
return 4
|
||
finally:
|
||
# 清理临时解压目录
|
||
shutil.rmtree(tmp_extract, ignore_errors=True)
|
||
_log(f"cleaned temp: {tmp_extract}")
|
||
|
||
else:
|
||
_log(f"FATAL: unknown mode {mode}")
|
||
return 1
|
||
|
||
_log("=== DONE")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
try:
|
||
sys.exit(main())
|
||
except SystemExit:
|
||
raise
|
||
except Exception:
|
||
_log(f"UNCAUGHT: {traceback.format_exc()}")
|
||
sys.exit(99)
|