ModelForge/backend/src/modelforge_api/services/audit.py

"""Canonical, append-only audit writes, checkpoints and verification.

The chain head is a single serialization point. PostgreSQL writers use a transaction-scoped
advisory lock; SQLite test databases use an engine-scoped process lock held until the owning
transaction ends. An event and the durable singleton checkpoint advance in the same transaction.
"""

from __future__ import annotations

import hashlib
import json
import threading
import uuid
from collections.abc import Iterable
from dataclasses import dataclass
from datetime import UTC, datetime
from typing import Any, Protocol
from weakref import WeakKeyDictionary

from sqlalchemy import event, select, text
from sqlalchemy.engine import Connection, Engine
from sqlalchemy.orm import Session

from modelforge_api.domain.audit import (
    AUDIT_CHAIN_SINGLETON_ID,
    AUDIT_CURRENT_HASH_FORMAT,
    AUDIT_EMPTY_LEGACY_PREFIX_SEAL,
    AUDIT_HASH_FORMAT_V1,
    AUDIT_HASH_FORMAT_V2,
    AUDIT_LEGACY_PREFIX_DOMAIN,
    canonical_audit_payload_and_hash,
    canonical_audit_payload_text_and_hash,
    normalise_audit_event_id,
    normalise_audit_timestamp,
)
from modelforge_api.persistence.models import (
    AuditChainHead,
    AuditEvent,
    _append_canonical_audit_event,
)

AUDIT_CHAIN_POSTGRES_LOCK_KEY = int.from_bytes(b"MF_AUDIT", byteorder="big", signed=False)
_SQLITE_LOCK_INFO_KEY = "modelforge_audit_chain_lock"
_SQLITE_LOCKS: WeakKeyDictionary[Engine, threading.Lock] = WeakKeyDictionary()
_SQLITE_LOCKS_GUARD = threading.Lock()


class AuditChainIntegrityError(RuntimeError):
    """The writer cannot safely append to the authoritative audit chain."""


class AuditEventLike(Protocol):
    id: Any
    sequence: int
    occurred_at: Any
    correlation_id: str
    actor_type: str
    actor_id: str
    action: str
    resource_type: str
    resource_id: str | None
    outcome: str
    details: dict[str, Any]
    previous_event_hash: str | None
    event_hash: str
    hash_format: str
    canonical_payload: str | None


@dataclass(frozen=True, slots=True)
class AuditContext:
    """Authenticated/request identity seam for audit-producing services.

    Existing callers can keep passing their actor fields directly. The API authentication boundary
    can instead construct this context with its request correlation id without changing chain code.
    """

    actor_type: str
    actor_id: str
    correlation_id: str | None = None


@dataclass(frozen=True, slots=True)
class AuditChainCheckpoint:
    """Transport-neutral checkpoint used by local and restored-database verification."""

    singleton_id: int
    event_count: int
    last_sequence: int
    last_event_hash: str | None
    hash_format: str
    v2_start_sequence: int
    legacy_prefix_count: int
    legacy_prefix_seal: str


@dataclass(frozen=True, slots=True)
class AuditEventRecord:
    """Transport-neutral event used when recovery reads a database through psql."""

    id: Any
    sequence: int
    occurred_at: Any
    correlation_id: str
    actor_type: str
    actor_id: str
    action: str
    resource_type: str
    resource_id: str | None
    outcome: str
    details: dict[str, Any]
    previous_event_hash: str | None
    event_hash: str
    hash_format: str
    canonical_payload: str | None = None


