#!/usr/bin/env python3 """Create, verify and restore a hashed Python-environment directory checkpoint.""" from __future__ import annotations import argparse import hashlib import json import os import shutil import tempfile from pathlib import Path class CheckpointError(ValueError): pass def fault(point: str) -> None: if os.environ.get("NEXUS_ALLOW_TEST_FAULTS") == "1" \ and os.environ.get("NEXUS_TEST_CRASH_AT") == point: os._exit(97) def sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as stream: for chunk in iter(lambda: stream.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def inventory(root: Path) -> list[dict]: result = [] if not root.exists(): return result for path in sorted(root.rglob("*")): relative = path.relative_to(root).as_posix() mode = path.lstat().st_mode & 0o7777 if path.is_symlink(): result.append({"path": relative, "type": "symlink", "mode": mode, "target": os.readlink(path)}) elif path.is_file(): result.append({"path": relative, "type": "file", "mode": mode, "bytes": path.stat().st_size, "sha256": sha256(path)}) elif path.is_dir(): result.append({"path": relative, "type": "directory", "mode": mode}) else: raise CheckpointError(f"UNSUPPORTED_ENTRY:{relative}") return result def canonical(value: dict) -> bytes: return (json.dumps(value, sort_keys=True, separators=(",", ":")) + "\n").encode() def fsync_tree(root: Path) -> None: for path in root.rglob("*"): if path.is_file() and not path.is_symlink(): fd = os.open(path, os.O_RDONLY) try: os.fsync(fd) finally: os.close(fd) for path in sorted((p for p in root.rglob("*") if p.is_dir()), reverse=True): fd = os.open(path, os.O_RDONLY) try: os.fsync(fd) finally: os.close(fd) fd = os.open(root, os.O_RDONLY) try: os.fsync(fd) finally: os.close(fd) def create(target: Path, checkpoint: Path, operation_id: str) -> dict: if checkpoint.exists(): raise CheckpointError("CHECKPOINT_ALREADY_EXISTS") checkpoint.parent.mkdir(parents=True, exist_ok=True) temporary = Path(tempfile.mkdtemp(prefix=f".{checkpoint.name}.", dir=checkpoint.parent)) try: payload = temporary / "payload" existed = target.exists() if existed: if not target.is_dir() or target.is_symlink(): raise CheckpointError("TARGET_NOT_DIRECTORY") shutil.copytree(target, payload, symlinks=True) manifest = { "schema": "com.cezen.nexus-one.python-environment-checkpoint.v1", "operation_id": operation_id, "target": str(target), "target_existed": existed, "entries": inventory(payload), } manifest["inventory_sha256"] = hashlib.sha256(canonical(manifest)).hexdigest() (temporary / "manifest.json").write_bytes(canonical(manifest)) fsync_tree(temporary) fault("create_before_replace") os.replace(temporary, checkpoint) fault("create_after_replace") directory_fd = os.open(checkpoint.parent, os.O_RDONLY) try: os.fsync(directory_fd) finally: os.close(directory_fd) return manifest finally: if temporary.exists(): shutil.rmtree(temporary) def verify(checkpoint: Path, expected_target: Path | None = None) -> dict: manifest = json.loads((checkpoint / "manifest.json").read_text()) recorded = manifest.pop("inventory_sha256", None) if recorded != hashlib.sha256(canonical(manifest)).hexdigest(): raise CheckpointError("CHECKPOINT_MANIFEST_INVALID") manifest["inventory_sha256"] = recorded if expected_target is not None and manifest["target"] != str(expected_target): raise CheckpointError("CHECKPOINT_TARGET_MISMATCH") actual = inventory(checkpoint / "payload") if actual != manifest["entries"]: raise CheckpointError("CHECKPOINT_PAYLOAD_INVALID") return manifest def restore(checkpoint: Path, target: Path) -> dict: manifest = verify(checkpoint, target) target.parent.mkdir(parents=True, exist_ok=True) if not manifest["target_existed"]: if target.exists(): quarantine = target.with_name(f".{target.name}.rollback-discard") if quarantine.exists(): raise CheckpointError("ROLLBACK_QUARANTINE_EXISTS") os.replace(target, quarantine) fault("restore_after_quarantine") directory_fd = os.open(target.parent, os.O_RDONLY) try: os.fsync(directory_fd) finally: os.close(directory_fd) return manifest temporary = Path(tempfile.mkdtemp(prefix=f".{target.name}.restore.", dir=target.parent)) shutil.rmtree(temporary) shutil.copytree(checkpoint / "payload", temporary, symlinks=True) fsync_tree(temporary) if target.exists(): quarantine = target.with_name(f".{target.name}.rollback-discard") if quarantine.exists(): raise CheckpointError("ROLLBACK_QUARANTINE_EXISTS") os.replace(target, quarantine) fault("restore_after_quarantine") os.replace(temporary, target) fault("restore_after_replace") directory_fd = os.open(target.parent, os.O_RDONLY) try: os.fsync(directory_fd) finally: os.close(directory_fd) return manifest def recover_restore(checkpoint: Path, target: Path) -> dict: manifest = verify(checkpoint, target) quarantine = target.with_name(f".{target.name}.rollback-discard") if target.exists(): if manifest["target_existed"] and inventory(target) == manifest["entries"]: return manifest if not manifest["target_existed"]: raise CheckpointError("ROLLBACK_TARGET_SHOULD_BE_ABSENT") raise CheckpointError("ROLLBACK_TARGET_AMBIGUOUS") if not manifest["target_existed"]: return manifest if not quarantine.exists(): raise CheckpointError("ROLLBACK_SOURCE_STATE_MISSING") return restore(checkpoint, target) def main() -> int: parser = argparse.ArgumentParser() sub = parser.add_subparsers(dest="command", required=True) create_parser = sub.add_parser("create") create_parser.add_argument("target", type=Path) create_parser.add_argument("checkpoint", type=Path) create_parser.add_argument("--operation-id", required=True) verify_parser = sub.add_parser("verify") verify_parser.add_argument("checkpoint", type=Path) verify_parser.add_argument("--target", type=Path) restore_parser = sub.add_parser("restore") restore_parser.add_argument("checkpoint", type=Path) restore_parser.add_argument("target", type=Path) recover_parser = sub.add_parser("recover-restore") recover_parser.add_argument("checkpoint", type=Path) recover_parser.add_argument("target", type=Path) args = parser.parse_args() try: if args.command == "create": result = create(args.target, args.checkpoint, args.operation_id) elif args.command == "verify": result = verify(args.checkpoint, args.target) elif args.command == "restore": result = restore(args.checkpoint, args.target) else: result = recover_restore(args.checkpoint, args.target) except (CheckpointError, OSError, json.JSONDecodeError) as error: print(json.dumps({"ok": False, "code": str(error)}, sort_keys=True)) return 1 print(json.dumps({"ok": True, "inventory_sha256": result["inventory_sha256"], "entries": len(result["entries"])}, sort_keys=True)) return 0 if __name__ == "__main__": raise SystemExit(main())