vps-manager/scripts/restore_migration_bundle.py

"""迁移恢复：校验并解包迁移包到目标部署目录

在新机器运行（前提：已跑 deploy/setup.sh 完成代码/venv/units 部署）：
  sudo systemctl stop vps-manager vps-manager-update.timer
  .venv/bin/python scripts/restore_migration_bundle.py <bundle.tar.gz> [--app-dir /opt/vps-manager]
  sudo systemctl start vps-manager

行为：
  - 按 manifest.sha256 逐文件校验，任一不匹配即拒绝恢复（防传输损坏/调包）
  - 目标机现有 data/ 与 .env 先备份为 *.pre-restore-<ts>（可回滚）
  - 解包 data/（两库 + escrow + backups）与 .env（chmod 600）
  - 打印后续手工步骤（serve/agent/CORS/源机 timer 下线，见 docs/migration-runbook.md）
"""

import argparse
import hashlib
import shutil
import sys
import tarfile
import tempfile
from datetime import datetime
from pathlib import Path


def sha256_file(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def main() -> int:
    ap = argparse.ArgumentParser(description="校验并恢复 vps-manager 迁移包")
    ap.add_argument("bundle", help="迁移包 tar.gz 路径")
    ap.add_argument("--app-dir", default="/opt/vps-manager", help="目标部署目录")
    args = ap.parse_args()

    bundle = Path(args.bundle)
    app_dir = Path(args.app_dir)
    if not bundle.exists():
        print(f"[abort] 迁移包不存在：{bundle}", file=sys.stderr)
        return 2
    if not (app_dir / "app" / "main.py").exists():
        print("[abort] 目标目录不像已部署的 vps-manager（缺 app/main.py），请先跑 deploy/setup.sh",
              file=sys.stderr)
        return 2

    work = Path(tempfile.mkdtemp(prefix="vps-restore-"))
    try:
        with tarfile.open(bundle) as tar:
            try:
                # Python >= 3.11.4/3.12：官方数据过滤器（防路径穿越/硬链接等）
                tar.extractall(work, filter="data")
            except TypeError:
                # 新机自带旧版 Python（如 3.10）无 filter 参数：手工校验成员路径
                for m in tar.getmembers():
                    if m.name.startswith("/") or ".." in Path(m.name).parts:
                        print(f"[abort] 包内路径异常，拒绝解包：{m.name}", file=sys.stderr)
                        return 2
                tar.extractall(work)
        root = work / "bundle"
        manifest_sha = root / "manifest.sha256"
        if not manifest_sha.exists():
            print("[abort] 包内缺 manifest.sha256，拒绝恢复", file=sys.stderr)
            return 2
        print("[1/4] 校验包内文件…")
        listed: set[str] = set()
        for line in manifest_sha.read_text(encoding="utf-8").splitlines():
            if not line.strip():
                continue
            digest, name = line.split("  ", 1)
            listed.add(name)
            target = root / name
            if not target.exists():
                print(f"[abort] 包内缺文件：{name}", file=sys.stderr)
                return 2
            actual = sha256_file(target)
            if actual != digest:
                print(f"[abort] 校验和不匹配：{name}\n  期望 {digest}\n  实际 {actual}",
                      file=sys.stderr)
                return 2
            print(f"  ok {name}")

        # 必备文件预检：校验和只能证明「已列出的文件完好」，发现不了「压根没打进包」。
        # 缺 assets.db/.env 还往下走的话，目标机 data/ 与 .env 已被移走，
        # 恢复出的却是空库，而且还会报「恢复成功」
        missing = [
            name for name in ("data/assets.db", ".env")
            if name not in listed or not (root / name).exists()
        ]
        if missing:
            print(f"[abort] 迁移包缺必备文件：{', '.join(missing)}（尚未触碰目标机状态）",
                  file=sys.stderr)
            return 2

        ts = datetime.now().strftime("%Y%m%d_%H%M%S")
        print("[2/4] 备份目标机现有状态…")
        for src, tag in ((app_dir / "data", f"data.pre-restore-{ts}"),
                         (app_dir / ".env", f".env.pre-restore-{ts}")):
            if src.exists():
                dst = app_dir / tag
                shutil.move(str(src), str(dst))
                print(f"  {src} -> {dst}")

        print("[3/4] 恢复 data/ 与 .env…")
        shutil.move(str(root / "data"), str(app_dir / "data"))
        shutil.move(str(root / ".env"), str(app_dir / ".env"))
        (app_dir / ".env").chmod(0o600)
        escrow = app_dir / "data" / "master_key.escrow"
        if escrow.exists():
            escrow.chmod(0o600)

        print("[4/4] 恢复完成。后续手工步骤（详见 docs/migration-runbook.md）：")
        print("  1. systemctl start vps-manager && curl 127.0.0.1:8000/health 比对 version/commit")
        print("  2. tailscale serve --bg --https=443 http://127.0.0.1:8000（设备名沿用旧名可保持 URL 不变）")
        print("  3. .env 的 CORS_ORIGINS 加入新 HTTPS 域名（若 URL 变化）")
        print("  4. 各被管 VPS 的 /etc/vps-agent.env：VPS_MANAGER_URL 指向新地址后 restart vps-agent")
        print("  5. 确认新机数据无误后，源机 disable 全部 timer 并 stop 服务（防双写/双通知）")
        print("  6. 删除本迁移包与源机上的包副本（等同最高机密）")
        return 0
    finally:
        shutil.rmtree(work, ignore_errors=True)


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