#!/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 [--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()