AI System Audit Trail Design

"""
AI system audit trail implementation.
 
Provides structured logging for AI system events with
tamper-evident integrity verification.
"""
import hashlib
import json
import time
import uuid
from dataclasses import dataclass, field, asdict
from typing import Any
from enum import Enum
 
 
class AuditEventType(Enum):
    INFERENCE_REQUEST = "inference_request"
    INFERENCE_RESPONSE = "inference_response"
    GUARDRAIL_TRIGGER = "guardrail_trigger"
    MODEL_LOAD = "model_load"
    MODEL_UPDATE = "model_update"
    CONFIG_CHANGE = "config_change"
    ADMIN_ACTION = "admin_action"
    ERROR = "error"
    AUTHENTICATION = "authentication"
    RATE_LIMIT = "rate_limit"
 
 
@dataclass
class AuditEvent:
    """A single audit trail event."""
    event_id: str
    event_type: AuditEventType
    timestamp: float
    model_id: str
    model_version: str
    event_data: dict[str, Any]
    user_id: str | None = None
    session_id: str | None = None
    source_ip: str | None = None
    parent_event_id: str | None = None  # For correlating request/response pairs
    integrity_hash: str = ""  # Computed after creation
 
    def compute_integrity_hash(self, previous_hash: str = "") -> str:
        """
        Compute a chained integrity hash for tamper evidence.
 
        Each event's hash includes the previous event's hash,
        creating a hash chain similar to a blockchain. Tampering
        with any event breaks the chain.
        """
        content = json.dumps({
            "event_id": self.event_id,
            "event_type": self.event_type.value,
            "timestamp": self.timestamp,
            "model_id": self.model_id,
            "event_data": self.event_data,
            "user_id": self.user_id,
            "previous_hash": previous_hash,
        }, sort_keys=True)
        self.integrity_hash = hashlib.sha256(content.encode()).hexdigest()
        return self.integrity_hash
 
 
