cmshoppe/app/updater_entry.py

"""在主程序退出后执行程序根项目事务替换的独立更新器。"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import shutil
import subprocess
import tempfile
import time
import uuid
from dataclasses import dataclass
from pathlib import Path, PurePosixPath

from app import update_health


PACKAGE_FORMAT = "cmshopee-portable-v1"
UPDATER_PROTOCOL = 1
ALLOWED_ROOTS = {
    "cmshopee.exe",
    "_internal",
    "version.txt",
    "README.txt",
    "package-manifest.json",
    "cmshopee-updater.exe",
}
PROTECTED_ROOTS = {"data", ".cmshopee-update"}


class UpdaterError(RuntimeError):
    """独立更新器拒绝计划或无法安全完成事务。"""


@dataclass(frozen=True)
class UpdatePlan:
    parent_pid: int
    install_root: Path
    staging_root: Path
    target_version: str
    transaction_id: str
    log_path: Path
    package_sha256: str = ""


def create_plan(install_root, staging_root, target_version, parent_pid, package_sha256=""):
    install_root = Path(install_root).resolve()
    transaction_id = uuid.uuid4().hex
    update_root = install_root / ".cmshopee-update"
    plan = UpdatePlan(
        parent_pid=int(parent_pid),
        install_root=install_root,
        staging_root=Path(staging_root).resolve(),
        target_version=str(target_version),
        transaction_id=transaction_id,
        log_path=(update_root / "logs" / (transaction_id + ".log")).resolve(),
        package_sha256=str(package_sha256 or "").lower(),
    )
    validate_plan(plan)
    plan_path = update_root / "plans" / (transaction_id + ".json")
    _atomic_json(
        plan_path,
        {
            "parent_pid": plan.parent_pid,
            "install_root": str(plan.install_root),
            "staging_root": str(plan.staging_root),
            "target_version": plan.target_version,
            "transaction_id": plan.transaction_id,
            "log_path": str(plan.log_path),
            "package_sha256": plan.package_sha256,
        },
    )
    return plan_path


def _is_child(path, parent):
    try:
        Path(path).resolve().relative_to(Path(parent).resolve())
        return True
    except ValueError:
        return False


def _atomic_json(path, payload):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    os.replace(str(temporary), str(path))


def _append_log(path, message):
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    with Path(path).open("a", encoding="utf-8") as stream:
        stream.write("%s %s\n" % (time.strftime("%Y-%m-%d %H:%M:%S"), message))


def load_plan(plan_path):
    plan_path = Path(plan_path).resolve()
    try:
        payload = json.loads(plan_path.read_text(encoding="utf-8"))
        plan = UpdatePlan(
            parent_pid=int(payload["parent_pid"]),
            install_root=Path(payload["install_root"]).resolve(),
            staging_root=Path(payload["staging_root"]).resolve(),
            target_version=str(payload["target_version"]),
            transaction_id=str(payload["transaction_id"]),
            log_path=Path(payload["log_path"]).resolve(),
            package_sha256=str(payload.get("package_sha256") or "").lower(),
        )
    except (OSError, ValueError, KeyError, TypeError) as exc:
        raise UpdaterError("更新计划无效") from exc
    validate_plan(plan, plan_path)
    return plan


def validate_plan(plan, plan_path=None):
    if plan.parent_pid <= 0:
        raise UpdaterError("更新计划缺少主程序进程编号")
    if not plan.install_root.is_absolute() or not plan.install_root.is_dir():
        raise UpdaterError("程序安装目录无效")
    update_root = plan.install_root / ".cmshopee-update"
    staging_parent = update_root / "staging"
    if not plan.staging_root.is_dir() or plan.staging_root.parent != staging_parent.resolve():
        raise UpdaterError("新版暂存目录越界")
    if not re.fullmatch(r"\d+(?:\.\d+)*", plan.target_version):
        raise UpdaterError("目标版本号格式不正确")
    if not re.fullmatch(r"[A-Za-z0-9-]{8,64}", plan.transaction_id):
        raise UpdaterError("更新事务编号无效")
    if plan.package_sha256 and not re.fullmatch(r"[0-9a-f]{64}", plan.package_sha256):
        raise UpdaterError("更新安装包校验值无效")
    if not _is_child(plan.log_path, update_root / "logs"):
        raise UpdaterError("更新日志路径越界")
    if plan_path is not None and not _is_child(plan_path, update_root):
        raise UpdaterError("更新计划文件路径越界")
    probe = update_root / (".write-probe-" + plan.transaction_id)
    try:
        probe.parent.mkdir(parents=True, exist_ok=True)
        probe.write_bytes(b"")
        probe.unlink()
    except OSError as exc:
        raise UpdaterError("程序安装目录不可写") from exc


def _sha256(path):
    digest = hashlib.sha256()
    with Path(path).open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def validate_staging(plan):
    manifest_path = plan.staging_root / "package-manifest.json"
    try:
        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    except (OSError, ValueError) as exc:
        raise UpdaterError("新版程序清单无效") from exc
    if (
        manifest.get("package_format") != PACKAGE_FORMAT
        or manifest.get("app_version") != plan.target_version
        or manifest.get("entry_point") != "cmshopee.exe"
        or int(manifest.get("updater_protocol") or 0) != UPDATER_PROTOCOL
    ):
        raise UpdaterError("新版程序清单与更新计划不一致")
    roots = manifest.get("replace_roots")
    if not isinstance(roots, list) or not roots:
        raise UpdaterError("新版程序替换范围无效")
    roots = list(dict.fromkeys(str(root) for root in roots))
    if any(root not in ALLOWED_ROOTS or root in PROTECTED_ROOTS for root in roots):
        raise UpdaterError("新版程序包含未授权替换范围")
    if "cmshopee.exe" not in roots or "_internal" not in roots:
        raise UpdaterError("新版程序缺少必要替换项目")

    declared = {}
    for item in manifest.get("files") or []:
        relative = PurePosixPath(str(item.get("path") or ""))
        if not relative.parts or ".." in relative.parts or relative.parts[0] not in roots:
            raise UpdaterError("新版程序清单包含越界文件")
        key = relative.as_posix().casefold()
        if key in declared:
            raise UpdaterError("新版程序清单包含重复文件")
        try:
            size_bytes = int(item["size_bytes"])
        except (KeyError, TypeError, ValueError) as exc:
            raise UpdaterError("新版程序清单文件大小无效") from exc
        if size_bytes < 0:
            raise UpdaterError("新版程序清单文件大小无效")
        declared[key] = (relative, size_bytes, str(item.get("sha256") or ""))
    actual = {}
    for path in plan.staging_root.rglob("*"):
        if path.is_symlink():
            raise UpdaterError("新版程序不允许符号链接")
        if path.is_file() and path.name != "package-manifest.json":
            relative = PurePosixPath(path.relative_to(plan.staging_root).as_posix())
            actual[relative.as_posix().casefold()] = relative
    if set(actual) != set(declared):
        raise UpdaterError("新版程序文件与清单不一致")
    for key, (relative, size_bytes, expected_hash) in declared.items():
        path = plan.staging_root.joinpath(*relative.parts)
        if path.stat().st_size != size_bytes or _sha256(path).lower() != expected_hash.lower():
            raise UpdaterError("新版程序文件校验失败")
    if (plan.staging_root / "version.txt").read_text(encoding="utf-8-sig").strip() != plan.target_version:
        raise UpdaterError("新版程序版本文件不一致")
    return roots


def _pid_running(pid):
    if pid <= 0:
        return False
    if os.name == "nt":
        import ctypes

        handle = ctypes.windll.kernel32.OpenProcess(0x1000, False, pid)
        if handle:
            ctypes.windll.kernel32.CloseHandle(handle)
            return True
        return False
    try:
        os.kill(pid, 0)
        return True
    except OSError:
        return False


def wait_for_parent_exit(pid, timeout=60, poll_seconds=0.2):
    deadline = time.monotonic() + timeout
    while _pid_running(pid):
        if time.monotonic() >= deadline:
            raise UpdaterError("等待旧版程序退出超时")
        time.sleep(poll_seconds)


class TransactionLock:
    def __init__(self, path):
        self.path = Path(path)
        self.fd = None

    def __enter__(self):
        self.path.parent.mkdir(parents=True, exist_ok=True)
        try:
            self.fd = os.open(str(self.path), os.O_CREAT | os.O_EXCL | os.O_WRONLY)
            os.write(self.fd, str(os.getpid()).encode("ascii"))
        except FileExistsError as exc:
            raise UpdaterError("已有更新程序正在处理当前安装目录") from exc
        return self

    def __exit__(self, *_args):
        if self.fd is not None:
            os.close(self.fd)
        try:
            self.path.unlink()
        except FileNotFoundError:
            pass


def _default_launch(executable, args):
    creationflags = getattr(subprocess, "CREATE_NO_WINDOW", 0)
    return subprocess.Popen([str(executable), *args], close_fds=True, creationflags=creationflags)


def wait_for_health(plan, process, timeout=90, poll_seconds=0.25):
    deadline = time.monotonic() + timeout
    health_path = (
        plan.install_root
        / ".cmshopee-update"
        / "transactions"
        / plan.transaction_id
        / "health.json"
    )
    last_status = ""
    while time.monotonic() < deadline:
        try:
            payload = json.loads(health_path.read_text(encoding="utf-8"))
            if (
                payload.get("transaction_id") == plan.transaction_id
                and payload.get("target_version") == plan.target_version
            ):
                last_status = str(payload.get("status") or "")
                if last_status in {"main_window_ready", "environment_blocked"}:
                    return last_status
        except (OSError, ValueError):
            pass
        if process is not None and hasattr(process, "poll") and process.poll() is not None:
            raise UpdaterError("新版程序在主窗口就绪前退出")
        time.sleep(poll_seconds)
    if last_status == "process_started":
        raise UpdaterError("新版程序启动后未能显示主窗口")
    raise UpdaterError("等待新版程序启动确认超时")


def _cleanup_after_health(plan, backup_root, remove_backup):
    update_root = plan.install_root / ".cmshopee-update"
    for path in (
        update_root / "pending.json",
        update_root / "downloads" / (plan.target_version + ".zip"),
        update_root / "downloads" / (plan.target_version + ".zip.part"),
        update_root / "plans" / (plan.transaction_id + ".json"),
    ):
        try:
            path.unlink()
        except FileNotFoundError:
            pass
    if plan.staging_root.exists():
        shutil.rmtree(str(plan.staging_root), ignore_errors=True)
    if remove_backup and backup_root.exists():
        shutil.rmtree(str(backup_root), ignore_errors=True)


def prune_backups(update_root, keep=2):
    backup_parent = Path(update_root) / "backup"
    if not backup_parent.is_dir():
        return
    backups = sorted(
        (path for path in backup_parent.iterdir() if path.is_dir()),
        key=lambda path: path.stat().st_mtime,
        reverse=True,
    )
    for path in backups[max(0, int(keep)) :]:
        shutil.rmtree(str(path), ignore_errors=True)


def apply_update(
    plan,
    *,
    wait_parent=wait_for_parent_exit,
    move=os.replace,
    launcher=_default_launch,
    health_waiter=wait_for_health,
):
    validate_plan(plan)
    roots = validate_staging(plan)
    update_root = plan.install_root / ".cmshopee-update"
    old_version_path = plan.install_root / "version.txt"
    old_version = old_version_path.read_text(encoding="utf-8-sig").strip() if old_version_path.exists() else "unknown"
    backup_root = update_root / "backup" / (old_version + "-" + plan.transaction_id)
    failed_root = update_root / "failed" / plan.transaction_id
    journal_path = update_root / "transactions" / plan.transaction_id / "journal.json"
    operations = []

    def journal(status, error=""):
        _atomic_json(
            journal_path,
            {
                "transaction_id": plan.transaction_id,
                "target_version": plan.target_version,
                "status": status,
                "operations": operations,
                "error": error,
            },
        )

    with TransactionLock(update_root / "update.lock"):
        wait_parent(plan.parent_pid)
        backup_root.mkdir(parents=True, exist_ok=False)
        journal("moving_old")
        try:
            for root in roots:
                source = plan.install_root / root
                if source.exists():
                    destination = backup_root / root
                    destination.parent.mkdir(parents=True, exist_ok=True)
                    move(str(source), str(destination))
                    operations.append({"kind": "old_to_backup", "root": root})
                    journal("moving_old")
            for root in roots:
                source = plan.staging_root / root
                if not source.exists():
                    raise UpdaterError("新版程序缺少替换项目：%s" % root)
                destination = plan.install_root / root
                move(str(source), str(destination))
                operations.append({"kind": "new_to_install", "root": root})
                journal("moving_new")
            process = launcher(
                plan.install_root / "cmshopee.exe",
                ["--update-transaction", plan.transaction_id, "--update-target-version", plan.target_version],
            )
            journal("launched")
            _append_log(plan.log_path, "新版程序已替换并启动：%s" % plan.target_version)
            health_status = health_waiter(plan, process)
            if health_status == "environment_blocked":
                journal("environment_blocked")
                _append_log(plan.log_path, "新版程序已启动，但本地环境需要用户处理")
                _cleanup_after_health(plan, backup_root, remove_backup=False)
                prune_backups(update_root)
                return backup_root
            journal("healthy")
            _append_log(plan.log_path, "新版主窗口已就绪：%s" % plan.target_version)
            _cleanup_after_health(plan, backup_root, remove_backup=True)
            prune_backups(update_root)
            return backup_root
        except Exception as exc:
            if any(operation["kind"] == "new_to_install" for operation in operations):
                update_health.record_failed_release(
                    plan.install_root,
                    plan.target_version,
                    plan.package_sha256,
                    str(exc),
                )
            journal("rolling_back", str(exc))
            failed_root.mkdir(parents=True, exist_ok=True)
            rollback_errors = []
            for operation in reversed(operations):
                root = operation["root"]
                try:
                    if operation["kind"] == "new_to_install":
                        source = plan.install_root / root
                        if source.exists():
                            destination = failed_root / root
                            destination.parent.mkdir(parents=True, exist_ok=True)
                            move(str(source), str(destination))
                    else:
                        source = backup_root / root
                        if source.exists():
                            move(str(source), str(plan.install_root / root))
                except Exception as rollback_exc:
                    rollback_errors.append(str(rollback_exc))
            status = "rollback_failed" if rollback_errors else "rolled_back"
            journal(status, "；".join(rollback_errors) or str(exc))
            _append_log(plan.log_path, "更新失败，已尝试恢复旧版：%s" % exc)
            if rollback_errors:
                raise UpdaterError("更新失败且旧版恢复不完整，请联系管理员") from exc
            raise UpdaterError("更新失败，已恢复旧版程序") from exc


def copy_and_launch_updater(updater_path, plan_path):
    updater_path = Path(updater_path).resolve()
    if not updater_path.is_file():
        raise UpdaterError("独立更新程序不存在")
    temporary_dir = Path(tempfile.mkdtemp(prefix="cmshopee-update-"))
    external_updater = temporary_dir / updater_path.name
    shutil.copy2(str(updater_path), str(external_updater))
    creationflags = getattr(subprocess, "CREATE_NO_WINDOW", 0) | getattr(
        subprocess, "DETACHED_PROCESS", 0
    )
    return subprocess.Popen(
        [str(external_updater), "--plan", str(Path(plan_path).resolve())],
        close_fds=True,
        creationflags=creationflags,
    )


def main(argv=None):
    parser = argparse.ArgumentParser(description="蝦皮圈优化助手独立更新程序")
    parser.add_argument("--plan", required=True)
    args = parser.parse_args(argv)
    try:
        plan = load_plan(args.plan)
        apply_update(plan)
        return 0
    except Exception as exc:
        try:
            if "plan" in locals():
                _append_log(plan.log_path, "自动更新失败：%s" % exc)
        except OSError:
            pass
        return 1


if __name__ == "__main__":
    raise SystemExit(main())