geointel/scripts/restore_storage_quarantine.py

#!/usr/bin/env python3
"""Restore a traceable GeoIntel cleanup quarantine without overwriting data."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
from datetime import datetime, timezone
from pathlib import Path


CONFIRMATION = "RESTORE_QUARANTINED_ARTIFACTS"
CLEANUP_PREFIXES = ("exports", "previews", "tiles", "masks", "derived", "rasters/derived")


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


def write_manifest(path: Path, payload: dict[str, object]) -> None:
    temporary = path.with_suffix(".json.partial")
    temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    temporary.replace(path)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--storage-root", type=Path, required=True)
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--confirm", required=True)
    return parser.parse_args()


def _within(path: Path, root: Path, *, label: str) -> Path:
    resolved = path.resolve()
    try:
        resolved.relative_to(root)
    except ValueError as exc:
        raise RuntimeError(f"{label} escapes the storage root") from exc
    return resolved


def main() -> int:
    args = parse_args()
    if args.confirm != CONFIRMATION:
        raise RuntimeError(f"Refusing restore; pass --confirm {CONFIRMATION}")
    storage_root = args.storage_root.expanduser().resolve()
    manifest_path = _within(args.manifest.expanduser(), storage_root, label="Manifest")
    protected_quarantine_root = storage_root / "operator-evidence" / "cleanup-quarantine"
    try:
        manifest_path.relative_to(protected_quarantine_root.resolve())
    except ValueError as exc:
        raise RuntimeError("Manifest is outside the protected cleanup quarantine") from exc
    payload = json.loads(manifest_path.read_text(encoding="utf-8"))
    if payload.get("schema_version") != 1:
        raise RuntimeError("Unsupported quarantine manifest")
    if payload.get("state") not in {"complete", "in_progress", "restore_in_progress", "restored"}:
        raise RuntimeError("Quarantine manifest is not in a restorable state")
    raw_entries = payload.get("entries")
    if not isinstance(raw_entries, list):
        raise RuntimeError("Quarantine manifest entries are invalid")

    plans: list[tuple[str, dict[str, object], Path, Path]] = []
    for raw_entry in raw_entries:
        if not isinstance(raw_entry, dict):
            raise RuntimeError("Quarantine manifest entry is invalid")
        status = raw_entry.get("status")
        if status not in {"planned", "linked", "quarantined", "restore_linked", "restored"}:
            raise RuntimeError(f"Quarantine manifest entry has an invalid status: {status!r}")
        relative_path = raw_entry.get("relative_path")
        quarantine_relative_path = raw_entry.get("quarantine_relative_path")
        expected_hash = raw_entry.get("sha256")
        expected_size = raw_entry.get("size_bytes")
        if (
            not isinstance(relative_path, str)
            or not isinstance(expected_hash, str)
            or not isinstance(expected_size, int)
        ):
            raise RuntimeError("Quarantine manifest entry lacks recovery metadata")
        if not any(
            relative_path == prefix or relative_path.startswith(f"{prefix}/")
            for prefix in CLEANUP_PREFIXES
        ):
            raise RuntimeError(f"Original path is outside the cleanup allowlist: {relative_path}")
        original = _within(storage_root / relative_path, storage_root, label="Original path")
        if not isinstance(quarantine_relative_path, str):
            if status != "planned":
                raise RuntimeError("Quarantine manifest entry lacks its retained path")
            quarantine_relative_path = (
                manifest_path.parent / "files" / relative_path
            ).relative_to(storage_root).as_posix()
            raw_entry["quarantine_relative_path"] = quarantine_relative_path
        quarantined = _within(
            storage_root / quarantine_relative_path,
            storage_root,
            label="Quarantine path",
        )
        try:
            quarantined.relative_to(manifest_path.parent.resolve())
        except ValueError as exc:
            raise RuntimeError("Quarantine entry escapes its operation directory") from exc
        original_exists = original.exists()
        quarantined_exists = quarantined.exists()
        if original_exists:
            if not original.is_file() or original.stat().st_size != expected_size or sha256(original) != expected_hash:
                raise RuntimeError(f"Restore destination already exists with different bytes: {relative_path}")
        if quarantined_exists:
            if (
                not quarantined.is_file()
                or quarantined.stat().st_size != expected_size
                or sha256(quarantined) != expected_hash
            ):
                raise RuntimeError(f"Quarantined artifact checksum mismatch: {quarantine_relative_path}")
        if original_exists and quarantined_exists:
            if not os.path.samefile(original, quarantined):
                raise RuntimeError(f"Restore destination already exists: {relative_path}")
            plans.append(("remove_duplicate_link", raw_entry, quarantined, original))
        elif original_exists:
            plans.append(("mark_restored", raw_entry, quarantined, original))
        elif quarantined_exists:
            plans.append(("restore", raw_entry, quarantined, original))
        else:
            raise RuntimeError(f"Both original and quarantined artifacts are missing: {relative_path}")

    payload["state"] = "restore_in_progress"
    write_manifest(manifest_path, payload)
    for action, entry, quarantined, original in plans:
        if action == "restore":
            original.parent.mkdir(parents=True, exist_ok=True)
            _within(original, storage_root, label="Original path")
            try:
                os.link(quarantined, original, follow_symlinks=False)
            except FileExistsError as exc:
                raise RuntimeError(f"Restore destination was created concurrently: {original}") from exc
            if not os.path.samefile(quarantined, original):
                original.unlink(missing_ok=True)
                raise RuntimeError(f"Restore link verification failed: {original}")
            entry["status"] = "restore_linked"
            write_manifest(manifest_path, payload)
            quarantined.unlink()
        elif action == "remove_duplicate_link":
            quarantined.unlink()
        entry["status"] = "restored"
        entry["restored_at"] = datetime.now(timezone.utc).isoformat()
        write_manifest(manifest_path, payload)
    payload["state"] = "restored"
    payload["restored_at"] = datetime.now(timezone.utc).isoformat()
    write_manifest(manifest_path, payload)
    print(json.dumps({"state": "restored", "restored_count": len(plans)}, sort_keys=True))
    return 0


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