class AuditTrailLogger:
    """
    Structured audit trail logger for AI systems.
 
    Provides tamper-evident logging with hash chaining,
    structured event types, and queryable output.
    """
 
    def __init__(self, log_sink: "AuditLogSink"):
        self.sink = log_sink
        self.previous_hash = ""
 
    def log_inference_request(
        self,
        model_id: str,
        model_version: str,
        input_text: str | None = None,
        input_hash: str | None = None,
        user_id: str | None = None,
        session_id: str | None = None,
        source_ip: str | None = None,
        parameters: dict | None = None,
    ) -> str:
        """
        Log an inference request.
 
        For privacy-sensitive deployments, log the input_hash
        instead of the full input_text.
        """
        event_id = str(uuid.uuid4())
 
        event_data = {
            "input_hash": input_hash or (
                hashlib.sha256(input_text.encode()).hexdigest()
                if input_text else None
            ),
            "input_length": len(input_text) if input_text else None,
            "parameters": parameters or {},
        }
 
        # Only include full text if retention policy allows
        if input_text and self._should_retain_full_text():
            event_data["input_text"] = input_text
 
        event = AuditEvent(
            event_id=event_id,
            event_type=AuditEventType.INFERENCE_REQUEST,
            timestamp=time.time(),
            model_id=model_id,
            model_version=model_version,
            event_data=event_data,
            user_id=user_id,
            session_id=session_id,
            source_ip=source_ip,
        )
 
        self._emit(event)
        return event_id
 
    def log_inference_response(
        self,
        request_event_id: str,
        model_id: str,
        model_version: str,
        output_text: str | None = None,
        output_hash: str | None = None,
        latency_ms: float = 0.0,
        token_count: int = 0,
        finish_reason: str = "",
        guardrail_results: dict | None = None,
    ) -> str:
        event_id = str(uuid.uuid4())
 
        event_data = {
            "output_hash": output_hash or (
                hashlib.sha256(output_text.encode()).hexdigest()
                if output_text else None
            ),
            "output_length": len(output_text) if output_text else None,
            "latency_ms": latency_ms,
            "token_count": token_count,
            "finish_reason": finish_reason,
            "guardrail_results": guardrail_results,
        }
 
        event = AuditEvent(
            event_id=event_id,
            event_type=AuditEventType.INFERENCE_RESPONSE,
            timestamp=time.time(),
            model_id=model_id,
            model_version=model_version,
            event_data=event_data,
            parent_event_id=request_event_id,
        )
 
        self._emit(event)
        return event_id
 
    def log_guardrail_trigger(
        self,
        request_event_id: str,
        model_id: str,
        model_version: str,
        guardrail_name: str,
        trigger_reason: str,
        action_taken: str,
        scores: dict | None = None,
    ) -> str:
        event_id = str(uuid.uuid4())
 
        event = AuditEvent(
            event_id=event_id,
            event_type=AuditEventType.GUARDRAIL_TRIGGER,
            timestamp=time.time(),
            model_id=model_id,
            model_version=model_version,
            event_data={
                "guardrail_name": guardrail_name,
                "trigger_reason": trigger_reason,
                "action_taken": action_taken,
                "scores": scores or {},
            },
            parent_event_id=request_event_id,
        )
 
        self._emit(event)
        return event_id
 
    def log_model_update(
        self,
        model_id: str,
        old_version: str,
        new_version: str,
        update_type: str,
        authorized_by: str,
        change_description: str,
        weight_hash: str | None = None,
    ) -> str:
        event_id = str(uuid.uuid4())
 
        event = AuditEvent(
            event_id=event_id,
            event_type=AuditEventType.MODEL_UPDATE,
            timestamp=time.time(),
            model_id=model_id,
            model_version=new_version,
            event_data={
                "old_version": old_version,
                "new_version": new_version,
                "update_type": update_type,
                "authorized_by": authorized_by,
                "change_description": change_description,
                "weight_hash": weight_hash,
            },
        )
 
        self._emit(event)
        return event_id
 
    def log_config_change(
        self,
        model_id: str,
        model_version: str,
        config_key: str,
        old_value: Any,
        new_value: Any,
        changed_by: str,
    ) -> str:
        event_id = str(uuid.uuid4())
 
        event = AuditEvent(
            event_id=event_id,
            event_type=AuditEventType.CONFIG_CHANGE,
            timestamp=time.time(),
            model_id=model_id,
            model_version=model_version,
            event_data={
                "config_key": config_key,
                "old_value": str(old_value),
                "new_value": str(new_value),
                "changed_by": changed_by,
            },
        )
 
        self._emit(event)
        return event_id
 
    def _emit(self, event: AuditEvent) -> None:
        """Compute integrity hash and write to sink."""
        self.previous_hash = event.compute_integrity_hash(self.previous_hash)
        self.sink.write(event)
 
    def _should_retain_full_text(self) -> bool:
        """Check data retention policy for full input/output text."""
        # This should be configurable based on deployment context
        # and data classification requirements
        return True

@dataclass
class ModelLifecycleEvent:
    """Events in the AI model lifecycle that must be audited."""
    event_type: str  # "training_start", "training_end", "evaluation", "deployment", etc.
    timestamp: float
    model_id: str
    version: str
    actor: str  # Who initiated the action
    details: dict[str, Any]
 
# Key lifecycle events to audit:
LIFECYCLE_EVENTS = {
    "training_start": "Model training initiated",
    "training_end": "Model training completed",
    "evaluation": "Model evaluated on benchmark",
    "checkpoint_saved": "Model checkpoint written to storage",
    "deployment": "Model deployed to serving infrastructure",
    "rollback": "Model rolled back to previous version",
    "retirement": "Model removed from serving",
    "fine_tuning": "Model fine-tuned on new data",
    "quantization": "Model quantized for deployment",
    "safety_eval": "Safety evaluation completed",
}

class AuditChainVerifier:
    """Verify the integrity of an audit trail hash chain."""
 
    def verify_chain(self, events: list[AuditEvent]) -> dict:
        """
        Verify that the audit trail hash chain is intact.
 
        Recomputes each event's hash using the previous event's hash
        and compares against the stored hash.
        """
        if not events:
            return {"status": "EMPTY", "events_checked": 0}
 
        previous_hash = ""
        breaks = []
 
        for i, event in enumerate(events):
            expected_hash = self._recompute_hash(event, previous_hash)
 
            if expected_hash != event.integrity_hash:
                breaks.append({
                    "position": i,
                    "event_id": event.event_id,
                    "expected_hash": expected_hash,
                    "stored_hash": event.integrity_hash,
                })
 
            previous_hash = event.integrity_hash
 
        return {
            "status": "INTACT" if not breaks else "BROKEN",
            "events_checked": len(events),
            "chain_breaks": breaks,
            "first_break_position": breaks[0]["position"] if breaks else None,
        }
 
    def _recompute_hash(self, event: AuditEvent, previous_hash: str) -> str:
        content = json.dumps({
            "event_id": event.event_id,
            "event_type": event.event_type.value,
            "timestamp": event.timestamp,
            "model_id": event.model_id,
            "event_data": event.event_data,
            "user_id": event.user_id,
            "previous_hash": previous_hash,
        }, sort_keys=True)
        return hashlib.sha256(content.encode()).hexdigest()

