#!/usr/bin/env python3
"""aperture_session — the client-side flight recorder for an autonomous agent run.

`aperture.session/v1`

You cannot hand us your 12-hour agent to watch, and we do not host it. But every time your agent
calls Aperture, it gets back a signed receipt. This module is the primitive that chains *your own*
receipts, on *your* side, into a tamper-evident **session**: an ordered trail with a single head you
can publish (and later witness in the public transparency log). The counterpart verifier lives in
`transparency_log.py` (`verify_session`); this is the emitter that produces what it consumes.

A session head binds the trail two independent ways:

  • **merkle_root** — RFC 6962 Merkle root over the receipt leaves (`session_root`). Binds content and
    order; supports inclusion proofs against a witness log. Session-id-agnostic by construction.
  • **chain_tip** — a domain-separated rolling hash seeded by the session id and folded over each leaf
    in turn (a `prev_receipt_hash` linked list). Binds content, order, **and the session identity**,
    and can be extended incrementally as receipts arrive without recomputing a tree.

Altering, reordering, dropping, or inserting a receipt moves both. Lifting an intact trail into a
different session id moves the chain_tip (not the root) — which is exactly why the head carries both.

This module signs nothing new. Its security rests entirely on the Aperture signatures already on each
receipt; the head only fixes their **sequence** and **session**. Dependency-light: `cryptography`
(via aperture_verify) + stdlib. Covered by `selftest`.

CLI:
    python3 aperture_session.py open      <bundle.json> [--id SID]     # start an empty session bundle
    python3 aperture_session.py record    <bundle.json> <receipt.json> # append a receipt, re-seal the head
    python3 aperture_session.py head       <bundle.json>               # print the publishable session head
    python3 aperture_session.py verify     <bundle.json>               # re-verify a saved bundle end-to-end
    python3 aperture_session.py selftest                               # 5/5
"""
from __future__ import annotations

import argparse
import hashlib
import json
import sys
import uuid

import aperture_verify
import transparency_log

ALG = "aperture.session/v1"
_CHAIN = b"aperture.session.chain/v1\x00"
_WRAP_FIELDS = ("receipt", "cert")


def _unwrap(obj: dict) -> dict:
    """Accept either a bare signed receipt/cert or a permalink envelope ({"receipt": {...}})."""
    for w in _WRAP_FIELDS:
        if isinstance(obj, dict) and isinstance(obj.get(w), dict) and "signature" in obj[w]:
            return obj[w]
    return obj


def chain_tip(session_id: str, receipts: list[dict]) -> str:
    """Rolling, domain-separated hash: h0 = H(TAG || session_id); h_i = H(TAG || h_{i-1} || leaf_i).
    Binds the session id *and* the ordered leaves into one 32-byte tip."""
    h = hashlib.sha256(_CHAIN + session_id.encode("utf-8")).digest()
    for r in receipts:
        leaf = bytes.fromhex(aperture_verify.leaf_hash(r))
        h = hashlib.sha256(_CHAIN + h + leaf).digest()
    return h.hex()


def head(session_id: str, receipts: list[dict]) -> dict:
    """The publishable session head — the one small object a customer records / witnesses."""
    return {
        "alg": ALG,
        "session_id": session_id,
        "n": len(receipts),
        "merkle_root": transparency_log.session_root(receipts) if receipts else "",
        "chain_tip": chain_tip(session_id, receipts),
    }


