dp-zp-agent/tools/evaluation/archive.py

122 lines
4.8 KiB
Python

"""Verify historical experiments or restore them to an explicit directory.
No experiment is regenerated and no historical fingerprint is rewritten.
Restore uses the reorganized experiment paths listed in catalog.json.
"""
from __future__ import annotations
import argparse
import hashlib
import json
from pathlib import Path, PurePosixPath
import tarfile
EXPERIMENTS = Path(__file__).resolve().parents[2] / "experiments"
def digest(content: bytes) -> str:
return hashlib.sha256(content).hexdigest()
def safe_path(base: Path, name: str) -> Path:
relative = PurePosixPath(name)
if not name or relative.is_absolute() or ".." in relative.parts or "\\" in name:
raise ValueError(f"Unsafe archive path: {name}")
target = base / relative
if not target.resolve().is_relative_to(base.resolve()):
raise ValueError(f"Archive path escapes destination: {name}")
return target
def load_catalog(experiments: Path = EXPERIMENTS) -> dict:
catalog = json.loads((experiments / "catalog.json").read_text())
if catalog.get("schema_version") != 1:
raise ValueError("Unsupported archive catalog")
return catalog
def archive_blobs(catalog: dict, experiments: Path) -> dict[str, bytes]:
archive = safe_path(experiments, catalog["archive"])
if digest(archive.read_bytes()) != catalog["archive_sha256"]:
raise ValueError("Archive checksum mismatch")
expected = {e["member"] for e in catalog["entries"] if e["location"] == "archive"}
with tarfile.open(archive) as stream:
members = stream.getmembers()
if {m.name for m in members} != expected or len(members) != len(expected):
raise ValueError("Archive member list mismatch")
if any(not member.isfile() for member in members):
raise ValueError("Only regular archive files are supported")
return {member.name: stream.extractfile(member).read() for member in members}
def verify(experiments: Path = EXPERIMENTS) -> dict:
catalog = load_catalog(experiments)
blobs = archive_blobs(catalog, experiments)
for entry in catalog["entries"]:
content = (
blobs[entry["member"]]
if entry["location"] == "archive"
else safe_path(experiments.parent, entry["path"]).read_bytes()
)
if len(content) != entry["bytes"] or digest(content) != entry["sha256"]:
raise ValueError(f"Experiment checksum mismatch: {entry['original_path']}")
return {
"entries": len(catalog["entries"]),
"archived": sum(e["location"] == "archive" for e in catalog["entries"]),
"source_commit": catalog["source_commit"],
}
def restore(
destination: Path, originals: list[str] | None = None, *, experiments: Path = EXPERIMENTS
) -> list[Path]:
"""Validate the complete selection before writing; never overwrite a file."""
catalog = load_catalog(experiments)
wanted = set(originals or [])
known = {entry["original_path"] for entry in catalog["entries"]}
if wanted - known:
raise ValueError("Unknown original experiment path")
blobs = archive_blobs(catalog, experiments)
plan = {}
for entry in catalog["entries"]:
if entry["location"] != "archive" or (wanted and entry["original_path"] not in wanted):
continue
target = safe_path(destination, entry["path"])
content = blobs[entry["member"]]
if len(content) != entry["bytes"] or digest(content) != entry["sha256"]:
raise ValueError("Archived experiment checksum mismatch")
if target.exists():
raise FileExistsError(f"Restore refuses to replace: {target}")
if target in plan and plan[target] != content:
raise ValueError("Conflicting restore destinations")
plan[target] = content
for target, content in plan.items():
target.parent.mkdir(parents=True, exist_ok=True)
with target.open("xb") as stream:
stream.write(content)
return list(plan)
def main():
parser = argparse.ArgumentParser(description=__doc__)
commands = parser.add_subparsers(dest="command", required=True)
commands.add_parser("verify")
commands.add_parser("list")
restore_parser = commands.add_parser("restore")
restore_parser.add_argument("--destination", required=True, type=Path)
restore_parser.add_argument("--original-path", action="append")
args = parser.parse_args()
if args.command == "verify":
print(json.dumps(verify()))
elif args.command == "list":
for entry in load_catalog()["entries"]:
print(f"{entry['location']:7} {entry['original_path']} -> {entry['path']}")
else:
restored = restore(args.destination, args.original_path)
print(json.dumps({"restored": len(restored), "destination": str(args.destination)}))
if __name__ == "__main__":
main()