class AuditLogSink:
    """Abstract base for audit log storage backends."""
 
    def write(self, event: AuditEvent) -> None:
        raise NotImplementedError
 
    def query(
        self,
        start_time: float | None = None,
        end_time: float | None = None,
        event_types: list[AuditEventType] | None = None,
        model_id: str | None = None,
        user_id: str | None = None,
        limit: int = 1000,
    ) -> list[AuditEvent]:
        raise NotImplementedError
 
 
class FileAuditLogSink(AuditLogSink):
    """
    Append-only file-based audit log sink.
 
    Suitable for single-node deployments or as a local buffer
    before forwarding to a centralized log system.
    """
 
    def __init__(self, log_dir: str):
        from pathlib import Path
        self.log_dir = Path(log_dir)
        self.log_dir.mkdir(parents=True, exist_ok=True)
        self.current_file = self.log_dir / f"audit_{int(time.time())}.jsonl"
 
    def write(self, event: AuditEvent) -> None:
        line = json.dumps(asdict(event), default=str)
        with open(self.current_file, "a") as f:
            f.write(line + "\n")
 
    def query(
        self,
        start_time: float | None = None,
        end_time: float | None = None,
        event_types: list[AuditEventType] | None = None,
        model_id: str | None = None,
        user_id: str | None = None,
        limit: int = 1000,
    ) -> list[dict]:
        results = []
        for log_file in sorted(self.log_dir.glob("audit_*.jsonl")):
            with open(log_file) as f:
                for line in f:
                    event = json.loads(line)
                    if start_time and event["timestamp"] < start_time:
                        continue
                    if end_time and event["timestamp"] > end_time:
                        continue
                    if event_types and event["event_type"] not in [
                        et.value for et in event_types
                    ]:
                        continue
                    if model_id and event["model_id"] != model_id:
                        continue
                    if user_id and event.get("user_id") != user_id:
                        continue
                    results.append(event)
                    if len(results) >= limit:
                        return results
        return results

def reconstruct_incident_timeline(
    sink: AuditLogSink,
    incident_start: float,
    incident_end: float,
    model_id: str,
) -> list[dict]:
    """
    Query the audit trail to reconstruct an incident timeline.
 
    Returns a chronologically ordered list of all events
    related to the specified model during the incident window.
    """
    events = sink.query(
        start_time=incident_start,
        end_time=incident_end,
        model_id=model_id,
        limit=10000,
    )
 
    # Sort by timestamp
    events.sort(key=lambda e: e["timestamp"])
 
    # Annotate events with time deltas
    timeline = []
    for i, event in enumerate(events):
        entry = {
            "sequence": i + 1,
            "timestamp": event["timestamp"],
            "time_since_start": event["timestamp"] - incident_start,
            "event_type": event["event_type"],
            "event_id": event["event_id"],
            "summary": _summarize_event(event),
        }
        if i > 0:
            entry["time_since_previous"] = (
                event["timestamp"] - events[i-1]["timestamp"]
            )
        timeline.append(entry)
 
    return timeline
 
 
def _summarize_event(event: dict) -> str:
    """Generate a human-readable summary of an audit event."""
    etype = event["event_type"]
    data = event.get("event_data", {})
 
    if etype == "inference_request":
        return f"Inference request (input length: {data.get('input_length', '?')})"
    if etype == "guardrail_trigger":
        return f"Guardrail '{data.get('guardrail_name')}' triggered: {data.get('trigger_reason')}"
    if etype == "model_update":
        return f"Model updated: {data.get('old_version')} -> {data.get('new_version')}"
    if etype == "config_change":
        return f"Config changed: {data.get('config_key')}"
    return f"Event: {etype}"