114 lines
4.3 KiB
Python
114 lines
4.3 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import shutil
|
|
import sqlite3
|
|
import tempfile
|
|
import zipfile
|
|
from collections.abc import Sequence
|
|
from contextlib import closing
|
|
from pathlib import Path
|
|
|
|
from tools.backup import sha256
|
|
|
|
|
|
def _safe_members(archive: zipfile.ZipFile) -> list[zipfile.ZipInfo]:
|
|
members = archive.infolist()
|
|
for member in members:
|
|
path = Path(member.filename)
|
|
if path.is_absolute() or ".." in path.parts:
|
|
raise ValueError("backup contains an unsafe path")
|
|
return members
|
|
|
|
|
|
def restore_backup(
|
|
archive_path: Path,
|
|
database: Path,
|
|
*,
|
|
private_skills: Path | None = None,
|
|
environment_file: Path | None = None,
|
|
encryption_key_file: Path | None = None,
|
|
overwrite: bool = False,
|
|
) -> dict:
|
|
archive_path = archive_path.resolve()
|
|
database = database.resolve()
|
|
destinations = [database]
|
|
destinations.extend(path.resolve() for path in (environment_file, encryption_key_file) if path)
|
|
if private_skills:
|
|
destinations.append(private_skills.resolve())
|
|
if not overwrite and any(path.exists() for path in destinations):
|
|
raise FileExistsError(
|
|
"restore destination exists; explicit overwrite confirmation is required"
|
|
)
|
|
with tempfile.TemporaryDirectory(prefix="xiaobai-restore-") as raw:
|
|
staging = Path(raw)
|
|
with zipfile.ZipFile(archive_path) as archive:
|
|
_safe_members(archive)
|
|
archive.extractall(staging)
|
|
manifest = json.loads((staging / "manifest.json").read_text(encoding="utf-8"))
|
|
if manifest.get("format") != 1 or not isinstance(manifest.get("files"), dict):
|
|
raise ValueError("backup manifest is invalid")
|
|
for name, metadata in manifest["files"].items():
|
|
path = staging / name
|
|
if not path.is_file() or sha256(path) != metadata.get("sha256"):
|
|
raise ValueError(f"backup checksum mismatch: {name}")
|
|
restored_db = staging / "database.sqlite3"
|
|
with closing(
|
|
sqlite3.connect(f"file:{restored_db.as_posix()}?mode=ro", uri=True)
|
|
) as connection:
|
|
if connection.execute("PRAGMA integrity_check").fetchone()[0] != "ok":
|
|
raise ValueError("restored database failed integrity check")
|
|
database.parent.mkdir(parents=True, exist_ok=True)
|
|
shutil.copy2(restored_db, database)
|
|
for destination, name in (
|
|
(environment_file, "secrets/environment"),
|
|
(encryption_key_file, "secrets/app-encryption.key"),
|
|
):
|
|
source = staging / name
|
|
if destination and source.is_file():
|
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
|
shutil.copy2(source, destination)
|
|
skills_source = staging / "private-mentor-skills"
|
|
if private_skills and skills_source.is_dir():
|
|
private_skills.mkdir(parents=True, exist_ok=True)
|
|
shutil.copytree(skills_source, private_skills, dirs_exist_ok=True)
|
|
return manifest
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(description="Validate and restore a Xiaobai backup")
|
|
parser.add_argument("--archive", type=Path, required=True)
|
|
parser.add_argument("--database", type=Path, required=True)
|
|
parser.add_argument("--private-skills", type=Path)
|
|
parser.add_argument("--environment-file", type=Path)
|
|
parser.add_argument("--encryption-key-file", type=Path)
|
|
parser.add_argument("--confirm-restore", action="store_true")
|
|
parser.add_argument("--confirm-overwrite", action="store_true")
|
|
return parser
|
|
|
|
|
|
def main(arguments: Sequence[str] | None = None) -> int:
|
|
parsed = build_parser().parse_args(arguments)
|
|
if not parsed.confirm_restore:
|
|
raise SystemExit("restore requires --confirm-restore")
|
|
manifest = restore_backup(
|
|
parsed.archive,
|
|
parsed.database,
|
|
private_skills=parsed.private_skills,
|
|
environment_file=parsed.environment_file,
|
|
encryption_key_file=parsed.encryption_key_file,
|
|
overwrite=parsed.confirm_overwrite,
|
|
)
|
|
print(
|
|
json.dumps(
|
|
{"restored": str(parsed.database.resolve()), "files": len(manifest["files"])},
|
|
ensure_ascii=False,
|
|
)
|
|
)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|