def _decode_canonical_audit_payload(value: str) -> dict[str, Any]:
    def reject_duplicate_keys(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
        result: dict[str, Any] = {}
        for key, item in pairs:
            if key in result:
                raise ValueError(f"duplicate canonical audit payload key {key!r}")
            result[key] = item
        return result

    decoded = json.loads(value, object_pairs_hook=reject_duplicate_keys)
    if not isinstance(decoded, dict):
        raise ValueError("canonical audit payload is not an object")
    return decoded


def _event_payload_violations(audit_event: AuditEventLike) -> list[str]:
    """Verify one event's semantic payload and exact stored hash bytes."""

    identity = str(audit_event.id)
    if audit_event.hash_format == AUDIT_HASH_FORMAT_V1:
        violations: list[str] = []
        if audit_event.canonical_payload is not None:
            violations.append(f"legacy audit event {identity} unexpectedly stores a v2 payload")
        try:
            _, expected_hash = canonical_audit_payload_and_hash(
                correlation_id=audit_event.correlation_id,
                actor_type=audit_event.actor_type,
                actor_id=audit_event.actor_id,
                action=audit_event.action,
                resource_type=audit_event.resource_type,
                resource_id=audit_event.resource_id,
                outcome=audit_event.outcome,
                details=audit_event.details,
                previous_event_hash=audit_event.previous_event_hash,
                hash_format=audit_event.hash_format,
                event_id=audit_event.id,
                occurred_at=audit_event.occurred_at,
            )
        except (TypeError, ValueError):
            violations.append(f"audit event {identity} has a non-canonical payload")
        else:
            if audit_event.event_hash != expected_hash:
                violations.append(
                    f"audit event {identity} content hash does not match its payload"
                )
        return violations

    if audit_event.hash_format != AUDIT_HASH_FORMAT_V2:
        return [f"audit event {identity} has an unsupported hash format"]
    canonical_payload = audit_event.canonical_payload
    if not isinstance(canonical_payload, str):
        return [f"audit event {identity} has no exact v2 canonical payload"]
    violations = []
    if hashlib.sha256(canonical_payload.encode("utf-8")).hexdigest() != audit_event.event_hash:
        violations.append(f"audit event {identity} content hash does not match stored bytes")
    try:
        decoded_payload = _decode_canonical_audit_payload(canonical_payload)
        _, expected_text, _expected_hash = canonical_audit_payload_text_and_hash(
            correlation_id=audit_event.correlation_id,
            actor_type=audit_event.actor_type,
            actor_id=audit_event.actor_id,
            action=audit_event.action,
            resource_type=audit_event.resource_type,
            resource_id=audit_event.resource_id,
            outcome=audit_event.outcome,
            details=audit_event.details,
            previous_event_hash=audit_event.previous_event_hash,
            hash_format=audit_event.hash_format,
            event_id=audit_event.id,
            occurred_at=audit_event.occurred_at,
        )
        expected_payload = _decode_canonical_audit_payload(expected_text)
    except (json.JSONDecodeError, TypeError, ValueError):
        violations.append(f"audit event {identity} has a non-canonical payload")
    else:
        canonical_decoded = json.dumps(
            decoded_payload, sort_keys=True, separators=(",", ":"), ensure_ascii=False
        )
        canonical_expected = json.dumps(
            expected_payload, sort_keys=True, separators=(",", ":"), ensure_ascii=False
        )
        if canonical_decoded != canonical_expected:
            violations.append(
                f"audit event {identity} content hash bytes do not describe its event columns"
            )
    return violations


def _legacy_prefix_entry(event_row: AuditEventLike) -> bytes:
    payload = {
        "sequence": int(event_row.sequence),
        "id": normalise_audit_event_id(event_row.id),
        "occurred_at": normalise_audit_timestamp(event_row.occurred_at),
        "event_hash": str(event_row.event_hash),
    }
    return json.dumps(payload, sort_keys=True, separators=(",", ":")).encode("utf-8") + b"\n"


def legacy_audit_prefix_seal(events: Iterable[AuditEventLike]) -> str:
    """Seal immutable v1 identities/timestamps without rewriting approved legacy events."""

    digest = hashlib.sha256()
    digest.update(AUDIT_LEGACY_PREFIX_DOMAIN)
    for audit_event in events:
        digest.update(_legacy_prefix_entry(audit_event))
    return digest.hexdigest()


EMPTY_LEGACY_PREFIX_SEAL = AUDIT_EMPTY_LEGACY_PREFIX_SEAL


def audit_chain_violations(
    events: Iterable[AuditEventLike], checkpoint: AuditChainHead | AuditChainCheckpoint | None
) -> list[str]:
    """Verify sequence, links, versioned hashes, the prefix seal and durable checkpoint."""

    violations: list[str] = []
    if checkpoint is None:
        violations.append("audit chain checkpoint is missing")
        cutover = 1
        prefix_count = 0
    else:
        cutover = int(checkpoint.v2_start_sequence)
        prefix_count = int(checkpoint.legacy_prefix_count)
        if checkpoint.singleton_id != AUDIT_CHAIN_SINGLETON_ID:
            violations.append("audit chain checkpoint has an invalid singleton id")
        if checkpoint.hash_format != AUDIT_CURRENT_HASH_FORMAT:
            violations.append("audit chain checkpoint names an unsupported current hash format")
        if cutover < 1 or prefix_count != cutover - 1:
            violations.append("audit chain checkpoint has an invalid v2 cutover")

    previous: AuditEventLike | None = None
    seen_sequences: set[int] = set()
    seen_event_ids: set[str] = set()
    observed = 0
    prefix_observed = 0
    prefix_digest = hashlib.sha256()
    prefix_digest.update(AUDIT_LEGACY_PREFIX_DOMAIN)
    for expected_sequence, audit_event in enumerate(events, start=1):
        observed += 1
        identity = str(audit_event.id)
        try:
            normalised_identity = normalise_audit_event_id(audit_event.id)
        except ValueError:
            violations.append(f"audit event {identity} has a non-canonical UUID")
        else:
            if normalised_identity in seen_event_ids:
                violations.append(f"audit event UUID {normalised_identity} is duplicated")
            seen_event_ids.add(normalised_identity)
        sequence = int(audit_event.sequence)
        if sequence in seen_sequences:
            violations.append(f"audit sequence {sequence} is duplicated at event {identity}")
        seen_sequences.add(sequence)
        if sequence != expected_sequence:
            violations.append(
                f"audit event {identity} has sequence {sequence}; expected {expected_sequence}"
            )

        expected_link = previous.event_hash if previous is not None else None
        if audit_event.previous_event_hash != expected_link:
            position = "first event" if previous is None else f"event after {previous.id}"
            violations.append(
                f"audit event {identity} has an invalid previous hash for the {position}"
            )

        expected_format = AUDIT_HASH_FORMAT_V1 if sequence < cutover else AUDIT_HASH_FORMAT_V2
        if audit_event.hash_format != expected_format:
            violations.append(
                f"audit event {identity} uses {audit_event.hash_format!r}; "
                f"expected {expected_format!r} at sequence {sequence}"
            )
        violations.extend(_event_payload_violations(audit_event))

        if sequence <= prefix_count:
            prefix_observed += 1
            try:
                prefix_digest.update(_legacy_prefix_entry(audit_event))
            except (TypeError, ValueError):
                violations.append(
                    f"legacy audit event {identity} has a non-canonical identity or timestamp"
                )
        previous = audit_event

    if checkpoint is not None:
        observed_last_sequence = int(previous.sequence) if previous is not None else 0
        observed_last_hash = previous.event_hash if previous is not None else None
        if int(checkpoint.event_count) != observed:
            violations.append(
                f"audit checkpoint records {checkpoint.event_count} events; observed {observed}"
            )
        if int(checkpoint.last_sequence) != observed_last_sequence:
            violations.append(
                "audit checkpoint last sequence does not match the retained event suffix"
            )
        if checkpoint.last_event_hash != observed_last_hash:
            violations.append("audit checkpoint last hash does not match the retained event suffix")
        if prefix_observed != prefix_count:
            violations.append(
                f"audit checkpoint seals {prefix_count} legacy events; observed {prefix_observed}"
            )
        if checkpoint.legacy_prefix_seal != prefix_digest.hexdigest():
            violations.append("audit legacy-prefix seal does not match immutable legacy history")
    return violations


def _audit_engine(session: Session) -> Engine:
    bind = session.get_bind()
    if isinstance(bind, Connection):
        return bind.engine
    return bind


def _acquire_audit_write_lock(session: Session) -> None:
    """Hold the chain-head lock until the current root transaction completes."""

    engine = _audit_engine(session)
    dialect = engine.dialect.name
    if dialect == "postgresql":
        session.execute(
            text("SELECT pg_advisory_xact_lock(:lock_key)"),
            {"lock_key": AUDIT_CHAIN_POSTGRES_LOCK_KEY},
        )
        return
    if dialect != "sqlite":
        raise RuntimeError(f"audit-chain writes do not support the {dialect!r} database dialect")

    if _SQLITE_LOCK_INFO_KEY in session.info:
        return
    if not session.in_transaction():
        session.begin()
    with _SQLITE_LOCKS_GUARD:
        lock = _SQLITE_LOCKS.setdefault(engine, threading.Lock())
    lock.acquire()
    session.info[_SQLITE_LOCK_INFO_KEY] = lock


@event.listens_for(Session, "after_transaction_end")
def _release_sqlite_audit_write_lock(session: Session, transaction: Any) -> None:
    if transaction.parent is not None:
        return
    lock = session.info.pop(_SQLITE_LOCK_INFO_KEY, None)
    if lock is not None:
        lock.release()


def _checkpoint_from_mapping(row: Any) -> AuditChainCheckpoint:
    return AuditChainCheckpoint(
        singleton_id=int(row.singleton_id),
        event_count=int(row.event_count),
        last_sequence=int(row.last_sequence),
        last_event_hash=row.last_event_hash,
        hash_format=str(row.hash_format),
        v2_start_sequence=int(row.v2_start_sequence),
        legacy_prefix_count=int(row.legacy_prefix_count),
        legacy_prefix_seal=str(row.legacy_prefix_seal),
    )


def _load_checkpoint(session: Session) -> AuditChainCheckpoint:
    row = session.execute(
        select(
            AuditChainHead.singleton_id,
            AuditChainHead.event_count,
            AuditChainHead.last_sequence,
            AuditChainHead.last_event_hash,
            AuditChainHead.hash_format,
            AuditChainHead.v2_start_sequence,
            AuditChainHead.legacy_prefix_count,
            AuditChainHead.legacy_prefix_seal,
        ).where(AuditChainHead.singleton_id == AUDIT_CHAIN_SINGLETON_ID)
    ).one_or_none()
    if row is not None:
        return _checkpoint_from_mapping(row)
    raise AuditChainIntegrityError(
        "audit chain checkpoint is missing; only schema creation or the audited migration may seed it"
    )


def _assert_checkpoint_is_appendable(
    session: Session, checkpoint: AuditChainCheckpoint
) -> AuditEvent | None:
    """Validate the locked checkpoint and constant-size retained tail before append.

    Migration and recovery seal/verify the complete immutable prefix. Runtime append therefore
    proves the checkpoint shape, current tail hash/link and compare-and-set predecessor in O(1).
    A privileged edit in older middle history is intentionally the responsibility of the explicit
    strict invariant/recovery gates; ordinary SQLAlchemy audit DML is blocked separately.
    """

    violations: list[str] = []
    if checkpoint.singleton_id != AUDIT_CHAIN_SINGLETON_ID:
        violations.append("checkpoint singleton id is invalid")
    if checkpoint.hash_format != AUDIT_CURRENT_HASH_FORMAT:
        violations.append("checkpoint hash format is unsupported")
    if checkpoint.event_count < 0 or checkpoint.last_sequence < 0:
        violations.append("checkpoint counts cannot be negative")
    if checkpoint.event_count != checkpoint.last_sequence:
        violations.append("checkpoint event count and last sequence disagree")
    if checkpoint.v2_start_sequence < 1:
        violations.append("checkpoint v2 cutover is invalid")
    if checkpoint.legacy_prefix_count != checkpoint.v2_start_sequence - 1:
        violations.append("checkpoint legacy-prefix count and v2 cutover disagree")
    if checkpoint.legacy_prefix_count > checkpoint.event_count:
        violations.append("checkpoint legacy prefix exceeds the retained event count")
    if (
        len(checkpoint.legacy_prefix_seal) != 64
        or any(character not in "0123456789abcdef" for character in checkpoint.legacy_prefix_seal)
    ):
        violations.append("checkpoint legacy-prefix seal is malformed")
    if (
        checkpoint.legacy_prefix_count == 0
        and checkpoint.legacy_prefix_seal != EMPTY_LEGACY_PREFIX_SEAL
    ):
        violations.append("empty legacy-prefix checkpoint has the wrong seal")

    tail = list(
        session.scalars(
        select(AuditEvent)
            .order_by(AuditEvent.sequence.desc(), AuditEvent.id.desc())
            .limit(2)
            .execution_options(populate_existing=True)
        )
    )
    latest = tail[0] if tail else None
    if checkpoint.event_count == 0:
        if latest is not None:
            violations.append("empty checkpoint has a retained audit tail")
        if checkpoint.last_event_hash is not None:
            violations.append("empty checkpoint carries a last-event hash")
    else:
        if latest is None:
            violations.append("non-empty checkpoint has no retained audit tail")
        elif (
            latest.sequence != checkpoint.last_sequence
            or latest.event_hash != checkpoint.last_event_hash
        ):
            violations.append("checkpoint does not identify the current retained audit tail")
        if checkpoint.last_event_hash is None:
            violations.append("non-empty checkpoint has no last-event hash")
        elif len(checkpoint.last_event_hash) != 64 or any(
            character not in "0123456789abcdef"
            for character in checkpoint.last_event_hash
        ):
            violations.append("checkpoint last-event hash is malformed")

    if latest is not None:
        expected_format = (
            AUDIT_HASH_FORMAT_V1
            if latest.sequence < checkpoint.v2_start_sequence
            else AUDIT_HASH_FORMAT_V2
        )
        if latest.hash_format != expected_format:
            violations.append("retained audit tail uses the wrong hash format")
        payload_violations = _event_payload_violations(latest)
        violations.extend(
            f"retained audit tail: {violation}" for violation in payload_violations
        )

        if latest.sequence == 1:
            if latest.previous_event_hash is not None:
                violations.append("first retained audit event has a previous hash")
            if len(tail) != 1:
                violations.append("checkpoint count one has more than one retained event")
        elif latest.sequence > 1:
            if len(tail) != 2:
                violations.append("retained audit tail has no predecessor")
            else:
                predecessor = tail[1]
                if predecessor.sequence != latest.sequence - 1:
                    violations.append("retained audit tail predecessor is not contiguous")
                if latest.previous_event_hash != predecessor.event_hash:
                    violations.append("retained audit tail link does not match its predecessor")

    if violations:
        raise AuditChainIntegrityError(
            "audit checkpoint/tail failed pre-append verification: "
            + "; ".join(violations[:5])
        )
    return latest


class AuditWriter:
    def __init__(
        self,
        session: Session,
        actor_type: str | None = None,
        actor_id: str | None = None,
        *,
        context: AuditContext | None = None,
    ) -> None:
        if context is not None:
            if actor_type is not None or actor_id is not None:
                raise ValueError("pass either an audit context or actor fields, not both")
            resolved = context
        else:
            if actor_type is None or actor_id is None:
                raise ValueError("audit actor type and id are required")
            resolved = AuditContext(actor_type=actor_type, actor_id=actor_id)
        self.session = session
        self.context = resolved

    def write(
        self,
        action: str,
        resource_type: str,
        resource_id: str | None,
        details: dict[str, Any],
        outcome: str = "success",
    ) -> AuditEvent:
        _acquire_audit_write_lock(self.session)
        try:
            checkpoint = _load_checkpoint(self.session)
            previous = _assert_checkpoint_is_appendable(self.session, checkpoint)
            event_id = uuid.uuid4()
            occurred_at = datetime.now(UTC)
            return _append_canonical_audit_event(
                self.session,
                event_id=event_id,
                occurred_at=occurred_at,
                correlation_id=self.context.correlation_id or str(uuid.uuid4()),
                actor_type=self.context.actor_type,
                actor_id=self.context.actor_id,
                action=action,
                resource_type=resource_type,
                resource_id=resource_id,
                outcome=outcome,
                details=details,
                previous_event_hash=previous.event_hash if previous else None,
                expected_event_count=checkpoint.event_count,
                expected_last_sequence=checkpoint.last_sequence,
                expected_last_event_hash=checkpoint.last_event_hash,
                expected_hash_format=checkpoint.hash_format,
                expected_v2_start_sequence=checkpoint.v2_start_sequence,
                expected_legacy_prefix_count=checkpoint.legacy_prefix_count,
                expected_legacy_prefix_seal=checkpoint.legacy_prefix_seal,
            )
        except Exception:
            # A caller must never be able to catch an audit failure and commit an unaudited domain
            # mutation or a detached event. Roll back the complete owning transaction fail-closed.
            self.session.rollback()
            raise