class Session:
    """A client-side agent-run recorder. Drop it in your agent loop; feed it each Aperture receipt."""

    def __init__(self, session_id: str | None = None, *, verify: bool = True, pinned: dict | None = None):
        self.session_id = session_id or ("sess_" + uuid.uuid4().hex)
        self.verify = verify          # fail-fast: reject an unsigned/forged receipt at record() time
        self.pinned = pinned
        self.receipts: list[dict] = []

    def record(self, receipt: dict) -> dict:
        """Append one receipt (bare or enveloped). If verify=True, its Aperture signature is checked now
        so a bad receipt never silently enters the trail. Returns the running head."""
        r = _unwrap(receipt)
        if self.verify:
            aperture_verify.verify_receipt(r, pinned=self.pinned)   # raises VerifyError on any failure
        self.receipts.append(r)
        return self.head()

    def head(self) -> dict:
        return head(self.session_id, self.receipts)

    def bundle(self) -> dict:
        """The full portable artifact: the head plus the ordered receipts it commits to."""
        return {"head": self.head(), "receipts": self.receipts}

    def save(self, path: str) -> None:
        with open(path, "w", encoding="utf-8") as f:
            json.dump(self.bundle(), f, ensure_ascii=False, indent=2)

    @classmethod
    def load(cls, path: str, *, verify: bool = False, pinned: dict | None = None) -> "Session":
        with open(path, "r", encoding="utf-8") as f:
            b = json.load(f)
        s = cls(session_id=b["head"]["session_id"], verify=False, pinned=pinned)
        s.receipts = [_unwrap(r) for r in b.get("receipts", [])]
        s.verify = verify
        return s


def verify_bundle(bundle: dict, pinned: dict | None = None) -> dict:
    """Re-verify a saved session bundle end-to-end and independently:
      1. every receipt's Aperture signature (via transparency_log.verify_session),
      2. the Merkle root matches the ordered leaves,
      3. the chain_tip matches the recomputed session-bound rolling hash,
      4. the head's n matches the number of receipts.
    Never raises — returns a diagnosable report so a partial/broken trail is legible."""
    h = bundle.get("head", {}) or {}
    receipts = [_unwrap(r) for r in bundle.get("receipts", [])]
    sid = h.get("session_id", "")
    sess = transparency_log.verify_session(receipts, h.get("merkle_root", ""), pinned=pinned)
    tip = chain_tip(sid, receipts)
    tip_ok = (tip == (h.get("chain_tip") or "").strip().lower())
    n_ok = (h.get("n") == len(receipts))
    return {
        "ok": bool(sess["ok"] and tip_ok and n_ok),
        "session_id": sid,
        "n": len(receipts),
        "n_claimed": h.get("n"),
        "n_matches": n_ok,
        "sigs_ok": sess["n_bad_sig"] == 0,
        "bad_sig_indices": sess["bad_sig_indices"],
        "merkle_root": sess["root"],
        "root_matches": sess["root_matches"],
        "chain_tip": tip,
        "chain_tip_matches": tip_ok,
    }


# ── CLI ─────────────────────────────────────────────────────────────────────────────────────────────
def _load(path: str) -> dict:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def main(argv: list[str]) -> int:
    p = argparse.ArgumentParser(prog="aperture_session", description="client-side agent-run flight recorder")
    sub = p.add_subparsers(dest="cmd")
    o = sub.add_parser("open"); o.add_argument("bundle"); o.add_argument("--id", default=None)
    r = sub.add_parser("record"); r.add_argument("bundle"); r.add_argument("receipt")
    r.add_argument("--no-verify", action="store_true")
    hd = sub.add_parser("head"); hd.add_argument("bundle")
    v = sub.add_parser("verify"); v.add_argument("bundle")
    sub.add_parser("selftest")
    a = p.parse_args(argv)

    if a.cmd == "open":
        s = Session(session_id=a.id)
        s.save(a.bundle)
        print(json.dumps(s.head(), indent=2))
        print("session opened: %s" % s.session_id, file=sys.stderr)
        return 0
    if a.cmd == "record":
        s = Session.load(a.bundle, verify=not a.no_verify)
        try:
            s.record(_load(a.receipt))
        except aperture_verify.VerifyError as e:
            print("✗ refused — receipt failed verification: %s" % e, file=sys.stderr)
            return 1
        s.save(a.bundle)
        h = s.head()
        print("✓ recorded — session %s now holds %d receipt(s)" % (h["session_id"], h["n"]), file=sys.stderr)
        print("  merkle_root %s…  chain_tip %s…" % (h["merkle_root"][:16], h["chain_tip"][:16]), file=sys.stderr)
        return 0
    if a.cmd == "head":
        print(json.dumps(_load(a.bundle).get("head", {}), indent=2))
        return 0
    if a.cmd == "verify":
        rep = verify_bundle(_load(a.bundle))
        if rep["ok"]:
            print("✓ SESSION VERIFIED — %d receipts, every signature valid, root + chain-tip match "
                  "(session %s)" % (rep["n"], rep["session_id"]))
            return 0
        why = []
        if not rep["sigs_ok"]:
            why.append("%d receipt(s) failed signature %s" % (len(rep["bad_sig_indices"]), rep["bad_sig_indices"]))
        if not rep["root_matches"]:
            why.append("merkle root mismatch")
        if not rep["chain_tip_matches"]:
            why.append("chain-tip mismatch (session id or order altered)")
        if not rep["n_matches"]:
            why.append("receipt count mismatch (%s claimed vs %d present)" % (rep["n_claimed"], rep["n"]))
        print("✗ SESSION NOT VERIFIED — " + "; ".join(why))
        return 1
    if a.cmd == "selftest":
        return _selftest()
    p.print_help()
    return 2


