xieyi-weixin-/installer_payload.py

"""Verified installer payloads and reversible same-volume installation swaps."""
from __future__ import annotations

import hashlib
import json
import os
import shutil
import tempfile
import uuid
from pathlib import Path

MANIFEST_NAME = "installation-manifest.json"
REQUIRED_FILES = (
    "ZhenAIUpdater.exe", "_internal/base_library.zip",
    "_internal/shiboken6/Shiboken.pyd", "_internal/shiboken6/shiboken6.abi3.dll",
    "_internal/PySide6/QtCore.pyd", "_internal/PySide6/QtWidgets.pyd",
    "_internal/PySide6/QtWebEngineCore.pyd", "_internal/PySide6/QtWebEngineProcess.exe",
    "_internal/PySide6/plugins/platforms/qwindows.dll",
    "_internal/PySide6/resources/qtwebengine_resources.pak",
    "_internal/assets/ui/zhenyang-ai-console.html",
    "_internal/assets/ui/console-app.js",
    "_internal/assets/ui/console-responsive.css",
    "_internal/assets/ui/console-review.css",
    "_internal/assets/ui/console-design.css",
    "_internal/assets/ui/proactive-contact.js",
    "_internal/assets/ui/proactive-contact.css",
)

def file_hash(path: Path) -> str:
    result = hashlib.sha256()
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(1024 * 1024), b""):
            result.update(chunk)
    return result.hexdigest()

def fresh_process_environment() -> dict[str, str]:
    environment = {key: value for key, value in os.environ.items()
                   if not key.startswith("_PYI_") and key != "_MEIPASS2"}
    environment["PYINSTALLER_RESET_ENVIRONMENT"] = "1"
    return environment

def _contained_file(root: Path, relative: str) -> Path:
    path = root / relative
    if not relative or Path(relative).is_absolute() or ".." in Path(relative).parts:
        raise RuntimeError(f"Invalid payload path: {relative}")
    if not path.resolve().is_relative_to(root.resolve()):
        raise RuntimeError(f"Payload path escapes its directory: {relative}")
    return path

def validate_installation(root: Path, version: str = "", *, hashes: bool = True) -> dict:
    root = root.resolve()
    manifest = json.loads((root / MANIFEST_NAME).read_text(encoding="utf-8"))
    if version and manifest.get("version") != version:
        raise RuntimeError("Installation payload version does not match the update")
    expected_exe = f"ZhenAI-WeCom-Assistant-v{manifest['version']}.exe"
    if manifest.get("executable") != expected_exe:
        raise RuntimeError("Invalid application executable in installation manifest")
    files = manifest.get("files")
    if not isinstance(files, dict) or not files:
        raise RuntimeError("Installation manifest contains no files")
    for name in (expected_exe, *REQUIRED_FILES):
        if name not in files:
            raise RuntimeError(f"Required runtime missing from payload: {name}")
    if not any(name.startswith("_internal/python3") and name.endswith(".dll") for name in files):
        raise RuntimeError("Python runtime missing from payload")
    for name, expected in files.items():
        path = _contained_file(root, name)
        if not path.is_file() or path.stat().st_size != expected['size']:
            raise RuntimeError(f"Installation file missing or incomplete: {name}")
        if hashes and file_hash(path) != expected['sha256']:
            raise RuntimeError(f"Installation checksum mismatch: {name}")
    return manifest

def write_manifest(root: Path, version: str) -> Path:
    root = root.resolve()
    files = {}
    for path in sorted(root.rglob("*")):
        if path.is_file() and path.name != MANIFEST_NAME:
            files[path.relative_to(root).as_posix()] = {"size": path.stat().st_size, "sha256": file_hash(path)}
    manifest = {"version": version, "executable": f"ZhenAI-WeCom-Assistant-v{version}.exe", "files": files}
    target = root / MANIFEST_NAME
    target.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
    validate_installation(root, version)
    return target

def apply_payload(payload: Path, install_dir: Path, version: str, close_application, *, move_payload: bool = False) -> Path | None:
    payload = payload.resolve()
    install_dir = install_dir.absolute()
    if install_dir == install_dir.parent or install_dir.resolve() != install_dir:
        raise RuntimeError("Installation directory must not be a drive root or directory link")
    if payload.is_relative_to(install_dir) or install_dir.is_relative_to(payload):
        raise RuntimeError("Installer payload must be outside the installation directory")
    validate_installation(payload, version)
    if install_dir.exists() and any(install_dir.iterdir()) and not any(install_dir.glob("ZhenAI-WeCom-Assistant-v*.exe")):
        raise RuntimeError("Refusing to replace a directory that is not an application installation")
    install_dir.parent.mkdir(parents=True, exist_ok=True)
    stage = Path(tempfile.mkdtemp(prefix=f".{install_dir.name}.stage-", dir=install_dir.parent))
    backup = install_dir.with_name(f".{install_dir.name}.backup-{uuid.uuid4().hex[:12]}")
    moved_old = False
    committed = False
    try:
        if move_payload and payload.drive.lower() == install_dir.drive.lower():
            # NSIS has already unpacked onto this volume; rename instead of
            # duplicating more than a gigabyte of runtime files.
            stage.rmdir()
            payload.rename(stage)
        else:
            required = sum(item['size'] for item in validate_installation(payload, version, hashes=False)['files'].values())
            if shutil.disk_usage(install_dir.parent).free < required + 64 * 1024 * 1024:
                raise RuntimeError("Not enough disk space to stage the update; existing installation preserved")
            shutil.copytree(payload, stage, dirs_exist_ok=True)
        validate_installation(stage, version)
        # All slow extraction, copying and checks finish before the old app stops.
        close_application()
        if install_dir.exists():
            install_dir.rename(backup)
            moved_old = True
        try:
            stage.rename(install_dir)
            committed = True
        except BaseException:
            if moved_old:
                backup.rename(install_dir)
            raise
        return backup if moved_old else None
    finally:
        # Delete only the exact temporary directory created by this invocation.
        if not committed and stage.exists():
            assert stage.resolve().parent == install_dir.parent.resolve()
            assert stage.name.startswith(f".{install_dir.name}.stage-")
            shutil.rmtree(stage)