Ontwerp van audit trails voor AI-systemen

"""
Implementatie van een audit trail voor AI-systemen.
 
Biedt gestructureerde logging voor AI-systeemgebeurtenissen met
manipulatiebestendige integriteitsverificatie.
"""
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:
    """Een enkele audit trail-gebeurtenis."""
    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  # Voor het correleren van request/response-paren
    integrity_hash: str = ""  # Berekend na aanmaak
 
    def compute_integrity_hash(self, previous_hash: str = "") -> str:
        """
        Bereken een geketende integriteitshash als manipulatiebewijs.
 
        De hash van elke gebeurtenis bevat de hash van de vorige gebeurtenis,
        waardoor een hash chain ontstaat vergelijkbaar met een blockchain. Manipuleren
        van een gebeurtenis verbreekt de keten.
        """
        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:
    """
    Gestructureerde audit trail-logger voor AI-systemen.
 
    Biedt manipulatiebestendige logging met hash chaining,
    gestructureerde gebeurtenistypen en bevraagbare 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 een inferentieverzoek.
 
        Voor privacygevoelige deployments: log de input_hash
        in plaats van de volledige 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 {},
        }
 
        # Neem volledige tekst alleen op als het bewaarbeleid dit toestaat
        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:
        """Bereken de integriteitshash en schrijf naar de sink."""
        self.previous_hash = event.compute_integrity_hash(self.previous_hash)
        self.sink.write(event)
 
    def _should_retain_full_text(self) -> bool:
        """Controleer het databewaarbeleid voor volledige input-/outputtekst."""
        # Dit moet configureerbaar zijn op basis van de deploymentcontext
        # en de vereisten voor dataclassificatie
        return True

@dataclass
class ModelLifecycleEvent:
    """Gebeurtenissen in de levenscyclus van het AI-model die geaudit moeten worden."""
    event_type: str  # "training_start", "training_end", "evaluation", "deployment", enz.
    timestamp: float
    model_id: str
    version: str
    actor: str  # Wie de actie initieerde
    details: dict[str, Any]
 
# Belangrijke levenscyclusgebeurtenissen om te auditen:
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:
    """Verifieer de integriteit van een audit trail-hash chain."""
 
    def verify_chain(self, events: list[AuditEvent]) -> dict:
        """
        Verifieer dat de hash chain van de audit trail intact is.
 
        Herberekent de hash van elke gebeurtenis met de hash van de vorige gebeurtenis
        en vergelijkt deze met de opgeslagen 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:
    """Abstracte basis voor opslagbackends van auditlogs."""
 
    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 bestandsgebaseerde auditlog-sink.
 
    Geschikt voor single-node deployments of als lokale buffer
    voordat doorgestuurd wordt naar een gecentraliseerd logsysteem.
    """
 
    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]:
    """
    Bevraag de audit trail om een incidenttijdlijn te reconstrueren.
 
    Retourneert een chronologisch geordende lijst van alle gebeurtenissen
    gerelateerd aan het opgegeven model gedurende het incidentvenster.
    """
    events = sink.query(
        start_time=incident_start,
        end_time=incident_end,
        model_id=model_id,
        limit=10000,
    )
 
    # Sorteer op tijdstempel
    events.sort(key=lambda e: e["timestamp"])
 
    # Annoteer gebeurtenissen met tijddelta's
    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:
    """Genereer een door mensen leesbare samenvatting van een auditgebeurtenis."""
    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}"