audit-core/audit_core/postgres_backend.py

"""Durable PostgreSQL audit backend (AUDIT-WP-0005-T01).

Production custody. Runs on the Railiance shared PostgreSQL platform
(`rapp-postgres`), inside the database-per-consumer boundary fixed by that
repo's ADR-0001.

Two properties are pushed into the database rather than the application:

* **Idempotency** — ``INSERT ... ON CONFLICT DO NOTHING RETURNING`` resolves
  insert-or-detect in one statement, so two concurrent submissions of the same
  event cannot both be told they were first. The SQLite backend learned this
  the hard way; see its test.
* **Append-only custody** — a trigger rejects ``UPDATE`` and ``DELETE`` on the
  events table, so a leaked runtime credential can add records but cannot
  rewrite or erase existing ones. This is what lets ``RetentionPolicy`` declare
  ``immutable=True`` honestly; see :meth:`retention_policy` for its limits.
"""

from __future__ import annotations

import json
import os
from datetime import datetime, timezone
from typing import Any

from audit_core.interface import (
    AcceptResult,
    AuditEvent,
    BackendUnavailableError,
    EventConflictError,
    EventValidationError,
    RetentionPolicy,
    validate_event,
)
from audit_core.credentials import CredentialDirectory
from audit_core.integrity import (
    CHAIN_LOCK_KEY,
    GENESIS,
    ChainRow,
    attestation_from_report,
    chain_link,
    verify_rows,
)
from audit_core.redaction import Finding

try:  # pragma: no cover - import guard
    import psycopg
    from psycopg_pool import ConnectionPool
except ImportError as exc:  # pragma: no cover
    raise ImportError(
        "the postgres backend needs psycopg: pip install 'audit-core[postgres]'"
    ) from exc

DEFAULT_SCHEMA = "audit_core"

# Ordered, append-only. Each entry runs once and is recorded in
# schema_migrations. Never edit a released migration — add a new one.
MIGRATIONS: list[tuple[str, str]] = [
    (
        "0001-events",
        """
        CREATE TABLE IF NOT EXISTS {schema}.events (
          event_id       text PRIMARY KEY,
          payload_hash   text NOT NULL,
          accepted_at    timestamptz NOT NULL DEFAULT now(),
          observed_at    timestamptz,
          tenant         text NOT NULL,
          correlation_id text,
          source         text NOT NULL,
          action         text NOT NULL,
          record         jsonb NOT NULL
        );
        CREATE INDEX IF NOT EXISTS events_correlation_idx
          ON {schema}.events (correlation_id);
        -- Tenant keying is mandatory per business-app-service-contract 1.3;
        -- the index makes per-tenant export and retention tractable.
        CREATE INDEX IF NOT EXISTS events_tenant_idx ON {schema}.events (tenant);
        """,
    ),
    (
        "0002-append-only",
        """
        CREATE OR REPLACE FUNCTION {schema}.reject_mutation() RETURNS trigger AS $fn$
        BEGIN
          RAISE EXCEPTION 'audit events are append-only (attempted %)', TG_OP
            USING ERRCODE = 'restrict_violation';
        END;
        $fn$ LANGUAGE plpgsql;

        DROP TRIGGER IF EXISTS events_append_only ON {schema}.events;
        CREATE TRIGGER events_append_only
          BEFORE UPDATE OR DELETE ON {schema}.events
          FOR EACH ROW EXECUTE FUNCTION {schema}.reject_mutation();
        """,
    ),
    (
        "0003-dead-letters",
        """
        CREATE TABLE IF NOT EXISTS {schema}.dead_letters (
          id               bigserial PRIMARY KEY,
          event_id         text,
          received_at      timestamptz NOT NULL DEFAULT now(),
          sender           text,
          reason           text NOT NULL,
          payload_hash     text NOT NULL,
          payload          text,
          payload_withheld boolean NOT NULL DEFAULT false
        );
        CREATE INDEX IF NOT EXISTS dead_letters_event_idx
          ON {schema}.dead_letters (event_id);
        """,
    ),
    (
        "0004-secret-findings",
        """
        CREATE TABLE IF NOT EXISTS {schema}.secret_findings (
          sender      text NOT NULL,
          source      text NOT NULL,
          action      text NOT NULL,
          field_path  text NOT NULL,
          outcome     text NOT NULL,
          persisted   boolean NOT NULL DEFAULT false,
          occurrences bigint NOT NULL DEFAULT 0,
          first_seen  timestamptz NOT NULL DEFAULT now(),
          last_seen   timestamptz NOT NULL DEFAULT now(),
          PRIMARY KEY (sender, source, action, field_path, outcome)
        );
        """,
    ),
    (
        "0005-runtime-grants",
        """
        DO $grant$
        BEGIN
          IF EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'audit_core_app') THEN
            EXECUTE 'GRANT SELECT, INSERT ON {schema}.events TO audit_core_app';
            EXECUTE 'GRANT SELECT, INSERT ON {schema}.dead_letters TO audit_core_app';
            EXECUTE 'GRANT SELECT, INSERT, UPDATE ON {schema}.secret_findings TO audit_core_app';
            EXECUTE 'GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA {schema} TO audit_core_app';
          END IF;
        END
        $grant$;
        """,
    ),
    (
        "0006-chain",
        """
        ALTER TABLE {schema}.events ADD COLUMN IF NOT EXISTS chain_hash text;
        ALTER TABLE {schema}.events ADD COLUMN IF NOT EXISTS chain_prev text;
        """,
    ),
]