def _selftest() -> int:
    """Self-sign a throwaway trail (no network, no pinned prod key) and exercise the guarantees."""
    from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
    import base64
    sk = Ed25519PrivateKey.generate()
    pub = base64.b64encode(sk.public_key().public_bytes_raw()).decode()
    pin = {pub: "selftest"}

    def mk(q: str) -> dict:
        r = {"kind": "read", "query": q, "certificate": "GROUNDED"}
        canon = json.dumps(r, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode()
        r["signature"] = {"alg": "ed25519", "pub": pub, "sig": base64.b64encode(sk.sign(canon)).decode()}
        return r

    ok = total = 0

    # 1. a recorded session round-trips and verifies
    s = Session(verify=True, pinned=pin)
    for q in ("q1", "q2", "q3"):
        s.record(mk(q))
    b = s.bundle()
    hit = verify_bundle(b, pinned=pin)["ok"]
    total += 1; ok += hit
    print("  %s a 3-receipt recorded session verifies (sigs + root + chain-tip)" % ("✓" if hit else "✗"))

    # 2. record() fails fast on an unsigned/forged receipt
    forged = mk("q4"); forged["query"] = "tampered-after-signing"
    try:
        Session(verify=True, pinned=pin).record(forged); refused = False
    except aperture_verify.VerifyError:
        refused = True
    total += 1; ok += refused
    print("  %s record() refuses a receipt whose signature does not match its body" % ("✓" if refused else "✗"))

    # 3. tampering a receipt in a saved bundle is caught
    tb = json.loads(json.dumps(b)); tb["receipts"][1]["certificate"] = "OFF_MAP"
    hit = not verify_bundle(tb, pinned=pin)["ok"]
    total += 1; ok += hit
    print("  %s tampering a receipt in the bundle breaks verification" % ("✓" if hit else "✗"))

    # 4. reordering the trail moves the head (root + chain-tip)
    rb = json.loads(json.dumps(b)); rb["receipts"] = [rb["receipts"][0], rb["receipts"][2], rb["receipts"][1]]
    rep = verify_bundle(rb, pinned=pin)
    hit = (not rep["ok"]) and (not rep["root_matches"]) and (not rep["chain_tip_matches"])
    total += 1; ok += hit
    print("  %s reordering the trail breaks both the root and the chain-tip" % ("✓" if hit else "✗"))

    # 5. lifting an intact trail into a different session id moves the chain-tip (not the root)
    lb = json.loads(json.dumps(b)); lb["head"]["session_id"] = "sess_impostor"
    rep = verify_bundle(lb, pinned=pin)
    hit = (not rep["ok"]) and rep["root_matches"] and (not rep["chain_tip_matches"])
    total += 1; ok += hit
    print("  %s re-labeling the session id is caught by the chain-tip (root alone would not)" % ("✓" if hit else "✗"))

    print("  session self-test: %d/%d passed" % (ok, total))
    return 0 if ok == total else 1


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))