from __future__ import annotations import argparse import hashlib import json import sqlite3 import tempfile import zipfile from collections.abc import Sequence from contextlib import closing from datetime import datetime from pathlib import Path def sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: for block in iter(lambda: handle.read(1024 * 1024), b""): digest.update(block) return digest.hexdigest() def snapshot_database(source: Path, destination: Path) -> None: source_uri = f"file:{source.resolve().as_posix()}?mode=ro" with closing(sqlite3.connect(source_uri, uri=True)) as origin: with closing(sqlite3.connect(destination)) as snapshot: origin.backup(snapshot) if snapshot.execute("PRAGMA integrity_check").fetchone()[0] != "ok": raise RuntimeError("database snapshot failed integrity check") def create_backup( database: Path, output: Path, *, private_skills: Path | None = None, environment_file: Path | None = None, encryption_key_file: Path | None = None, ) -> dict: database = database.resolve() if not database.is_file(): raise FileNotFoundError(database) output = output.resolve() output.parent.mkdir(parents=True, exist_ok=True) with tempfile.TemporaryDirectory(prefix="xiaobai-backup-") as raw: staging = Path(raw) database_copy = staging / "database.sqlite3" snapshot_database(database, database_copy) files: list[tuple[Path, str]] = [(database_copy, "database.sqlite3")] for source, archive_name in ( (environment_file, "secrets/environment"), (encryption_key_file, "secrets/app-encryption.key"), ): if source: resolved = source.resolve() if not resolved.is_file(): raise FileNotFoundError(resolved) files.append((resolved, archive_name)) if private_skills and private_skills.exists(): files.extend( (item, f"private-mentor-skills/{item.relative_to(private_skills).as_posix()}") for item in private_skills.rglob("*") if item.is_file() ) manifest = { "format": 1, "created_at": datetime.now().astimezone().isoformat(timespec="seconds"), "files": { name: {"sha256": sha256(path), "size": path.stat().st_size} for path, name in files }, } manifest_path = staging / "manifest.json" manifest_path.write_text( json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8" ) temporary = output.with_suffix(output.suffix + ".tmp") with zipfile.ZipFile(temporary, "w", compression=zipfile.ZIP_DEFLATED) as archive: archive.write(manifest_path, "manifest.json") for path, name in files: archive.write(path, name) temporary.replace(output) return manifest def apply_retention(directory: Path, *, keep_daily: int = 7, keep_weekly: int = 4) -> list[Path]: backups = sorted( directory.glob("xiaobai-*.zip"), key=lambda item: item.stat().st_mtime, reverse=True ) keep = set(backups[:keep_daily]) weeks: set[str] = set() for backup in backups: week = datetime.fromtimestamp(backup.stat().st_mtime).strftime("%G-%V") if week not in weeks and len(weeks) < keep_weekly: weeks.add(week) keep.add(backup) removed = [] for backup in backups: if backup not in keep: backup.unlink() removed.append(backup) return removed def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description="Create a consistent Xiaobai backup archive") parser.add_argument("--database", type=Path, required=True) parser.add_argument("--output", 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("--apply-retention", action="store_true") return parser def main(arguments: Sequence[str] | None = None) -> int: parsed = build_parser().parse_args(arguments) manifest = create_backup( parsed.database, parsed.output, private_skills=parsed.private_skills, environment_file=parsed.environment_file, encryption_key_file=parsed.encryption_key_file, ) removed = apply_retention(parsed.output.parent) if parsed.apply_retention else [] print( json.dumps( { "backup": str(parsed.output.resolve()), "files": len(manifest["files"]), "removed_by_retention": len(removed), }, ensure_ascii=False, ) ) return 0 if __name__ == "__main__": raise SystemExit(main())