Create replayproof_receipts.py

replayproof_receipts.py

@@ -0,0 +1,423 @@
+from __future__ import annotations
+
+import os
+import json
+import base64
+import zipfile
+import hashlib
+import tempfile
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import Dict, Any, List, Optional, Tuple
+
+from cryptography.hazmat.primitives.asymmetric.ed25519 import (
+    Ed25519PrivateKey,
+    Ed25519PublicKey,
+)
+from cryptography.hazmat.primitives import serialization
+
+from replayproof_sim import (
+    SimConfig,
+    SimState,
+    reset_sim,
+    step_sim,
+    render_world_image,
+    render_pov_image,
+    observation_sha256,
+)
+from replayproof_media import export_gif, sha256_file
+
+EVENT_SPEC = "replayproof-event-v0"
+ROOT_SPEC = "replayproof-root-v0"
+
+
+def now_utc_iso() -> str:
+    return datetime.now(timezone.utc).isoformat()
+
+
+def canon(obj) -> bytes:
+    return json.dumps(obj, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8")
+
+
+def sha256_hex(b: bytes) -> str:
+    return hashlib.sha256(b).hexdigest()
+
+
+def pubkey_id(pk_bytes: bytes) -> str:
+    return sha256_hex(pk_bytes)[:16]
+
+
+def b64e(b: bytes) -> str:
+    return base64.b64encode(b).decode("ascii")
+
+
+def b64d(s: str) -> bytes:
+    return base64.b64decode(s.encode("ascii"))
+
+
+@dataclass
+class Recorder:
+    session_id: str
+    cfg: Dict[str, Any]
+    seed: int
+    sk: Ed25519PrivateKey
+    pk: Ed25519PublicKey
+    pk_bytes: bytes
+    pk_id: str
+    events: List[Dict[str, Any]]
+
+    @staticmethod
+    def new_session(sim_spec: str, cfg: Dict[str, Any], seed: int) -> "Recorder":
+        sk = Ed25519PrivateKey.generate()
+        pk = sk.public_key()
+
+        pk_bytes = pk.public_bytes(
+            encoding=serialization.Encoding.Raw,
+            format=serialization.PublicFormat.Raw,
+        )
+        pid = pubkey_id(pk_bytes)
+
+        return Recorder(
+            session_id=os.urandom(16).hex(),
+            cfg={"sim_spec": sim_spec, "cfg": cfg},
+            seed=int(seed),
+            sk=sk,
+            pk=pk,
+            pk_bytes=pk_bytes,
+            pk_id=pid,
+            events=[],
+        )
+
+    def _sign(self, event_sha256_hex: str) -> str:
+        sig = self.sk.sign(bytes.fromhex(event_sha256_hex))
+        return b64e(sig)
+
+    def record_event(self, state: SimState, action: str) -> Dict[str, Any]:
+        prev = self.events[-1]["event_sha256"] if self.events else None
+
+        body = {
+            "spec": EVENT_SPEC,
+            "session_id": self.session_id,
+            "ts": now_utc_iso(),
+            "step": int(state.step),
+            "action": str(action),
+            "seed": int(self.seed),
+            "cfg_sha256": sha256_hex(canon(self.cfg)),
+            "obs_sha256": observation_sha256(state),
+            "state_sha256": state.last_state_sha256,
+            "prev_event_sha256": prev,
+            "pubkey_id": self.pk_id,
+        }
+
+        eh = sha256_hex(canon(body))
+        event = dict(body)
+        event["event_sha256"] = eh
+        event["sig"] = self._sign(eh)
+
+        self.events.append(event)
+        return event
+
+    def root_manifest(self, final_state: SimState, media_sha256: Optional[str] = None) -> Dict[str, Any]:
+        first = self.events[0]["event_sha256"] if self.events else None
+        last = self.events[-1]["event_sha256"] if self.events else None
+
+        root_body = {
+            "spec": ROOT_SPEC,
+            "session_id": self.session_id,
+            "created_utc": now_utc_iso(),
+            "seed": int(self.seed),
+            "cfg": self.cfg,
+            "pubkey": b64e(self.pk_bytes),
+            "pubkey_id": self.pk_id,
+            "events_count": len(self.events),
+            "first_event_sha256": first,
+            "final_event_sha256": last,
+            "final_state_sha256": final_state.last_state_sha256,
+            "media_sha256": media_sha256,
+        }
+
+        root_sha = sha256_hex(canon(root_body))
+        root_sig = b64e(self.sk.sign(bytes.fromhex(root_sha)))
+
+        root = dict(root_body)
+        root["root_sha256"] = root_sha
+        root["root_sig"] = root_sig
+        return root
+
+
+def _verify_event_signature(event: Dict[str, Any], pk: Ed25519PublicKey) -> Tuple[bool, str]:
+    try:
+        sig = b64d(event["sig"])
+        eh = event["event_sha256"]
+        pk.verify(sig, bytes.fromhex(eh))
+        return True, "ok"
+    except Exception as e:
+        return False, f"signature_failed: {type(e).__name__}: {e}"
+
+
+def _recompute_event_hash(event: Dict[str, Any]) -> str:
+    body = dict(event)
+    body.pop("sig", None)
+    body.pop("event_sha256", None)
+    return sha256_hex(canon(body))
+
+
+def build_receipt_bundle_zip(
+    zip_path: str,
+    recorder: Recorder,
+    cfg: SimConfig,
+    final_state: SimState,
+    include_gif: bool = True,
+    watermark: bool = True,
+) -> str:
+    os.makedirs(os.path.dirname(zip_path), exist_ok=True)
+
+    gif_path = None
+    media_sha = None
+
+    if include_gif:
+        # Replay frames deterministically for a POV clip
+        frames = []
+        st = reset_sim(cfg, recorder.seed)
+        frames.append(render_pov_image(st))
+
+        for ev in recorder.events:
+            if ev["action"] == "RESET":
+                continue
+            st, _ = step_sim(cfg, st)  # deterministic policy matches recorded run
+            frames.append(render_pov_image(st))
+            if st.done:
+                break
+
+        wm = None
+        if watermark:
+            wm = f"final:{(final_state.last_state_sha256 or '')[:12]}"
+
+        gif_path = os.path.join(os.path.dirname(zip_path), f"run_{recorder.session_id}.gif")
+        export_gif(frames, gif_path, fps=10, watermark=wm)
+        media_sha = sha256_file(gif_path)
+
+    root = recorder.root_manifest(final_state=final_state, media_sha256=media_sha)
+
+    events_jsonl = "\n".join(json.dumps(e, ensure_ascii=False) for e in recorder.events) + ("\n" if recorder.events else "")
+    root_json = json.dumps(root, ensure_ascii=False, indent=2)
+
+    with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as z:
+        z.writestr("events.jsonl", events_jsonl)
+        z.writestr("root.json", root_json)
+        if gif_path and os.path.exists(gif_path):
+            z.write(gif_path, arcname=os.path.basename(gif_path))
+
+    return zip_path
+
+
+def _load_bundle(zip_path: str) -> Dict[str, Any]:
+    with zipfile.ZipFile(zip_path, "r") as z:
+        names = set(z.namelist())
+        if "root.json" not in names or "events.jsonl" not in names:
+            raise ValueError("bundle_missing_root_or_events")
+
+        root = json.loads(z.read("root.json").decode("utf-8"))
+        events_lines = z.read("events.jsonl").decode("utf-8").splitlines()
+        events = [json.loads(line) for line in events_lines if line.strip()]
+
+        media_name = None
+        for n in names:
+            if n.lower().endswith(".gif") or n.lower().endswith(".mp4"):
+                media_name = n
+                break
+        media_bytes = z.read(media_name) if media_name else None
+
+    return {"root": root, "events": events, "media_name": media_name, "media_bytes": media_bytes}
+
+
+def verify_receipt_bundle_zip(zip_path: str) -> Dict[str, Any]:
+    report: Dict[str, Any] = {
+        "ok": False,
+        "zip": os.path.basename(zip_path),
+        "checks": {},
+        "errors": [],
+    }
+
+    try:
+        bundle = _load_bundle(zip_path)
+        root = bundle["root"]
+        events = bundle["events"]
+        media_name = bundle["media_name"]
+        media_bytes = bundle["media_bytes"]
+
+        report["checks"]["has_root_and_events"] = True
+
+        # Root signature verification
+        pk_bytes = b64d(root["pubkey"])
+        pk = Ed25519PublicKey.from_public_bytes(pk_bytes)
+
+        root_body = dict(root)
+        root_sig = root_body.pop("root_sig", None)
+        root_sha = root_body.pop("root_sha256", None)
+
+        recomputed_root_sha = sha256_hex(canon(root_body))
+        report["checks"]["root_sha_match"] = (root_sha == recomputed_root_sha)
+        if root_sha != recomputed_root_sha:
+            report["errors"].append("root_hash_mismatch")
+
+        try:
+            pk.verify(b64d(root_sig), bytes.fromhex(root_sha))
+            report["checks"]["root_sig_valid"] = True
+        except Exception as e:
+            report["checks"]["root_sig_valid"] = False
+            report["errors"].append(f"root_sig_invalid: {type(e).__name__}: {e}")
+
+        # Event chain + signatures + hash recomputation
+        prev = None
+        chain_ok = True
+        sig_ok = True
+        hash_ok = True
+
+        for ev in events:
+            if ev.get("prev_event_sha256") != prev:
+                chain_ok = False
+
+            reh = _recompute_event_hash(ev)
+            if reh != ev.get("event_sha256"):
+                hash_ok = False
+
+            ok, _ = _verify_event_signature(ev, pk)
+            if not ok:
+                sig_ok = False
+
+            prev = ev.get("event_sha256")
+
+        report["checks"]["event_chain_ok"] = chain_ok
+        report["checks"]["event_hashes_ok"] = hash_ok
+        report["checks"]["event_sigs_ok"] = sig_ok
+
+        if not chain_ok:
+            report["errors"].append("event_chain_broken")
+        if not hash_ok:
+            report["errors"].append("event_hash_mismatch")
+        if not sig_ok:
+            report["errors"].append("event_signature_invalid")
+
+        # Optional media hash
+        if media_name and media_bytes is not None:
+            media_sha = hashlib.sha256(media_bytes).hexdigest()
+            expected = root.get("media_sha256")
+            report["checks"]["media_present"] = True
+            report["checks"]["media_sha_match"] = (expected == media_sha)
+            if expected and expected != media_sha:
+                report["errors"].append("media_hash_mismatch")
+        else:
+            report["checks"]["media_present"] = False
+
+        # Replay-verify state hashes (strongest integrity check)
+        cfg_dict = root["cfg"]["cfg"]
+        cfg = SimConfig(
+            size=int(cfg_dict["size"]),
+            walls_pct=float(cfg_dict["walls_pct"]),
+            coins=int(cfg_dict["coins"]),
+            hazards=int(cfg_dict["hazards"]),
+            pov_radius=int(cfg_dict["pov_radius"]),
+            max_steps=int(cfg_dict.get("max_steps", 2000)),
+        )
+        seed = int(root["seed"])
+
+        st = reset_sim(cfg, seed)
+
+        mismatch_at = None
+        for ev in events:
+            if ev["action"] == "RESET":
+                if ev["state_sha256"] != st.last_state_sha256:
+                    mismatch_at = ("RESET", ev["step"])
+                    break
+                continue
+
+            st, _ = step_sim(cfg, st)
+            if ev["state_sha256"] != st.last_state_sha256:
+                mismatch_at = (ev["action"], ev["step"])
+                break
+
+            if st.done:
+                break
+
+        report["checks"]["replay_state_hashes_ok"] = (mismatch_at is None)
+        if mismatch_at is not None:
+            report["errors"].append(f"replay_state_hash_mismatch_at: action={mismatch_at[0]} step={mismatch_at[1]}")
+
+        last_event = events[-1]["event_sha256"] if events else None
+        report["checks"]["final_event_matches_root"] = (root.get("final_event_sha256") == last_event)
+        if root.get("final_event_sha256") != last_event:
+            report["errors"].append("root_final_event_mismatch")
+
+        report["ok"] = (len(report["errors"]) == 0)
+        return report
+
+    except Exception as e:
+        report["errors"].append(f"exception: {type(e).__name__}: {e}")
+        return report
+
+
+def replay_from_receipt_bundle_zip(zip_path: str, export_gif: bool = True, watermark: bool = True) -> Dict[str, Any]:
+    out: Dict[str, Any] = {
+        "ok": False,
+        "message": "",
+        "world_img": None,
+        "pov_img": None,
+        "gif_path": None,
+        "final_hash": None,
+        "expected_hash": None,
+    }
+
+    ver = verify_receipt_bundle_zip(zip_path)
+    bundle = _load_bundle(zip_path)
+    root = bundle["root"]
+    events = bundle["events"]
+
+    cfg_dict = root["cfg"]["cfg"]
+    cfg = SimConfig(
+        size=int(cfg_dict["size"]),
+        walls_pct=float(cfg_dict["walls_pct"]),
+        coins=int(cfg_dict["coins"]),
+        hazards=int(cfg_dict["hazards"]),
+        pov_radius=int(cfg_dict["pov_radius"]),
+        max_steps=int(cfg_dict.get("max_steps", 2000)),
+    )
+    seed = int(root["seed"])
+
+    st = reset_sim(cfg, seed)
+    frames = [render_pov_image(st)]
+
+    for ev in events:
+        if ev["action"] == "RESET":
+            continue
+        st, _ = step_sim(cfg, st)
+        frames.append(render_pov_image(st))
+        if st.done:
+            break
+
+    out["world_img"] = render_world_image(st)
+    out["pov_img"] = render_pov_image(st)
+    out["final_hash"] = st.last_state_sha256
+    out["expected_hash"] = root.get("final_state_sha256")
+
+    if export_gif:
+        tmp_dir = tempfile.mkdtemp(prefix="replayproof_replay_")
+        gif_path = os.path.join(tmp_dir, f"replay_{root['session_id']}.gif")
+        wm = None
+        if watermark:
+            wm = f"final:{(st.last_state_sha256 or '')[:12]}"
+        export_gif(frames, gif_path, fps=10, watermark=wm)
+        out["gif_path"] = gif_path
+
+    if ver["ok"] and out["final_hash"] == out["expected_hash"]:
+        out["ok"] = True
+        out["message"] = f"✅ Replay OK. final_state_sha256 matches. ({(out['final_hash'] or '')[:12]})"
+    else:
+        out["ok"] = False
+        out["message"] = (
+            f"❌ Replay mismatch or verification failed.\n"
+            f"verify_ok={ver['ok']}  replay_final={(out['final_hash'] or '')[:12]}  expected={(out['expected_hash'] or '')[:12]}\n"
+            f"errors={ver.get('errors')}"
+        )
+
+    return out