scripts/restore_from_dataset_atomic.py

#!/usr/bin/env python3
"""Atomic Dataset restore for HermesFace secure state snapshots."""
from __future__ import annotations

import hashlib
import json
import logging
import shutil
import sys
import tempfile
import time
from datetime import datetime, timezone
from pathlib import Path, PurePosixPath
from typing import Any, Dict, List, Optional, Sequence

from huggingface_hub import HfApi, hf_hub_download
from huggingface_hub.utils import EntryNotFoundError

logging.basicConfig(
    level=logging.INFO,
    format='{"timestamp": "%(asctime)s", "level": "%(levelname)s", "module": "atomic-restore", "message": "%(message)s"}',
)
logger = logging.getLogger(__name__)


def _safe_dataset_path(value: str) -> str:
    pure = PurePosixPath(str(value).replace("\\", "/"))
    if pure.is_absolute() or not pure.parts or any(part in {"", ".", ".."} for part in pure.parts):
        raise ValueError(f"unsafe dataset path: {value!r}")
    return pure.as_posix()


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 _assert_safe_target(target: Path, protected_paths: Sequence[Path]) -> None:
    resolved_target = target.resolve()
    for protected in protected_paths:
        resolved_protected = Path(protected).resolve()
        if (
            resolved_target == resolved_protected
            or resolved_target in resolved_protected.parents
            or resolved_protected in resolved_target.parents
        ):
            raise RuntimeError(
                f"refusing restore target overlapping application source: {resolved_target}"
            )