# Rejection reasons whose payload must never be persisted — storing the body of
# an event rejected *for carrying secret-shaped material* would write that
# material into the audit store.
WITHHOLD_PAYLOAD_REASONS = frozenset({"secret_shaped_field"})


class PostgresAuditBackend:
    """Audit custody in PostgreSQL."""

    def __init__(
        self,
        dsn: str | None = None,
        *,
        schema: str = DEFAULT_SCHEMA,
        retention_days: int | None = None,
        recoverable_days: int | None = 30,
        recoverable_source: str | None = (
            "resource-control/data/capability/platform-audit-storage.json"
            "#provisions[capability=data.backup]"
        ),
        recoverable_basis: str | None = "measured",
        min_size: int = 1,
        max_size: int = 8,
        statement_timeout_ms: int = 30_000,
        migrate: bool = True,
        credential_dir: str | None = None,
    ) -> None:
        # A mounted credential directory takes precedence: it is the only
        # source that can change while the process runs, which is what dynamic
        # leases require.
        self.credentials = (
            CredentialDirectory(credential_dir) if credential_dir else None
        )
        # An empty conninfo is valid: libpq then reads PGHOST/PGUSER/PGPASSWORD/
        # PGPORT/PGDATABASE from the environment. That is exactly the shape the
        # railiance-platform credential broker injects into a child process, so
        # a brokered lease needs no DSN assembled by hand — and no credential
        # ever passes through audit-core's own configuration.
        self.dsn = dsn if dsn is not None else os.environ.get("AUDIT_CORE_DATABASE_URL", "")
        if not self.dsn and not self.credentials and not _libpq_env_present():
            raise ValueError(
                "no connection information: set AUDIT_CORE_DATABASE_URL, or supply "
                "PGHOST/PGUSER/PGDATABASE (as the credential broker does)"
            )
        if not schema.isidentifier():
            raise ValueError(f"unsafe schema name: {schema!r}")
        self.schema = schema
        self.retention_days = retention_days
        self.recoverable_days = recoverable_days
        self.recoverable_source = recoverable_source
        self.recoverable_basis = recoverable_basis
        base_kwargs = {
            "autocommit": True,
            # A stalled write must surface as unavailable rather than hold a
            # request open indefinitely.
            "options": f"-c statement_timeout={int(statement_timeout_ms)}",
        }
        # psycopg_pool resolves a callable on every connection attempt, so
        # passing one is what makes a rotated lease take effect without a
        # restart. A fixed dict would freeze the credential at startup.
        kwargs: Any = base_kwargs
        if self.credentials:
            def kwargs() -> dict:  # type: ignore[misc]
                params = self.credentials.read()
                if "dsn" in params:
                    raise ValueError(
                        "a 'dsn' file is not supported with pooled connections; "
                        "supply username/password/host/port/dbname files"
                    )
                return {**base_kwargs, **params}

        try:
            self.pool = ConnectionPool(
                self.dsn,
                min_size=min_size,
                max_size=max_size,
                kwargs=kwargs,
                open=True,
                timeout=10,
            )
        except Exception as exc:  # psycopg raises a wide family here
            raise BackendUnavailableError(f"cannot connect: {exc}") from exc
        if migrate:
            self.migrate()

    # --- schema ------------------------------------------------------------

    def migrate(self) -> list[str]:
        """Apply pending migrations. Returns the ids applied this call."""
        applied: list[str] = []
        try:
            with self.pool.connection() as conn:
                migrate_role = os.environ.get("AUDIT_CORE_MIGRATE_ROLE", "").strip()
                if migrate_role:
                    if not migrate_role.isidentifier():
                        raise ValueError(f"unsafe migrate role: {migrate_role!r}")
                    # Leased login roles must SET ROLE to the group so new
                    # objects are owned by audit_core_migrate, not the lease.
                    conn.execute(f'SET ROLE "{migrate_role}"')
                conn.execute(f'CREATE SCHEMA IF NOT EXISTS "{self.schema}"')
                conn.execute(
                    f'CREATE TABLE IF NOT EXISTS "{self.schema}".schema_migrations ('
                    "  id text PRIMARY KEY, applied_at timestamptz NOT NULL DEFAULT now())"
                )
                done = {
                    row[0]
                    for row in conn.execute(
                        f'SELECT id FROM "{self.schema}".schema_migrations'
                    ).fetchall()
                }
                for migration_id, body in MIGRATIONS:
                    if migration_id in done:
                        continue
                    conn.execute(body.format(schema=f'"{self.schema}"'))
                    conn.execute(
                        f'INSERT INTO "{self.schema}".schema_migrations (id) VALUES (%s)',
                        (migration_id,),
                    )
                    applied.append(migration_id)
                self._backfill_chain(conn)
        except psycopg.Error as exc:
            raise BackendUnavailableError(f"migration failed: {exc}") from exc
        return applied

    @property
    def _events(self) -> str:
        return f'"{self.schema}".events'

    # --- contract ----------------------------------------------------------

    @property
    def retention_policy(self) -> RetentionPolicy:
        """Declared custody guarantees.

        ``immutable`` is True because migration 0002 installs a trigger that
        rejects UPDATE and DELETE, so no consumer credential can alter a stored
        record. It is not a claim against the database owner or a superuser,
        who can drop the trigger. ``tamper_evidence`` is True because a
        hash chain plus verify detects a rewritten payload, and a chain-head
        attestation outside this database detects a suffix rewrite that
        stays inside Postgres. It is not WORM or ``data.archive``.

        ``custody_class`` is ``operational``, not ``archive``. This store is
        durable append-only Postgres recovered through the platform
        ``data.backup`` provision. It is not ITC-CAP ``data.archive`` (WORM
        object storage, manifests, retrieval tests). Recoverable history is
        the cited platform window, not ``retention_days``.
        """
        return RetentionPolicy(
            custody_class="operational",
            retention_days=self.retention_days,
            immutable=True,
            tamper_evidence=True,
            durable=True,
            recoverable_days=self.recoverable_days,
            recoverable_source=self.recoverable_source,
            recoverable_basis=self.recoverable_basis,
        )

    def emit(self, event: AuditEvent) -> str:
        return self.accept(event, payload_hash=_record_hash(event)).reference

    def accept(self, event: AuditEvent, payload_hash: str) -> AcceptResult:
        validate_event(event)
        reference = f"audit:{event.event_id}"
        details = event.details if isinstance(event.details, dict) else {}
        record = event.as_record()
        try:
            with self.pool.connection() as conn, conn.transaction():
                conn.execute("SELECT pg_advisory_xact_lock(%s)", (CHAIN_LOCK_KEY,))
                head = conn.execute(
                    f"SELECT chain_hash FROM {self._events} "
                    "ORDER BY accepted_at DESC, event_id DESC LIMIT 1"
                ).fetchone()
                previous = head[0] if head and head[0] else GENESIS
                link = chain_link(previous, payload_hash, event.event_id)
                inserted = conn.execute(
                    f"""
                    INSERT INTO {self._events}
                      (event_id, payload_hash, observed_at, tenant, correlation_id,
                       source, action, record, chain_hash, chain_prev)
                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                    ON CONFLICT (event_id) DO NOTHING
                    RETURNING event_id
                    """,
                    (
                        event.event_id,
                        payload_hash,
                        _timestamp(event.observed_at),
                        event.tenant,
                        str(details.get("correlation_id") or "") or None,
                        event.source,
                        event.action,
                        json.dumps(record, sort_keys=True),
                        link,
                        previous,
                    ),
                ).fetchone()
                if inserted is not None:
                    return AcceptResult(duplicate=False, reference=reference)
                existing = conn.execute(
                    f"SELECT payload_hash FROM {self._events} WHERE event_id = %s",
                    (event.event_id,),
                ).fetchone()
        except psycopg.Error as exc:
            raise BackendUnavailableError(str(exc)) from exc

        if existing is None:
            raise BackendUnavailableError("event disappeared during accept")
        if existing[0] != payload_hash:
            raise EventConflictError(
                f"event_id {event.event_id} already held with a different payload"
            )
        return AcceptResult(duplicate=True, reference=reference)

    # --- operator surface --------------------------------------------------

    def get(self, event_id: str) -> dict | None:
        rows = self._query(
            f"SELECT record, accepted_at FROM {self._events} WHERE event_id = %s",
            (event_id,),
        )
        if not rows:
            return None
        return {"accepted_at": _iso(rows[0][1]), **rows[0][0]}

    def by_correlation(self, correlation_id: str, limit: int = 100) -> list[dict]:
        rows = self._query(
            f"SELECT record, accepted_at FROM {self._events} "
            "WHERE correlation_id = %s ORDER BY accepted_at, event_id LIMIT %s",
            (correlation_id, int(limit)),
        )
        return [{"accepted_at": _iso(at), **rec} for rec, at in rows]

    def replay(self, event_id: str) -> AcceptResult:
        """Re-submit a stored event through :meth:`accept`.

        Reconciliation, not re-creation: replaying an event already in custody
        must return ``duplicate=True`` against the same record. If it were to
        report a first acceptance, the store would be producing a second
        custody record for one source event — the exact failure the whole
        idempotency design exists to prevent.
        """
        rows = self._query(
            f"SELECT record, payload_hash FROM {self._events} WHERE event_id = %s",
            (event_id,),
        )
        if not rows:
            raise KeyError(event_id)
        record, payload_hash = rows[0]
        return self.accept(_event_from_record(record), payload_hash)

    def record_rejection(
        self,
        *,
        event_id: str | None,
        reason: str,
        payload_hash: str,
        sender: str | None = None,
        payload: str | None = None,
    ) -> None:
        withheld = reason in WITHHOLD_PAYLOAD_REASONS
        self._execute(
            f'INSERT INTO "{self.schema}".dead_letters '
            "(event_id, sender, reason, payload_hash, payload, payload_withheld) "
            "VALUES (%s, %s, %s, %s, %s, %s)",
            (event_id, sender, reason, payload_hash,
             None if withheld else payload, withheld),
        )

    def dead_letters(self, limit: int = 100) -> list[dict]:
        rows = self._query(
            "SELECT event_id, received_at, sender, reason, payload_hash, payload, "
            f'payload_withheld FROM "{self.schema}".dead_letters '
            "ORDER BY id DESC LIMIT %s",
            (int(limit),),
        )
        return [
            {
                "event_id": r[0], "received_at": _iso(r[1]), "sender": r[2],
                "reason": r[3], "payload_hash": r[4], "payload": r[5],
                "payload_withheld": bool(r[6]),
            }
            for r in rows
        ]

    def count_secret_findings(
        self, *, sender: str, source: str, action: str, outcome: str, findings
    ) -> None:
        for finding in findings:
            self._execute(
                f'INSERT INTO "{self.schema}".secret_findings '
                "(sender, source, action, field_path, outcome, persisted, occurrences) "
                "VALUES (%s, %s, %s, %s, %s, %s, 1) "
                "ON CONFLICT (sender, source, action, field_path, outcome) DO UPDATE "
                "SET occurrences = secret_findings.occurrences + 1, last_seen = now()",
                (sender, source, action, finding.path, outcome,
                 bool(getattr(finding, "in_persisted_data", False))),
            )

    def secret_findings(self, limit: int = 100) -> list[dict]:
        rows = self._query(
            "SELECT sender, source, action, field_path, outcome, persisted, "
            f'occurrences, first_seen, last_seen FROM "{self.schema}".secret_findings '
            "ORDER BY occurrences DESC, last_seen DESC LIMIT %s",
            (int(limit),),
        )
        return [
            {
                "sender": r[0], "source": r[1], "action": r[2], "field_path": r[3],
                "outcome": r[4], "persisted": bool(r[5]), "occurrences": r[6],
                "first_seen": _iso(r[7]), "last_seen": _iso(r[8]),
            }
            for r in rows
        ]

    # --- lifecycle ---------------------------------------------------------

    def verify_chain(self, attestation: dict | None = None):
        rows = self._query(
            f"SELECT event_id, payload_hash, chain_hash, chain_prev, accepted_at "
            f"FROM {self._events} ORDER BY accepted_at, event_id",
            (),
        )
        return verify_rows(
            [
                ChainRow(
                    event_id=r[0],
                    payload_hash=r[1],
                    chain_hash=r[2] or "",
                    chain_prev=r[3] or "",
                    accepted_at=_iso(r[4]),
                )
                for r in rows
            ],
            attestation=attestation,
        )

    def attest_chain(self) -> dict:
        return attestation_from_report(self.verify_chain())

    def health(self) -> None:
        self._query("SELECT 1", ())

    def _backfill_chain(self, conn) -> None:
        """Fill chain columns on rows accepted before migration 0006.

        The append-only trigger must be disabled for this UPDATE. A
        database owner can do that; this is a one-time migrate, not a
        runtime path.
        """
        missing = conn.execute(
            f"SELECT count(*) FROM {self._events} "
            "WHERE chain_hash IS NULL OR chain_prev IS NULL"
        ).fetchone()[0]
        if missing:
            conn.execute(
                f'ALTER TABLE {self._events} DISABLE TRIGGER events_append_only'
            )
            prev = GENESIS
            for event_id, payload_hash in conn.execute(
                f"SELECT event_id, payload_hash FROM {self._events} "
                "ORDER BY accepted_at, event_id"
            ).fetchall():
                link = chain_link(prev, payload_hash, event_id)
                conn.execute(
                    f"UPDATE {self._events} SET chain_prev = %s, chain_hash = %s "
                    "WHERE event_id = %s",
                    (prev, link, event_id),
                )
                prev = link
            conn.execute(
                f'ALTER TABLE {self._events} ENABLE TRIGGER events_append_only'
            )
        conn.execute(
            f"ALTER TABLE {self._events} ALTER COLUMN chain_hash SET NOT NULL"
        )
        conn.execute(
            f"ALTER TABLE {self._events} ALTER COLUMN chain_prev SET NOT NULL"
        )

    def close(self) -> None:
        self.pool.close()

    def _query(self, sql: str, params: tuple) -> list:
        try:
            with self.pool.connection() as conn:
                return conn.execute(sql, params).fetchall()
        except psycopg.Error as exc:
            raise BackendUnavailableError(str(exc)) from exc

    def _execute(self, sql: str, params: tuple) -> None:
        try:
            with self.pool.connection() as conn:
                conn.execute(sql, params)
        except psycopg.Error as exc:
            raise BackendUnavailableError(str(exc)) from exc


def _libpq_env_present() -> bool:
    """Whether libpq has enough in the environment to connect on its own."""
    return bool(os.environ.get("PGHOST") and os.environ.get("PGUSER"))


def _timestamp(value: str | None) -> datetime | None:
    if not value:
        return None
    try:
        return datetime.fromisoformat(str(value).replace("Z", "+00:00"))
    except ValueError:
        return None


def _iso(value: Any) -> str | None:
    if isinstance(value, datetime):
        return value.astimezone(timezone.utc).replace(microsecond=0).isoformat()
    return value


def _event_from_record(record: dict) -> AuditEvent:
    return AuditEvent(
        source=record["source"], action=record["action"], resource=record["resource"],
        outcome=record["outcome"], tenant=record["tenant"], scope=record["scope"],
        actor=record.get("actor"), reason=record.get("reason"),
        details=record.get("details") or {}, event_id=record["event_id"],
        observed_at=record["observed_at"], schema_version=record["schema_version"],
    )


def _record_hash(event: AuditEvent) -> str:
    import hashlib

    return hashlib.sha256(
        json.dumps(event.as_record(), sort_keys=True).encode("utf-8")
    ).hexdigest()