agentmetry 0.4.0__py3-none-any.whl

agentmetry/core/audit/trail_merkle.py

@@ -0,0 +1,332 @@
+"""Merkle tree over the hash-chained trail, for proofs that disclose one event.
+
+The linear chain in trail_chain.py answers "has this file been altered". It
+cannot answer "did this specific event happen" without handing over the entire
+file, because every record's hash depends on the one before it. For an audit
+that is the wrong shape: proving one tool call should not require disclosing a
+month of unrelated work to whoever is asking.
+
+A Merkle tree over the same `record_sha256` values answers it in O(log n)
+hashes. Publish the root somewhere the audited machine cannot rewrite, and any
+single event can later be proved against it while the rest of the trail stays
+private.
+
+**Additive by construction.** Nothing here changes the envelope, the canonical
+event JSON, or `record_sha256`. Leaves are the record hashes the chain already
+computes, so every trail ever written is compatible and no stored event needs
+migrating. The root lives in the sidecar and is derived, never authoritative:
+delete it and it recomputes from the file.
+
+Deliberately not maintained on the append path. The root is what you publish
+periodically, not per event, and `append_chained_line` is already the hottest
+and least protected code in the system (it has no cross-process lock). Adding
+an incremental tree update there would buy latency we do not need at a risk we
+should not take.
+
+## Why this is not the reference implementation it was modelled on
+
+Microsoft's Agent Governance Toolkit ships a `MerkleAuditChain` doing the same
+job, MIT licensed and readable. Two properties of that implementation are the
+reason this one was written from RFC 6962 instead:
+
+1. **No domain separation.** Leaves are hashed as the entry hash itself and
+   internal nodes as `sha256(left || right)`, so the two are indistinguishable.
+   That is the textbook Merkle second-preimage weakness: an internal node can be
+   presented as a leaf, and a proof for a node that is not an event verifies
+   against the root. RFC 6962 §2.1 prefixes leaves with 0x00 and internal nodes
+   with 0x01 for exactly this reason, which is what this module does.
+
+2. **Duplicated last node when a level is odd** (`right = left` where no sibling
+   exists). That is the shape behind CVE-2012-2459: two different leaf sequences
+   can produce the same root, so a proof can be valid for a tree that was never
+   written. RFC 6962 splits at the largest power of two below `n` instead and
+   never duplicates or pads, so distinct inputs always give distinct roots.
+
+Neither is exotic and both are well documented; the point is only that a
+tamper-evidence primitive is the wrong place to inherit a subtle weakness for
+the sake of matching someone else's bytes. Interoperating on the *format* is
+still worth doing, and a CT-style tree is the more interoperable choice anyway
+since it is what Certificate Transparency, Trillian and Sigstore already speak.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from agentmetry.core.audit.trail_chain import chain_sidecar_path, is_chained_record
+
+# RFC 6962 §2.1 domain separation. Without these two bytes a leaf and an
+# internal node hash identically and the tree loses second-preimage resistance.
+_LEAF_PREFIX = b"\x00"
+_NODE_PREFIX = b"\x01"
+
+MERKLE_VERSION = 1
+
+
+def _leaf_hash(record_sha256: str) -> bytes:
+    return hashlib.sha256(_LEAF_PREFIX + bytes.fromhex(record_sha256)).digest()
+
+
+def _node_hash(left: bytes, right: bytes) -> bytes:
+    return hashlib.sha256(_NODE_PREFIX + left + right).digest()
+
+
+def _split(n: int) -> int:
+    """Largest power of two strictly less than n (RFC 6962's k).
+
+    Splitting here rather than padding to a power of two is what keeps distinct
+    leaf sequences mapping to distinct roots.
+    """
+    k = 1
+    while k << 1 < n:
+        k <<= 1
+    return k
+
+
+def _root(leaves: list[bytes]) -> bytes:
+    if not leaves:
+        # MTH({}) = SHA-256() of the empty string, per RFC 6962.
+        return hashlib.sha256(b"").digest()
+    if len(leaves) == 1:
+        return leaves[0]
+    k = _split(len(leaves))
+    return _node_hash(_root(leaves[:k]), _root(leaves[k:]))
+
+
+def _audit_path(index: int, leaves: list[bytes]) -> list[bytes]:
+    if len(leaves) <= 1:
+        return []
+    k = _split(len(leaves))
+    if index < k:
+        return _audit_path(index, leaves[:k]) + [_root(leaves[k:])]
+    return _audit_path(index - k, leaves[k:]) + [_root(leaves[:k])]
+
+
+def _root_from_path(index: int, tree_size: int, leaf: bytes, path: list[bytes]) -> bytes:
+    """Recompute the root from a leaf and its audit path (RFC 6962 §2.1.1).
+
+    The index and tree size carry the shape, which is why a proof cannot be
+    replayed against a differently sized tree.
+    """
+    if index >= tree_size:
+        raise ValueError(f"leaf index {index} outside tree of size {tree_size}")
+    fn, sn = index, tree_size - 1
+    node = leaf
+    for sibling in path:
+        if fn & 1 or fn == sn:
+            node = _node_hash(sibling, node)
+            while fn & 1 == 0 and fn != 0:
+                fn >>= 1
+                sn >>= 1
+        else:
+            node = _node_hash(node, sibling)
+        fn >>= 1
+        sn >>= 1
+    if sn != 0:
+        raise ValueError("audit path too short for the stated tree size")
+    return node
+
+
+# ----------------------------------------------------------------------
+# Reading the trail
+# ----------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class TrailLeaves:
+    """Record hashes in file order, with the seq each came from.
+
+    Legacy unchained lines carry no `record_sha256` and are skipped rather than
+    invented for. A trail that predates chaining simply has fewer leaves than
+    lines, and `verify_trail_file` already reports that count separately.
+    """
+
+    record_hashes: list[str]
+    seqs: list[int]
+    skipped_legacy: int
+
+    def __len__(self) -> int:
+        return len(self.record_hashes)
+
+    def index_of_seq(self, seq: int) -> int | None:
+        try:
+            return self.seqs.index(seq)
+        except ValueError:
+            return None
+
+
+def read_leaves(trail_path: Path) -> TrailLeaves:
+    hashes: list[str] = []
+    seqs: list[int] = []
+    skipped = 0
+    if not trail_path.is_file():
+        return TrailLeaves([], [], 0)
+    with trail_path.open("r", encoding="utf-8", errors="replace") as fh:
+        for line in fh:
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                record = json.loads(line)
+            except json.JSONDecodeError:
+                skipped += 1
+                continue
+            if not is_chained_record(record):
+                skipped += 1
+                continue
+            trail = record["trail"]
+            digest = str(trail.get("record_sha256") or "")
+            if len(digest) != 64:
+                skipped += 1
+                continue
+            hashes.append(digest)
+            seqs.append(int(trail.get("seq") or 0))
+    return TrailLeaves(hashes, seqs, skipped)
+
+
+def merkle_root(trail_path: Path, tree_size: int | None = None) -> tuple[str, int]:
+    """(root hex, tree size) for a trail file, optionally at a historical size.
+
+    `tree_size` exists because the trail is append-only and live: the root
+    changes on every event. A proof issued this morning names the tree it was
+    issued against, and checking it against tonight's root fails for the most
+    boring possible reason. Recomputing over the first `tree_size` leaves gives
+    the root as it stood then, which is the thing the proof actually claims.
+
+    Without this the whole feature is unusable on a running recorder, which is
+    the only kind anybody has.
+    """
+    leaves = read_leaves(trail_path)
+    hashes = leaves.record_hashes
+    if tree_size is not None:
+        if tree_size > len(hashes):
+            raise ValueError(
+                f"trail holds {len(hashes)} chained records, cannot reconstruct "
+                f"a tree of size {tree_size}. The proof is either forged or from "
+                "a different trail."
+            )
+        hashes = hashes[:tree_size]
+    nodes = [_leaf_hash(h) for h in hashes]
+    return _root(nodes).hex(), len(nodes)
+
+
+# ----------------------------------------------------------------------
+# Proofs
+# ----------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class InclusionProof:
+    seq: int
+    leaf_index: int
+    tree_size: int
+    record_sha256: str
+    path: list[str]
+    root_sha256: str
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "v": MERKLE_VERSION,
+            "alg": "rfc6962-sha256",
+            "seq": self.seq,
+            "leaf_index": self.leaf_index,
+            "tree_size": self.tree_size,
+            "record_sha256": self.record_sha256,
+            "path": list(self.path),
+            "root_sha256": self.root_sha256,
+        }
+
+    @staticmethod
+    def from_dict(data: dict[str, Any]) -> InclusionProof:
+        return InclusionProof(
+            seq=int(data.get("seq") or 0),
+            leaf_index=int(data["leaf_index"]),
+            tree_size=int(data["tree_size"]),
+            record_sha256=str(data["record_sha256"]),
+            path=[str(p) for p in (data.get("path") or [])],
+            root_sha256=str(data.get("root_sha256") or ""),
+        )
+
+
+def build_proof(trail_path: Path, seq: int) -> InclusionProof:
+    leaves = read_leaves(trail_path)
+    index = leaves.index_of_seq(seq)
+    if index is None:
+        raise ValueError(f"seq {seq} is not a chained record in {trail_path}")
+    nodes = [_leaf_hash(h) for h in leaves.record_hashes]
+    path = _audit_path(index, nodes)
+    return InclusionProof(
+        seq=seq,
+        leaf_index=index,
+        tree_size=len(nodes),
+        record_sha256=leaves.record_hashes[index],
+        path=[p.hex() for p in path],
+        root_sha256=_root(nodes).hex(),
+    )
+
+
+def verify_proof(proof: InclusionProof, expected_root: str | None = None) -> tuple[bool, str]:
+    """Check a proof, optionally against an independently held root.
+
+    `expected_root` is the whole point. Verifying against the root carried
+    inside the proof only shows the proof is internally consistent, which a
+    forger can also arrange. The check that means something compares it to a
+    root recorded somewhere the audited machine could not rewrite.
+    """
+    try:
+        leaf = _leaf_hash(proof.record_sha256)
+        path = [bytes.fromhex(p) for p in proof.path]
+        computed = _root_from_path(proof.leaf_index, proof.tree_size, leaf, path).hex()
+    except (ValueError, TypeError) as exc:
+        return False, f"malformed proof: {exc}"
+
+    if proof.root_sha256 and computed != proof.root_sha256:
+        return False, (
+            f"proof does not reconstruct its own root (computed {computed[:16]}…, "
+            f"stated {proof.root_sha256[:16]}…)"
+        )
+    if expected_root:
+        if computed != expected_root:
+            return False, (
+                f"proof is valid but for a different tree (computed {computed[:16]}…, "
+                f"expected {expected_root[:16]}…)"
+            )
+        return True, f"verified against the supplied root {computed[:16]}…"
+    return True, (
+        f"internally consistent (root {computed[:16]}…). No external root was "
+        "supplied, so this shows the proof is well formed, not that the trail "
+        "was not rewritten."
+    )
+
+
+# ----------------------------------------------------------------------
+# Sidecar
+# ----------------------------------------------------------------------
+
+
+def record_root(trail_path: Path) -> dict[str, Any]:
+    """Recompute the root and store it alongside the chain head.
+
+    Written as extra keys in the existing sidecar. `_load_sidecar` reads only
+    `seq` and `last_sha256` and ignores everything else, so an older build reads
+    a newer sidecar without noticing, which is the property that makes this
+    additive rather than a format change.
+    """
+    root, size = merkle_root(trail_path)
+    sidecar = chain_sidecar_path(trail_path)
+    data: dict[str, Any] = {}
+    if sidecar.is_file():
+        try:
+            data = json.loads(sidecar.read_text(encoding="utf-8"))
+        except (OSError, json.JSONDecodeError):
+            data = {}
+    data["merkle_root"] = root
+    data["merkle_tree_size"] = size
+    data["merkle_alg"] = "rfc6962-sha256"
+    sidecar.parent.mkdir(parents=True, exist_ok=True)
+    sidecar.write_text(json.dumps(data, indent=2) + "\n", encoding="utf-8")
+    return {"root": root, "tree_size": size, "sidecar": str(sidecar)}

agentmetry/core/auth.py

@@ -0,0 +1,54 @@
+"""Optional API key authentication for mutating endpoints."""
+
+from __future__ import annotations
+
+import secrets
+
+from fastapi import HTTPException, Request, Security
+from fastapi.security import APIKeyHeader
+
+from agentmetry.core.config import settings
+
+
+def _token_matches(token: str | None) -> bool:
+    """Constant-time compare so the key can't be recovered a byte at a time.
+
+    A plain `==` short-circuits on the first differing byte, leaking the shared
+    key's prefix through response timing to anyone who can reach the endpoint.
+    `compare_digest` runs in time independent of where the mismatch is.
+    """
+    return secrets.compare_digest(token or "", settings.api_key)
+
+_api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
+
+
+async def require_api_key(
+    request: Request,
+    api_key: str | None = Security(_api_key_header),
+) -> None:
+    """Skip auth when AGENTMETRY_API_KEY is unset (local dev)."""
+    if not settings.api_key:
+        return
+
+    token = api_key
+    if not token:
+        auth = request.headers.get("Authorization", "")
+        if auth.lower().startswith("bearer "):
+            token = auth[7:].strip()
+
+    if not _token_matches(token):
+        raise HTTPException(status_code=401, detail="Invalid or missing API key")
+
+
+def verify_ws_token(query_token: str | None, request: Request) -> bool:
+    """Validate WebSocket connection when AGENTMETRY_API_KEY is set."""
+    if not settings.api_key:
+        return True
+    token = query_token
+    if not token:
+        token = request.headers.get("X-API-Key")
+    if not token:
+        auth = request.headers.get("Authorization", "")
+        if auth.lower().startswith("bearer "):
+            token = auth[7:].strip()
+    return _token_matches(token)

agentmetry/core/bus/__init__.py

@@ -0,0 +1,5 @@
+"""In-process event bus — decoupled pub/sub with a durable SQLite outbox.
+
+Import from the submodules (core.bus.bus, core.bus.events, core.bus.outbox);
+bridges live in core.bus.bridges and api.ws_bridge and are wired in lifespan.
+"""

agentmetry/core/bus/audit_exporter.py

@@ -0,0 +1,107 @@
+"""Subscribe to the event bus and forward canonical audit events to configured sinks."""
+
+from __future__ import annotations
+
+import logging
+
+from agentmetry.core.audit.alerts import AlertWebhookSink
+from agentmetry.core.audit.canonical import normalize_outbox_row
+from agentmetry.core.audit.sinks import MultiAuditSink, build_audit_sinks, parse_sink_modes
+from agentmetry.core.bus.bus import EventBus, bus
+from agentmetry.core.bus.events import Event, LLM_TOKEN
+from agentmetry.core.config import settings
+
+logger = logging.getLogger(__name__)
+
+
+def event_to_outbox_row(event: Event) -> dict:
+    return {
+        "seq": event.seq,
+        "ts": event.ts,
+        "topic": event.topic,
+        "session_id": event.session_id,
+        "thread_id": event.thread_id,
+        "payload": event.payload,
+    }
+
+
+def _describe_sink_targets(modes: set[str]) -> list[str]:
+    targets: list[str] = []
+    if "file" in modes:
+        targets.append(str(settings.audit_export_path))
+    if "webhook" in modes and settings.audit_webhook_url.strip():
+        targets.append(settings.audit_webhook_url.strip())
+    if "elastic" in modes and settings.audit_elastic_url.strip():
+        targets.append(f"elastic:{settings.audit_elastic_index}")
+    if "splunk" in modes and settings.audit_splunk_hec_url.strip():
+        targets.append(f"splunk:{settings.audit_splunk_index}")
+    return targets
+
+
+async def audit_exporter(
+    bus_: EventBus = bus,
+    *,
+    enabled: bool | None = None,
+) -> None:
+    """Forward normalized bus events to configured sinks (best-effort)."""
+    if enabled is None:
+        enabled = settings.audit_export_enabled
+    if not enabled:
+        logger.info("Audit export disabled (AGENTMETRY_AUDIT_EXPORT_ENABLED=0)")
+        return
+
+    modes = parse_sink_modes(settings.audit_sink)
+
+    sink = build_audit_sinks(
+        modes=modes,
+        file_path=settings.audit_export_path,
+        webhook_url=settings.audit_webhook_url,
+        webhook_timeout_seconds=settings.audit_webhook_timeout_seconds,
+        webhook_format=settings.audit_webhook_format,
+        elastic_url=settings.audit_elastic_url,
+        elastic_index=settings.audit_elastic_index,
+        elastic_api_key=settings.audit_elastic_api_key,
+        elastic_verify_tls=settings.audit_elastic_verify_tls,
+        splunk_hec_url=settings.audit_splunk_hec_url,
+        splunk_hec_token=settings.audit_splunk_hec_token,
+        splunk_index=settings.audit_splunk_index,
+        splunk_sourcetype=settings.audit_splunk_sourcetype,
+        splunk_verify_tls=settings.audit_splunk_verify_tls,
+    )
+    if sink is None and not settings.audit_alert_webhook_url.strip():
+        logger.warning(
+            "Audit export enabled but no sinks configured "
+            "(check AGENTMETRY_AUDIT_SINK and backend credentials)"
+        )
+        return
+
+    # Add Alert Sink if configured
+    if settings.audit_alert_webhook_url.strip():
+        alert_sink = AlertWebhookSink(
+            settings.audit_alert_webhook_url.strip(),
+            timeout_seconds=settings.audit_webhook_timeout_seconds,
+        )
+        if sink is None:
+            sink = alert_sink
+        elif isinstance(sink, MultiAuditSink):
+            sink._sinks.append(alert_sink)
+        else:
+            sink = MultiAuditSink([sink, alert_sink])
+
+    targets = _describe_sink_targets(modes)
+    if settings.audit_alert_webhook_url.strip():
+        targets.append(f"alerts:{settings.audit_alert_webhook_url.strip()}")
+    logger.info("Audit export sinks → %s", ", ".join(targets))
+
+    sub = bus_.subscribe("audit-export", exclude={LLM_TOKEN})
+    try:
+        while True:
+            event = await sub.get()
+            try:
+                canonical = normalize_outbox_row(event_to_outbox_row(event))
+                if canonical is not None:
+                    await sink.emit(canonical)
+            except Exception:
+                logger.exception("Audit export failed for %s seq=%s", event.topic, event.seq)
+    finally:
+        bus_.unsubscribe(sub)

agentmetry/core/bus/bridges.py

@@ -0,0 +1,26 @@
+"""Core-side bus subscribers: durable persistence."""
+
+from __future__ import annotations
+
+import logging
+
+from agentmetry.core.bus.bus import EventBus, bus
+from agentmetry.core.bus.events import LLM_TOKEN
+from agentmetry.core.bus.outbox import EventOutbox, get_outbox
+
+logger = logging.getLogger(__name__)
+
+
+async def outbox_persister(bus_: EventBus = bus, outbox: EventOutbox | None = None) -> None:
+    """Persist every event except the token firehose. Never drops."""
+    ob = outbox or get_outbox()
+    sub = bus_.subscribe("outbox", exclude={LLM_TOKEN})
+    try:
+        while True:
+            event = await sub.get()
+            try:
+                ob.append(event)
+            except Exception:
+                logger.exception("Outbox write failed for %s", event.topic)
+    finally:
+        bus_.unsubscribe(sub)

agentmetry/core/bus/bus.py

@@ -0,0 +1,102 @@
+"""EventBus — synchronous publish, per-subscriber queues, explicit drop policy.
+
+Loop-thread only: publish() must be called from the event loop (the vault
+watcher already marshals via call_soon_threadsafe). Publishing never blocks —
+each subscriber declares whether overflow drops its oldest events (WS bridge)
+or is unbounded (outbox persister, which must never lose audit events).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from dataclasses import dataclass, field
+from typing import Any
+
+from agentmetry.core.bus.events import Event
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Subscription:
+    name: str
+    queue: asyncio.Queue
+    topics: frozenset[str] | None      # None = all topics
+    exclude: frozenset[str]
+    drop_oldest: bool
+
+    async def get(self) -> Event:
+        return await self.queue.get()
+
+
+@dataclass
+class EventBus:
+    _subs: list[Subscription] = field(default_factory=list)
+    _seq: int = 0
+
+    def set_initial_seq(self, seq: int) -> None:
+        """Continue numbering after the outbox's last persisted event."""
+        self._seq = max(self._seq, seq)
+
+    def subscribe(
+        self,
+        name: str,
+        *,
+        topics: set[str] | None = None,
+        exclude: set[str] | None = None,
+        maxsize: int = 0,
+        drop_oldest: bool = False,
+    ) -> Subscription:
+        sub = Subscription(
+            name=name,
+            queue=asyncio.Queue(maxsize=maxsize),
+            topics=frozenset(topics) if topics is not None else None,
+            exclude=frozenset(exclude or ()),
+            drop_oldest=drop_oldest,
+        )
+        self._subs.append(sub)
+        return sub
+
+    def unsubscribe(self, sub: Subscription) -> None:
+        self._subs = [s for s in self._subs if s is not sub]
+
+    def publish(
+        self,
+        topic: str,
+        payload: dict[str, Any],
+        *,
+        session_id: str = "",
+        thread_id: str = "",
+    ) -> Event:
+        self._seq += 1
+        event = Event(
+            topic=topic,
+            payload=payload,
+            seq=self._seq,
+            session_id=session_id,
+            thread_id=thread_id,
+        )
+        for sub in self._subs:
+            if sub.topics is not None and topic not in sub.topics:
+                continue
+            if topic in sub.exclude:
+                continue
+            try:
+                sub.queue.put_nowait(event)
+            except asyncio.QueueFull:
+                if sub.drop_oldest:
+                    try:
+                        sub.queue.get_nowait()
+                    except asyncio.QueueEmpty:
+                        pass
+                    try:
+                        sub.queue.put_nowait(event)
+                    except asyncio.QueueFull:
+                        pass
+                else:
+                    logger.warning("Bus subscriber %s full — dropped %s", sub.name, topic)
+        return event
+
+
+bus = EventBus()

agentmetry/core/bus/events.py

@@ -0,0 +1,50 @@
+"""Event envelope and topic constants."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any
+
+# Run lifecycle
+RUN_STARTED = "run/started"
+RUN_NODE = "run/node"
+RUN_WAITING = "run/approval_required"
+RUN_APPROVAL_GRANTED = "run/approval_granted"
+RUN_APPROVAL_DENIED = "run/approval_denied"
+RUN_COMPLETED = "run/completed"
+RUN_FAILED = "run/failed"
+RUN_TERMINATED = "run/terminated"
+FLYWHEEL_CAPTURE = "run/flywheel_capture"
+ALERT_COST = "run/cost_alert"
+ALERT_DRIFT = "run/drift_alert"
+
+# High-volume, ephemeral — excluded from the outbox.
+LLM_TOKEN = "llm/token"
+
+# Vault
+VAULT_FILE_CHANGED = "vault/file_changed"
+
+# Interrupt vector table
+INTERRUPT_RAISED = "interrupt/raised"
+INTERRUPT_RESOLVED = "interrupt/resolved"
+
+# Driver subsystem (MCP)
+DRIVER_MOUNTED = "driver/mounted"
+DRIVER_FAILED = "driver/failed"
+TOOL_CALLED = "run/tool_called"
+TOOL_DENIED = "run/tool_denied"
+
+
+def _now_iso() -> str:
+    return datetime.now(timezone.utc).isoformat()
+
+
+@dataclass(frozen=True)
+class Event:
+    topic: str
+    payload: dict[str, Any]
+    seq: int = 0
+    session_id: str = ""   # WebSocket routing hint; empty = not for the UI
+    thread_id: str = ""
+    ts: str = field(default_factory=_now_iso)