class AtomicDatasetRestorer:
    def __init__(
        self,
        repo_id: str,
        dataset_path: str = "state",
        *,
        api: Optional[HfApi] = None,
        token: Optional[str] = None,
    ):
        self.repo_id = repo_id
        self.dataset_path = _safe_dataset_path(dataset_path)
        self.token = token
        self.api = api or HfApi(token=token)

    def calculate_checksum(self, file_path: Path) -> str:
        return _sha256(file_path)

    def validate_integrity(self, metadata: Dict[str, Any], state_files: List[Path]) -> bool:
        """Legacy metadata validation retained for compatibility."""
        try:
            if "checksum" not in metadata:
                logger.warning("no_checksum_in_metadata")
                return True
            calculated = hashlib.sha256(
                json.dumps(metadata.get("state_data", {}), sort_keys=True).encode("utf-8")
            ).hexdigest()
            expected = metadata["checksum"]
            valid = calculated == expected
            logger.info(f"integrity_check expected={expected} calculated={calculated} valid={valid}")
            return valid
        except Exception as exc:
            logger.error(f"integrity_validation_failed error={exc}")
            return False

    def create_backup_before_restore(self, target_dir: Path) -> Optional[Path]:
        """Legacy full-directory copy retained for callers of restore_from_commit."""
        try:
            if not target_dir.exists():
                return None
            timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
            backup_dir = target_dir.parent / f"state_backup_{timestamp}"
            shutil.copytree(target_dir, backup_dir)
            return backup_dir
        except Exception as exc:
            logger.error(f"local_backup_failed error={exc}")
            return None

    def restore_secure_from_commit(
        self,
        commit_sha: str,
        target_dir: Path,
        *,
        last_known_good_path: Optional[Path] = None,
        protected_paths: Sequence[Path] = (),
    ) -> Dict[str, Any]:
        """Verify a secure snapshot completely before atomically activating it."""
        try:
            from secure_backup import build_backup, extract_validated, restore_backup
        except ImportError:
            from scripts.secure_backup import (  # type: ignore
                build_backup,
                extract_validated,
                restore_backup,
            )

        operation_id = f"restore_{int(time.time())}"
        target = Path(target_dir)
        _assert_safe_target(target, protected_paths)
        if target.exists() and target.is_symlink():
            raise RuntimeError("refusing symlink restore target")

        had_existing_target = target.exists()
        try:
            self.api.repo_info(
                repo_id=self.repo_id,
                repo_type="dataset",
                revision=commit_sha,
            )
            archive_remote = f"{self.dataset_path}/secure_state.tar.gz"
            metadata_remote = f"{self.dataset_path}/metadata.json"

            with tempfile.TemporaryDirectory(prefix="hermes-remote-restore-") as tmpdir:
                temp_root = Path(tmpdir)
                downloaded_archive = Path(
                    hf_hub_download(
                        repo_id=self.repo_id,
                        repo_type="dataset",
                        filename=archive_remote,
                        revision=commit_sha,
                        token=self.token,
                    )
                )
                downloaded_metadata = Path(
                    hf_hub_download(
                        repo_id=self.repo_id,
                        repo_type="dataset",
                        filename=metadata_remote,
                        revision=commit_sha,
                        token=self.token,
                    )
                )
                archive = temp_root / "secure_state.tar.gz"
                metadata_path = temp_root / "metadata.json"
                shutil.copy2(downloaded_archive, archive)
                shutil.copy2(downloaded_metadata, metadata_path)

                metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
                if metadata.get("format") != "hermes-secure-backup" or metadata.get("format_version") != 1:
                    raise RuntimeError("unsupported or missing secure backup metadata")
                if metadata.get("archive") != "secure_state.tar.gz":
                    raise RuntimeError("secure backup metadata references an unexpected archive")
                actual_archive_sha = _sha256(archive)
                if actual_archive_sha != metadata.get("archive_sha256"):
                    raise RuntimeError("secure backup archive SHA-256 mismatch")

                # Explicit staging validation occurs before any live state or LKG
                # snapshot is touched.
                staged = temp_root / "staged"
                manifest = extract_validated(archive, staged)
                canonical_manifest_sha = hashlib.sha256(
                    json.dumps(manifest, sort_keys=True, separators=(",", ":")).encode("utf-8")
                ).hexdigest()
                if canonical_manifest_sha != metadata.get("manifest_sha256"):
                    raise RuntimeError("secure backup manifest SHA-256 mismatch")
                if manifest != metadata.get("manifest"):
                    raise RuntimeError("secure backup manifest metadata mismatch")

                lkg_path = Path(last_known_good_path) if last_known_good_path else target.parent / ".hermes-last-known-good.tar.gz"
                resolved_lkg_parent = lkg_path.parent.resolve()
                resolved_target = target.resolve()
                if resolved_lkg_parent == resolved_target or resolved_target in resolved_lkg_parent.parents:
                    raise RuntimeError("last-known-good snapshot must be outside the restore target")
                if target.exists():
                    build_backup(target, lkg_path)

                activation = restore_backup(archive, target)
                result = {
                    "success": True,
                    "operation_id": operation_id,
                    "commit_sha": commit_sha,
                    "timestamp": datetime.now(timezone.utc).isoformat(),
                    "archive_sha256": actual_archive_sha,
                    "last_known_good": str(lkg_path) if had_existing_target else None,
                    "activation": activation,
                }
                logger.info(f"atomic_secure_restore_completed op={operation_id} commit={commit_sha}")
                return result
        except EntryNotFoundError as exc:
            # Expected on the very first restore attempt against a freshly
            # created (or not-yet-saved-to) dataset repo: the commit exists
            # but no secure_state.tar.gz has ever been uploaded to it yet.
            # This is a normal "starting fresh" condition, not a failure, so
            # it is logged at INFO rather than ERROR.
            logger.info(
                f"atomic_secure_restore_no_prior_backup op={operation_id} "
                f"commit={commit_sha} detail={exc}"
            )
            return {
                "success": False,
                "operation_id": operation_id,
                "commit_sha": commit_sha,
                "error": str(exc),
                "no_prior_backup": True,
                "timestamp": datetime.now(timezone.utc).isoformat(),
            }
        except Exception as exc:
            logger.error(f"atomic_secure_restore_failed op={operation_id} error={exc}")
            return {
                "success": False,
                "operation_id": operation_id,
                "commit_sha": commit_sha,
                "error": str(exc),
                "timestamp": datetime.now(timezone.utc).isoformat(),
            }

    def restore_secure_latest(
        self,
        target_dir: Path,
        *,
        last_known_good_path: Optional[Path] = None,
        protected_paths: Sequence[Path] = (),
    ) -> Dict[str, Any]:
        try:
            repo_info = self.api.repo_info(repo_id=self.repo_id, repo_type="dataset")
            if not repo_info.sha:
                return {
                    "success": False,
                    "error": "No commit found in repository",
                    "timestamp": datetime.now(timezone.utc).isoformat(),
                }
            return self.restore_secure_from_commit(
                repo_info.sha,
                Path(target_dir),
                last_known_good_path=last_known_good_path,
                protected_paths=protected_paths,
            )
        except Exception as exc:
            return {
                "success": False,
                "error": f"Failed to get latest commit: {exc}",
                "timestamp": datetime.now(timezone.utc).isoformat(),
            }

    def restore_from_commit(
        self, commit_sha: str, target_dir: Path, force: bool = False
    ) -> Dict[str, Any]:
        """Legacy flat-file restore retained for compatibility only."""
        operation_id = f"restore_{int(time.time())}"
        try:
            self.api.repo_info(repo_id=self.repo_id, repo_type="dataset", revision=commit_sha)
        except Exception as exc:
            return {
                "success": False,
                "operation_id": operation_id,
                "error": f"Invalid commit: {exc}",
                "timestamp": datetime.now(timezone.utc).isoformat(),
            }

        backup_dir = self.create_backup_before_restore(target_dir)
        try:
            files = self.api.list_repo_files(
                repo_id=self.repo_id, repo_type="dataset", revision=commit_sha
            )
            prefix = f"{self.dataset_path}/"
            state_files = [path for path in files if path.startswith(prefix)]
            if not state_files:
                raise RuntimeError("No state files found in commit")

            downloaded_files: List[Path] = []
            metadata = None
            for remote_path in state_files:
                local = hf_hub_download(
                    repo_id=self.repo_id,
                    repo_type="dataset",
                    filename=remote_path,
                    revision=commit_sha,
                    token=self.token,
                )
                downloaded_files.append(Path(local))
                if remote_path.endswith("metadata.json"):
                    metadata = json.loads(Path(local).read_text(encoding="utf-8"))
            if not metadata or not self.validate_integrity(metadata, downloaded_files):
                raise RuntimeError("Data integrity validation failed")

            target_dir.mkdir(parents=True, exist_ok=True)
            restored = []
            for source in downloaded_files:
                if source.name != "metadata.json":
                    destination = target_dir / source.name
                    shutil.copy2(source, destination)
                    restored.append(str(destination))
            return {
                "success": True,
                "operation_id": operation_id,
                "commit_sha": commit_sha,
                "backup_dir": str(backup_dir) if backup_dir else None,
                "timestamp": datetime.now(timezone.utc).isoformat(),
                "restored_files": restored,
                "metadata": metadata,
            }
        except Exception as exc:
            return {
                "success": False,
                "operation_id": operation_id,
                "error": str(exc),
                "timestamp": datetime.now(timezone.utc).isoformat(),
            }

    def restore_latest(self, target_dir: Path, force: bool = False) -> Dict[str, Any]:
        try:
            repo_info = self.api.repo_info(repo_id=self.repo_id, repo_type="dataset")
            if not repo_info.sha:
                return {"success": False, "error": "No commit found in repository"}
            return self.restore_from_commit(repo_info.sha, target_dir, force)
        except Exception as exc:
            return {"success": False, "error": f"Failed to get latest commit: {exc}"}


def main() -> None:
    if len(sys.argv) < 3:
        print(json.dumps({
            "error": "Usage: python restore_from_dataset_atomic.py <repo_id> <target_dir> [--force]",
            "status": "error",
        }, indent=2))
        sys.exit(1)

    repo_id = sys.argv[1]
    target_dir = Path(sys.argv[2])
    force = "--force" in sys.argv
    try:
        result = AtomicDatasetRestorer(repo_id).restore_latest(target_dir, force)
        print(json.dumps(result, indent=2))
        if not result.get("success", False):
            sys.exit(1)
    except Exception as exc:
        print(json.dumps({"error": str(exc), "status": "error"}, indent=2))
        sys.exit(1)


if __name__ == "__main__":
    main()