216 lines
7.9 KiB
Python
216 lines
7.9 KiB
Python
#!/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())
|