ITWorx_MCP_Hub-Public/packages/hub_policy/audit.py

"""Sanitized append-only audit ingestion, search, retention, and export metadata."""

from __future__ import annotations

import hashlib
import json
import re
from collections.abc import Mapping
from datetime import UTC, datetime, timedelta
from typing import Any, Protocol
from uuid import UUID, uuid4

from pydantic import BaseModel, ConfigDict, Field

from .engine import DecisionTrace

_SECRET_VALUE = re.compile(r"(?:bearer\s+|eyJ[A-Za-z0-9_-]{20,}|-----BEGIN[^-]+-----)", re.I)

_SENSITIVE = re.compile(
    r"(?:password|secret|token|private|authorization|cookie|credential|api[_-]?key)", re.I
)


def sanitize_arguments(value: Any, *, max_bytes: int = 16_384, _depth: int = 0) -> Any:
    if _depth > 5:
        return "<depth-limited>"
    result: Any
    if isinstance(value, Mapping):
        result = {
            str(key)[:80]: (
                "<redacted>"
                if _SENSITIVE.search(str(key))
                else sanitize_arguments(item, max_bytes=max_bytes, _depth=_depth + 1)
            )
            for key, item in list(value.items())[:64]
        }
    elif isinstance(value, list | tuple):
        result = [
            sanitize_arguments(item, max_bytes=max_bytes, _depth=_depth + 1)
            for item in list(value)[:64]
        ]
    elif isinstance(value, str | int | float | bool) or value is None:
        if isinstance(value, str) and _SECRET_VALUE.search(value):
            result = "<redacted>"
        else:
            result = value if not isinstance(value, str) else value[:2048]
    else:
        result = f"<{type(value).__name__}>"
    if len(json.dumps(result, sort_keys=True, default=str).encode()) > max_bytes:
        return "<arguments-truncated>"
    return result


class AuditEvent(BaseModel):
    model_config = ConfigDict(extra="forbid", frozen=True)

    id: UUID = Field(default_factory=uuid4)
    occurred_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
    trace_id: str = Field(min_length=1, max_length=128)
    event_type: str = Field(default="mcp.tool.invoked", max_length=128)
    workspace: str = Field(min_length=1, max_length=128)
    project: str = Field(min_length=1, max_length=128)
    environment: str = Field(min_length=1, max_length=32)
    principal: str = Field(min_length=1, max_length=256)
    client_id: UUID | None = None
    tool: str = Field(min_length=1, max_length=256)
    policy_decision: str
    matched_rule_ids: tuple[str, ...] = ()
    policy_trace: tuple[dict[str, Any], ...] = ()
    arguments_hash: str = Field(pattern=r"^[0-9a-f]{64}$")
    sanitized_arguments: dict[str, Any] = Field(default_factory=dict)
    duration_ms: int = Field(ge=0, le=120_000, default=0)
    outcome: str = Field(pattern=r"^(success|denied|error|rate_limited)$")
    error_code: str | None = None
    retention_until: datetime | None = None

    @classmethod
    def from_context(
        cls,
        context: Any,
        trace: DecisionTrace,
        *,
        outcome: str,
        started: float,
        error_code: str | None = None,
        result_summary: dict[str, Any] | None = None,
    ) -> AuditEvent:
        import time

        raw = json.dumps(dict(context.arguments), sort_keys=True, default=str).encode()
        return cls(
            trace_id=trace.trace_id,
            workspace=context.workspace,
            project=context.project,
            environment=context.environment,
            principal=context.principal,
            client_id=context.client_id,
            tool=context.tool,
            policy_decision=trace.decision.value,
            matched_rule_ids=trace.matched_rule_ids,
            policy_trace=tuple(item.model_dump(mode="json") for item in trace.steps),
            arguments_hash=hashlib.sha256(raw).hexdigest(),
            sanitized_arguments=sanitize_arguments(
                {
                    **dict(context.arguments),
                    **({"_result": result_summary} if result_summary else {}),
                }
            ),
            duration_ms=min(120_000, int((time.monotonic() - started) * 1000)),
            outcome=outcome,
            error_code=error_code,
            retention_until=datetime.now(UTC) + timedelta(days=30),
        )


class AuditPage(BaseModel):
    model_config = ConfigDict(extra="forbid", frozen=True)

    items: tuple[AuditEvent, ...]
    next_offset: int | None = None


class AuditLog(Protocol):
    def append(self, event: AuditEvent) -> None: ...
    def search(
        self, *, workspace: str, project: str, limit: int = 100, offset: int = 0
    ) -> AuditPage: ...
    def export_metadata(self, *, workspace: str, project: str) -> dict[str, Any]: ...


class InMemoryAuditLog:
    """Reference sink used by services and tests; it deliberately has no update/delete API."""

    def __init__(self) -> None:
        self.events: list[AuditEvent] = []

    def append(self, event: AuditEvent) -> None:
        if any(item.id == event.id for item in self.events):
            raise ValueError("audit event IDs are immutable")
        self.events.append(event)

    def search(
        self, *, workspace: str, project: str, limit: int = 100, offset: int = 0
    ) -> AuditPage:
        if not 1 <= limit <= 500 or offset < 0:
            raise ValueError("invalid audit pagination")
        items = tuple(
            item for item in self.events if item.workspace == workspace and item.project == project
        )
        selected = items[offset : offset + limit]
        next_offset = offset + len(selected) if offset + len(selected) < len(items) else None
        return AuditPage(items=selected, next_offset=next_offset)

    def export_metadata(self, *, workspace: str, project: str) -> dict[str, Any]:
        events = self.search(workspace=workspace, project=project, limit=500).items
        digest = hashlib.sha256(
            "".join(item.model_dump_json() for item in events).encode()
        ).hexdigest()
        return {
            "workspace": workspace,
            "project": project,
            "eventCount": len(events),
            "sha256": digest,
            "generatedAt": datetime.now(UTC).isoformat(),
            "format": "audit-event-v1",
            "access": "policy-restricted",
        }