agenthacker 0.1.0__py3-none-any.whl

firewall_sdk/event_store.py

@@ -0,0 +1,660 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2026 AgentHacker
+
+"""EventStore Protocol, NullStore, and PostgresStore implementations.
+
+Defines the contract for persisting firewall events and agent invocations.
+The SDK ships NullStore (no-op default) and PostgresStore (requires psycopg2).
+Consumers can implement their own store (e.g. DynamoDB) without touching SDK code.
+"""
+
+from __future__ import annotations
+
+import logging
+import uuid
+from concurrent.futures import ThreadPoolExecutor
+from datetime import datetime
+from typing import Any, Protocol, runtime_checkable
+
+logger = logging.getLogger(__name__)
+
+
+@runtime_checkable
+class EventStore(Protocol):
+    """Contract for persisting firewall telemetry.
+
+    Two pairs of methods:
+    - record_*  — synchronous, for scripts/tests/batch operations.
+    - submit_*  — non-blocking fire-and-forget, for the logger in async routes.
+    """
+
+    def record_firewall_event(
+        self,
+        *,
+        timestamp: datetime,
+        checkpoint: str,
+        rule_id: str,
+        rule_name: str,
+        excerpt: str | None = None,
+        user_hash: str | None = None,
+        ip: str | None = None,
+        agent: str | None = None,
+        source: str = "runtime",
+        actor_role: str | None = None,
+        invocation_id: str | None = None,
+    ) -> None: ...
+
+    def record_invocation(
+        self,
+        *,
+        invocation_id: str,
+        timestamp: datetime,
+        user_hash: str | None = None,
+        question_preview: str | None = None,
+        blocked: bool,
+        checkpoint: str | None = None,
+        rule_id: str | None = None,
+        tool_calls: int = 0,
+        tokens: int = 0,
+        latency_ms: float = 0.0,
+        trace: Any = None,
+        agent: str | None = None,
+        source: str = "runtime",
+        actor_role: str | None = None,
+    ) -> None: ...
+
+    def submit_firewall_event(
+        self,
+        *,
+        timestamp: datetime,
+        checkpoint: str,
+        rule_id: str,
+        rule_name: str,
+        excerpt: str | None = None,
+        user_hash: str | None = None,
+        ip: str | None = None,
+        agent: str | None = None,
+        source: str = "runtime",
+        actor_role: str | None = None,
+        invocation_id: str | None = None,
+    ) -> None: ...
+
+    def submit_invocation(
+        self,
+        *,
+        invocation_id: str,
+        timestamp: datetime,
+        user_hash: str | None = None,
+        question_preview: str | None = None,
+        blocked: bool,
+        checkpoint: str | None = None,
+        rule_id: str | None = None,
+        tool_calls: int = 0,
+        tokens: int = 0,
+        latency_ms: float = 0.0,
+        trace: Any = None,
+        agent: str | None = None,
+        source: str = "runtime",
+        actor_role: str | None = None,
+    ) -> None: ...
+
+
+class NullStore:
+    """Default store — all methods are no-ops. Zero overhead, no persistence."""
+
+    def record_firewall_event(self, **kwargs: Any) -> None:
+        pass
+
+    def record_invocation(self, **kwargs: Any) -> None:
+        pass
+
+    def submit_firewall_event(self, **kwargs: Any) -> None:
+        pass
+
+    def submit_invocation(self, **kwargs: Any) -> None:
+        pass
+
+    def insert_batch_firewall_events(self, rows: list[dict]) -> None:
+        pass
+
+    def insert_batch_invocations(self, rows: list[dict]) -> None:
+        pass
+
+    def upsert_risk_score(self, risk_score: Any) -> None:
+        pass
+
+    def get_risk_score(self, user_hash: str) -> None:
+        return None
+
+    def pardon_user(self, user_hash: str, hours: int = 1) -> None:
+        pass
+
+
+# ── SQL ──────────────────────────────────────────────────────────────
+
+_CREATE_FIREWALL_EVENTS = """
+CREATE TABLE IF NOT EXISTS firewall_events (
+    id              UUID PRIMARY KEY,
+    invocation_id   UUID,
+    timestamp       TIMESTAMPTZ NOT NULL,
+    checkpoint      TEXT NOT NULL,
+    rule_id         TEXT NOT NULL,
+    rule_name       TEXT NOT NULL,
+    excerpt         TEXT,
+    user_hash       TEXT,
+    ip              TEXT,
+    agent           TEXT,
+    source          TEXT NOT NULL DEFAULT 'runtime',
+    actor_role      TEXT
+)
+"""
+
+_CREATE_INVOCATIONS = """
+CREATE TABLE IF NOT EXISTS invocations (
+    id              UUID PRIMARY KEY,
+    timestamp       TIMESTAMPTZ NOT NULL,
+    user_hash       TEXT,
+    question_preview TEXT,
+    blocked         BOOLEAN NOT NULL,
+    checkpoint      TEXT,
+    rule_id         TEXT,
+    tool_calls      INTEGER NOT NULL,
+    tokens          INTEGER NOT NULL,
+    latency_ms      DOUBLE PRECISION NOT NULL,
+    trace           JSONB,
+    agent           TEXT,
+    source          TEXT NOT NULL DEFAULT 'runtime',
+    actor_role      TEXT
+)
+"""
+
+_CREATE_USER_RISK_SCORES = """
+CREATE TABLE IF NOT EXISTS user_risk_scores (
+    user_hash    TEXT PRIMARY KEY,
+    score        DOUBLE PRECISION NOT NULL,
+    level        TEXT NOT NULL,
+    factors      JSONB NOT NULL DEFAULT '[]',
+    computed_at  TIMESTAMPTZ NOT NULL,
+    pardoned_until TIMESTAMPTZ
+)
+"""
+
+_ALTER_ADD_PARDONED_UNTIL = """
+ALTER TABLE user_risk_scores ADD COLUMN IF NOT EXISTS pardoned_until TIMESTAMPTZ
+"""
+
+_RISK_SCORE_INDEXES = [
+    "CREATE INDEX IF NOT EXISTS idx_urs_score ON user_risk_scores (score DESC)",
+    "CREATE INDEX IF NOT EXISTS idx_urs_level ON user_risk_scores (level)",
+    "CREATE INDEX IF NOT EXISTS idx_urs_computed_at ON user_risk_scores (computed_at)",
+]
+
+_UPSERT_RISK_SCORE = """
+INSERT INTO user_risk_scores (user_hash, score, level, factors, computed_at)
+VALUES (%s, %s, %s, %s, %s)
+ON CONFLICT (user_hash) DO UPDATE SET
+    score       = EXCLUDED.score,
+    level       = EXCLUDED.level,
+    factors     = EXCLUDED.factors,
+    computed_at = EXCLUDED.computed_at
+"""
+
+_SELECT_RISK_SCORE = """
+SELECT user_hash, score, level, factors, computed_at
+FROM user_risk_scores
+WHERE user_hash = %s
+"""
+
+_INDEXES = [
+    "CREATE INDEX IF NOT EXISTS idx_fe_timestamp ON firewall_events (timestamp)",
+    "CREATE INDEX IF NOT EXISTS idx_fe_invocation_id ON firewall_events (invocation_id)",
+    "CREATE INDEX IF NOT EXISTS idx_fe_checkpoint ON firewall_events (checkpoint)",
+    "CREATE INDEX IF NOT EXISTS idx_fe_rule_id ON firewall_events (rule_id)",
+    "CREATE INDEX IF NOT EXISTS idx_fe_user_hash ON firewall_events (user_hash)",
+    "CREATE INDEX IF NOT EXISTS idx_fe_agent ON firewall_events (agent)",
+    "CREATE INDEX IF NOT EXISTS idx_fe_source ON firewall_events (source)",
+    "CREATE INDEX IF NOT EXISTS idx_inv_timestamp ON invocations (timestamp)",
+    "CREATE INDEX IF NOT EXISTS idx_inv_user_hash ON invocations (user_hash)",
+    "CREATE INDEX IF NOT EXISTS idx_inv_blocked ON invocations (blocked)",
+    "CREATE INDEX IF NOT EXISTS idx_inv_agent ON invocations (agent)",
+    "CREATE INDEX IF NOT EXISTS idx_inv_source ON invocations (source)",
+]
+
+_INSERT_FIREWALL_EVENT = """
+INSERT INTO firewall_events
+    (id, invocation_id, timestamp, checkpoint, rule_id, rule_name,
+     excerpt, user_hash, ip, agent, source, actor_role)
+VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
+"""
+
+_INSERT_INVOCATION = """
+INSERT INTO invocations
+    (id, timestamp, user_hash, question_preview, blocked, checkpoint,
+     rule_id, tool_calls, tokens, latency_ms, trace, agent, source, actor_role)
+VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
+"""
+
+
+# ── PostgresStore ────────────────────────────────────────────────────
+
+
+class PostgresStore:
+    """Persists telemetry to PostgreSQL via psycopg2.
+
+    Lifecycle: __init__ (stores config) → connect() → use → close().
+    The constructor does NOT open a connection — call connect() after fork.
+    """
+
+    def __init__(
+        self,
+        dsn: str,
+        trace_mode: str = "blocked",
+        sync_commit: bool = True,
+        sample_rate: int = 10,
+    ) -> None:
+        self._dsn = dsn
+        self._trace_mode = trace_mode
+        self._sync_commit = sync_commit
+        self._sample_rate = sample_rate
+        self._conn: Any = None
+        self._executor: ThreadPoolExecutor | None = None
+
+    # ── Lifecycle ────────────────────────────────────────────────
+
+    def connect(self) -> None:
+        """Open connection, create tables/indexes, start write executor."""
+        import psycopg2
+
+        self._conn = psycopg2.connect(self._dsn)
+        if not self._sync_commit:
+            with self._conn.cursor() as cur:
+                cur.execute("SET synchronous_commit = off")
+            self._conn.commit()
+        self._create_tables()
+        self._executor = ThreadPoolExecutor(max_workers=1)
+
+    def close(self) -> None:
+        """Drain pending writes, then close the connection."""
+        if self._executor is not None:
+            self._executor.shutdown(wait=True)
+            self._executor = None
+        if self._conn is not None:
+            try:
+                self._conn.close()
+            except Exception:
+                pass
+            self._conn = None
+
+    def _reconnect(self) -> None:
+        """Replace a dead connection. Called on OperationalError."""
+        import psycopg2
+
+        if self._conn is not None:
+            try:
+                self._conn.close()
+            except Exception:
+                pass
+        self._conn = psycopg2.connect(self._dsn)
+        if not self._sync_commit:
+            with self._conn.cursor() as cur:
+                cur.execute("SET synchronous_commit = off")
+            self._conn.commit()
+
+    def _create_tables(self) -> None:
+        with self._conn.cursor() as cur:
+            cur.execute(_CREATE_FIREWALL_EVENTS)
+            cur.execute(_CREATE_INVOCATIONS)
+            cur.execute(_CREATE_USER_RISK_SCORES)
+            cur.execute(_ALTER_ADD_PARDONED_UNTIL)
+            for idx in _INDEXES + _RISK_SCORE_INDEXES:
+                cur.execute(idx)
+        self._conn.commit()
+
+    # ── Trace policy ─────────────────────────────────────────────
+
+    def _should_store_trace(self, blocked: bool, invocation_id: str) -> bool:
+        if self._trace_mode == "all":
+            return True
+        if self._trace_mode == "off":
+            return False
+        if self._trace_mode == "blocked":
+            return blocked
+        if self._trace_mode == "sampled":
+            if blocked:
+                return True
+            return hash(invocation_id) % self._sample_rate == 0
+        return blocked  # unknown mode → fall back to "blocked"
+
+    # ── Internal insert helpers (run on executor thread) ─────────
+
+    def _do_insert_firewall_event(self, params: tuple) -> None:
+        import psycopg2
+
+        try:
+            with self._conn.cursor() as cur:
+                cur.execute(_INSERT_FIREWALL_EVENT, params)
+            self._conn.commit()
+        except (psycopg2.OperationalError, psycopg2.InterfaceError):
+            try:
+                self._conn.rollback()
+            except Exception:
+                pass
+            self._reconnect()
+            try:
+                with self._conn.cursor() as cur:
+                    cur.execute(_INSERT_FIREWALL_EVENT, params)
+                self._conn.commit()
+            except Exception:
+                logger.warning(
+                    "Firewall event insert failed after reconnect", exc_info=True
+                )
+        except Exception:
+            try:
+                self._conn.rollback()
+            except Exception:
+                pass
+            logger.warning("Firewall event insert failed", exc_info=True)
+
+    def _do_insert_invocation(self, params: tuple) -> None:
+        import psycopg2
+
+        try:
+            with self._conn.cursor() as cur:
+                cur.execute(_INSERT_INVOCATION, params)
+            self._conn.commit()
+        except (psycopg2.OperationalError, psycopg2.InterfaceError):
+            try:
+                self._conn.rollback()
+            except Exception:
+                pass
+            self._reconnect()
+            try:
+                with self._conn.cursor() as cur:
+                    cur.execute(_INSERT_INVOCATION, params)
+                self._conn.commit()
+            except Exception:
+                logger.warning(
+                    "Invocation insert failed after reconnect", exc_info=True
+                )
+        except Exception:
+            try:
+                self._conn.rollback()
+            except Exception:
+                pass
+            logger.warning("Invocation insert failed", exc_info=True)
+
+    # ── Done callback for submit_* futures ───────────────────────
+
+    @staticmethod
+    def _on_done(fut: Any) -> None:
+        exc = fut.exception()
+        if exc is not None:
+            logger.warning("EventStore async write failed: %s", exc)
+
+    # ── Public API: synchronous ──────────────────────────────────
+
+    def record_firewall_event(self, **kwargs: Any) -> None:
+
+        params = (
+            str(uuid.uuid4()),
+            kwargs.get("invocation_id"),
+            kwargs["timestamp"],
+            kwargs["checkpoint"],
+            kwargs["rule_id"],
+            kwargs["rule_name"],
+            kwargs.get("excerpt"),
+            kwargs.get("user_hash"),
+            kwargs.get("ip"),
+            kwargs.get("agent"),
+            kwargs.get("source", "runtime"),
+            kwargs.get("actor_role"),
+        )
+        self._do_insert_firewall_event(params)
+
+    def record_invocation(self, **kwargs: Any) -> None:
+        from psycopg2.extras import Json
+
+        blocked = kwargs["blocked"]
+        invocation_id = kwargs["invocation_id"]
+        trace = kwargs.get("trace")
+        if trace is not None and not self._should_store_trace(blocked, invocation_id):
+            trace = None
+
+        params = (
+            invocation_id,
+            kwargs["timestamp"],
+            kwargs.get("user_hash"),
+            kwargs.get("question_preview"),
+            blocked,
+            kwargs.get("checkpoint"),
+            kwargs.get("rule_id"),
+            kwargs.get("tool_calls", 0),
+            kwargs.get("tokens", 0),
+            kwargs.get("latency_ms", 0.0),
+            Json(trace) if trace is not None else None,
+            kwargs.get("agent"),
+            kwargs.get("source", "runtime"),
+            kwargs.get("actor_role"),
+        )
+        self._do_insert_invocation(params)
+
+    # ── Public API: fire-and-forget ──────────────────────────────
+
+    def submit_firewall_event(self, **kwargs: Any) -> None:
+        if self._executor is None:
+            return
+        params = (
+            str(uuid.uuid4()),
+            kwargs.get("invocation_id"),
+            kwargs["timestamp"],
+            kwargs["checkpoint"],
+            kwargs["rule_id"],
+            kwargs["rule_name"],
+            kwargs.get("excerpt"),
+            kwargs.get("user_hash"),
+            kwargs.get("ip"),
+            kwargs.get("agent"),
+            kwargs.get("source", "runtime"),
+            kwargs.get("actor_role"),
+        )
+        fut = self._executor.submit(self._do_insert_firewall_event, params)
+        fut.add_done_callback(self._on_done)
+
+    def submit_invocation(self, **kwargs: Any) -> None:
+        if self._executor is None:
+            return
+        from psycopg2.extras import Json
+
+        blocked = kwargs["blocked"]
+        invocation_id = kwargs["invocation_id"]
+        trace = kwargs.get("trace")
+        if trace is not None and not self._should_store_trace(blocked, invocation_id):
+            trace = None
+
+        params = (
+            invocation_id,
+            kwargs["timestamp"],
+            kwargs.get("user_hash"),
+            kwargs.get("question_preview"),
+            blocked,
+            kwargs.get("checkpoint"),
+            kwargs.get("rule_id"),
+            kwargs.get("tool_calls", 0),
+            kwargs.get("tokens", 0),
+            kwargs.get("latency_ms", 0.0),
+            Json(trace) if trace is not None else None,
+            kwargs.get("agent"),
+            kwargs.get("source", "runtime"),
+            kwargs.get("actor_role"),
+        )
+        fut = self._executor.submit(self._do_insert_invocation, params)
+        fut.add_done_callback(self._on_done)
+
+    # ── Risk score persistence ────────────────────────────────────
+
+    def upsert_risk_score(self, risk_score: Any) -> None:
+        """Persist a RiskScore to user_risk_scores (upsert on user_hash PK)."""
+        from psycopg2.extras import Json
+
+        factors_data = [
+            {
+                "name": f.name,
+                "description": f.description,
+                "contribution": f.contribution,
+                "signal_value": f.signal_value,
+            }
+            for f in risk_score.factors
+        ]
+        params = (
+            risk_score.user_hash,
+            risk_score.score,
+            risk_score.level.value,
+            Json(factors_data),
+            risk_score.computed_at,
+        )
+        try:
+            with self._conn.cursor() as cur:
+                cur.execute(_UPSERT_RISK_SCORE, params)
+            self._conn.commit()
+        except Exception:
+            try:
+                self._conn.rollback()
+            except Exception:
+                pass
+            logger.warning(
+                "upsert_risk_score failed for %s", risk_score.user_hash, exc_info=True
+            )
+
+    def get_risk_score(self, user_hash: str) -> Any:
+        """Fetch the stored RiskScore for user_hash, or None if not yet computed."""
+        from firewall_sdk.anomaly import RiskFactor, RiskLevel, RiskScore
+
+        try:
+            with self._conn.cursor() as cur:
+                cur.execute(_SELECT_RISK_SCORE, (user_hash,))
+                row = cur.fetchone()
+        except Exception:
+            logger.warning("get_risk_score failed for %s", user_hash, exc_info=True)
+            return None
+
+        if row is None:
+            return None
+
+        uh, score, level_str, factors_json, computed_at = row
+        factors = [
+            RiskFactor(
+                name=f["name"],
+                description=f["description"],
+                contribution=f["contribution"],
+                signal_value=f["signal_value"],
+            )
+            for f in (factors_json or [])
+        ]
+        return RiskScore(
+            user_hash=uh,
+            score=float(score),
+            level=RiskLevel(level_str),
+            factors=factors,
+            computed_at=computed_at,
+        )
+
+    def pardon_user(self, user_hash: str, hours: int = 1) -> None:
+        """Set a manual pardon on user_hash — overrides computed risk to LOW for `hours` hours."""
+        from datetime import timedelta, timezone
+
+        pardoned_until = datetime.now(timezone.utc) + timedelta(hours=hours)
+        try:
+            with self._conn.cursor() as cur:
+                cur.execute(
+                    """
+                    INSERT INTO user_risk_scores (user_hash, score, level, factors, computed_at, pardoned_until)
+                    VALUES (%s, 0.0, 'LOW', '[]', NOW(), %s)
+                    ON CONFLICT (user_hash) DO UPDATE SET pardoned_until = EXCLUDED.pardoned_until
+                    """,
+                    (user_hash, pardoned_until),
+                )
+            self._conn.commit()
+        except Exception:
+            try:
+                self._conn.rollback()
+            except Exception:
+                pass
+            logger.warning("pardon_user failed for %s", user_hash, exc_info=True)
+
+    # ── Public API: batch inserts (synchronous, for scripts) ─────
+
+    def insert_batch_firewall_events(self, rows: list[dict]) -> None:
+        from psycopg2.extras import execute_values
+
+        if not rows:
+            return
+        values = [
+            (
+                str(uuid.uuid4()),
+                r.get("invocation_id"),
+                r["timestamp"],
+                r["checkpoint"],
+                r["rule_id"],
+                r["rule_name"],
+                r.get("excerpt"),
+                r.get("user_hash"),
+                r.get("ip"),
+                r.get("agent"),
+                r.get("source", "runtime"),
+                r.get("actor_role"),
+            )
+            for r in rows
+        ]
+        with self._conn.cursor() as cur:
+            execute_values(
+                cur,
+                """INSERT INTO firewall_events
+                    (id, invocation_id, timestamp, checkpoint, rule_id, rule_name,
+                     excerpt, user_hash, ip, agent, source, actor_role)
+                VALUES %s""",
+                values,
+            )
+        self._conn.commit()
+
+    def insert_batch_invocations(self, rows: list[dict]) -> None:
+        from psycopg2.extras import execute_values, Json
+
+        if not rows:
+            return
+        values = []
+        for r in rows:
+            blocked = r["blocked"]
+            invocation_id = r["invocation_id"]
+            trace = r.get("trace")
+            if trace is not None and not self._should_store_trace(
+                blocked, invocation_id
+            ):
+                trace = None
+            values.append(
+                (
+                    invocation_id,
+                    r["timestamp"],
+                    r.get("user_hash"),
+                    r.get("question_preview"),
+                    blocked,
+                    r.get("checkpoint"),
+                    r.get("rule_id"),
+                    r.get("tool_calls", 0),
+                    r.get("tokens", 0),
+                    r.get("latency_ms", 0.0),
+                    Json(trace) if trace is not None else None,
+                    r.get("agent"),
+                    r.get("source", "runtime"),
+                    r.get("actor_role"),
+                )
+            )
+        with self._conn.cursor() as cur:
+            execute_values(
+                cur,
+                """INSERT INTO invocations
+                    (id, timestamp, user_hash, question_preview, blocked, checkpoint,
+                     rule_id, tool_calls, tokens, latency_ms, trace, agent, source, actor_role)
+                VALUES %s""",
+                values,
+            )
+        self._conn.commit()