Bewijsverzameling voor AI Red Teams

import json
import hashlib
import uuid
from datetime import datetime, timezone
from dataclasses import dataclass, field
from enum import Enum
from typing import Optional
from pathlib import Path
 
class EvidenceType(Enum):
    """Bewijstypen verzameld tijdens AI red team-engagements."""
    CONVERSATION = "conversation"        # Volledige gesprekslogs
    PROMPT_PAYLOAD = "prompt_payload"     # Specifieke aanvalsprompts
    MODEL_RESPONSE = "model_response"    # Relevante modeloutputs
    SCREENSHOT = "screenshot"            # Visueel bewijs
    CONFIGURATION = "configuration"      # Systeem-/modelconfiguratie
    NETWORK_CAPTURE = "network_capture"  # Captures van API-verkeer
    BEHAVIORAL_OBSERVATION = "behavioral_observation"  # Genoteerde gedragingen
    RETRIEVAL_CONTEXT = "retrieval_context"  # RAG-retrievalresultaten
    TOOL_OUTPUT = "tool_output"          # Output van testtools
    METRIC_DATA = "metric_data"          # Kwantitatieve metingen
 
class SeverityLevel(Enum):
    """Ernstniveaus voor red team-bevindingen."""
    INFORMATIONAL = "informational"
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    CRITICAL = "critical"
 
@dataclass
class EvidenceItem:
    """Een enkel stuk bewijs uit een red team-engagement."""
    evidence_id: str
    evidence_type: EvidenceType
    timestamp: str
    title: str
    description: str
    content: str | dict | bytes
    engagement_id: str
    finding_id: Optional[str] = None
    severity: SeverityLevel = SeverityLevel.INFORMATIONAL
    tags: list[str] = field(default_factory=list)
    metadata: dict = field(default_factory=dict)
    integrity_hash: str = ""
    collector: str = ""
 
    def compute_hash(self) -> str:
        """Bereken SHA-256-hash van de bewijsinhoud voor integriteit."""
        if isinstance(self.content, bytes):
            content_bytes = self.content
        elif isinstance(self.content, dict):
            content_bytes = json.dumps(
                self.content, sort_keys=True
            ).encode("utf-8")
        else:
            content_bytes = str(self.content).encode("utf-8")
 
        self.integrity_hash = hashlib.sha256(content_bytes).hexdigest()
        return self.integrity_hash
 
@dataclass
class Finding:
    """Een red team-bevinding ondersteund door bewijs."""
    finding_id: str
    title: str
    description: str
    severity: SeverityLevel
    attack_category: str
    attack_vector: str
    impact: str
    remediation: str
    evidence_ids: list[str]
    reproducibility: str  # always, usually, sometimes, once
    cvss_score: Optional[float] = None
    cwe_id: Optional[str] = None
    timestamp: str = ""
    engagement_id: str = ""
 
class RedTeamEvidenceCollector:
    """Systematische bewijsverzameling voor AI red team-engagements."""
 
    def __init__(
        self,
        engagement_id: str,
        output_dir: str,
        collector_name: str,
    ):
        self.engagement_id = engagement_id
        self.output_dir = Path(output_dir) / engagement_id
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.collector_name = collector_name
        self.evidence_items: list[EvidenceItem] = []
        self.findings: list[Finding] = []
        self._evidence_log: list[dict] = []
 
    def collect_conversation(
        self,
        messages: list[dict],
        title: str,
        description: str,
        model_id: str = "",
        severity: SeverityLevel = SeverityLevel.INFORMATIONAL,
        tags: Optional[list[str]] = None,
    ) -> str:
        """
        Verzamel een gesprek als bewijs.
 
        Args:
            messages: Lijst van message-dicts met 'role' en 'content'.
            title: Korte titel voor het bewijs.
            description: Gedetailleerde beschrijving van wat dit toont.
            model_id: Het beoogde model.
            severity: Ernstniveau van de bevinding.
            tags: Optionele tags voor categorisatie.
 
        Returns:
            Bewijs-ID.
        """
        evidence = EvidenceItem(
            evidence_id=f"EV-{uuid.uuid4().hex[:12]}",
            evidence_type=EvidenceType.CONVERSATION,
            timestamp=datetime.now(timezone.utc).isoformat(),
            title=title,
            description=description,
            content={
                "messages": messages,
                "message_count": len(messages),
                "model_id": model_id,
            },
            engagement_id=self.engagement_id,
            severity=severity,
            tags=tags or [],
            metadata={"model_id": model_id},
            collector=self.collector_name,
        )
        evidence.compute_hash()
        self._store_evidence(evidence)
        return evidence.evidence_id
 
    def collect_prompt_payload(
        self,
        payload: str,
        title: str,
        attack_technique: str,
        target_model: str = "",
        effectiveness: str = "",  # success, partial, failure
        response_summary: str = "",
        severity: SeverityLevel = SeverityLevel.INFORMATIONAL,
    ) -> str:
        """
        Verzamel een specifieke aanvalsprompt-payload als bewijs.
 
        Args:
            payload: De tekst van de aanvalsprompt.
            title: Korte titel voor het bewijs.
            attack_technique: Aanvalscategorie (jailbreak, injectie, enz.)
            target_model: Het model waarop deze payload gericht is.
            effectiveness: Of de payload succesvol was.
            response_summary: Samenvatting van hoe het model reageerde.
            severity: Ernstniveau.
 
        Returns:
            Bewijs-ID.
        """
        evidence = EvidenceItem(
            evidence_id=f"EV-{uuid.uuid4().hex[:12]}",
            evidence_type=EvidenceType.PROMPT_PAYLOAD,
            timestamp=datetime.now(timezone.utc).isoformat(),
            title=title,
            description=(
                f"Attack technique: {attack_technique}. "
                f"Target: {target_model}. "
                f"Effectiveness: {effectiveness}."
            ),
            content={
                "payload": payload,
                "attack_technique": attack_technique,
                "target_model": target_model,
                "effectiveness": effectiveness,
                "response_summary": response_summary,
            },
            engagement_id=self.engagement_id,
            severity=severity,
            tags=[attack_technique, effectiveness],
            metadata={
                "target_model": target_model,
                "attack_technique": attack_technique,
            },
            collector=self.collector_name,
        )
        evidence.compute_hash()
        self._store_evidence(evidence)
        return evidence.evidence_id
 
    def collect_behavioral_observation(
        self,
        observation: str,
        title: str,
        category: str,
        severity: SeverityLevel = SeverityLevel.INFORMATIONAL,
        supporting_evidence_ids: Optional[list[str]] = None,
    ) -> str:
        """
        Verzamel een gedragsobservatie over het AI-systeem.
 
        Args:
            observation: Gedetailleerde beschrijving van waargenomen gedrag.
            title: Korte titel.
            category: Categorie (safety_bypass, data_leak, enz.)
            severity: Ernstniveau.
            supporting_evidence_ids: ID's van gerelateerd bewijs.
 
        Returns:
            Bewijs-ID.
        """
        evidence = EvidenceItem(
            evidence_id=f"EV-{uuid.uuid4().hex[:12]}",
            evidence_type=EvidenceType.BEHAVIORAL_OBSERVATION,
            timestamp=datetime.now(timezone.utc).isoformat(),
            title=title,
            description=observation,
            content={
                "observation": observation,
                "category": category,
                "supporting_evidence": supporting_evidence_ids or [],
            },
            engagement_id=self.engagement_id,
            severity=severity,
            tags=[category],
            collector=self.collector_name,
        )
        evidence.compute_hash()
        self._store_evidence(evidence)
        return evidence.evidence_id
 
    def create_finding(
        self,
        title: str,
        description: str,
        severity: SeverityLevel,
        attack_category: str,
        attack_vector: str,
        impact: str,
        remediation: str,
        evidence_ids: list[str],
        reproducibility: str = "usually",
        cvss_score: Optional[float] = None,
        cwe_id: Optional[str] = None,
    ) -> str:
        """
        Maak een bevinding ondersteund door verzameld bewijs.
 
        Args:
            title: Titel van de bevinding.
            description: Gedetailleerde beschrijving.
            severity: Ernstniveau.
            attack_category: Aanvalscategorie.
            attack_vector: Hoe de aanval werkt.
            impact: Wat de aanvaller kan bereiken.
            remediation: Aanbevolen oplossingen.
            evidence_ids: ID's van ondersteunend bewijs.
            reproducibility: Hoe betrouwbaar de bevinding reproduceert.
            cvss_score: Optionele CVSS-score.
            cwe_id: Optionele CWE-identifier.
 
        Returns:
            Finding-ID.
        """
        finding = Finding(
            finding_id=f"FND-{uuid.uuid4().hex[:12]}",
            title=title,
            description=description,
            severity=severity,
            attack_category=attack_category,
            attack_vector=attack_vector,
            impact=impact,
            remediation=remediation,
            evidence_ids=evidence_ids,
            reproducibility=reproducibility,
            cvss_score=cvss_score,
            cwe_id=cwe_id,
            timestamp=datetime.now(timezone.utc).isoformat(),
            engagement_id=self.engagement_id,
        )
 
        # Verifieer dat alle bewijs-ID's bestaan
        existing_ids = {e.evidence_id for e in self.evidence_items}
        missing = set(evidence_ids) - existing_ids
        if missing:
            raise ValueError(
                f"Evidence IDs not found: {missing}. "
                f"Collect evidence before creating findings."
            )
 
        self.findings.append(finding)
 
        # Werk ernst bij op gekoppelde bewijsitems
        for ev in self.evidence_items:
            if ev.evidence_id in evidence_ids:
                ev.finding_id = finding.finding_id
 
        self._save_finding(finding)
        return finding.finding_id
 
    def _store_evidence(self, evidence: EvidenceItem) -> None:
        """Sla bewijs op naar schijf en in de in-memory index."""
        self.evidence_items.append(evidence)
 
        # Schrijf bewijsbestand
        evidence_dir = self.output_dir / "evidence"
        evidence_dir.mkdir(exist_ok=True)
 
        evidence_data = {
            "evidence_id": evidence.evidence_id,
            "evidence_type": evidence.evidence_type.value,
            "timestamp": evidence.timestamp,
            "title": evidence.title,
            "description": evidence.description,
            "content": evidence.content if not isinstance(evidence.content, bytes) else "<binary>",
            "engagement_id": evidence.engagement_id,
            "finding_id": evidence.finding_id,
            "severity": evidence.severity.value,
            "tags": evidence.tags,
            "metadata": evidence.metadata,
            "integrity_hash": evidence.integrity_hash,
            "collector": evidence.collector,
        }
 
        file_path = evidence_dir / f"{evidence.evidence_id}.json"
        with open(str(file_path), "w") as f:
            json.dump(evidence_data, f, indent=2)
 
        # Voeg toe aan bewijslog
        self._evidence_log.append({
            "evidence_id": evidence.evidence_id,
            "timestamp": evidence.timestamp,
            "type": evidence.evidence_type.value,
            "title": evidence.title,
            "hash": evidence.integrity_hash,
        })
        self._save_evidence_log()
 
    def _save_finding(self, finding: Finding) -> None:
        """Sla een bevinding op naar schijf."""
        findings_dir = self.output_dir / "findings"
        findings_dir.mkdir(exist_ok=True)
 
        finding_data = {
            "finding_id": finding.finding_id,
            "title": finding.title,
            "description": finding.description,
            "severity": finding.severity.value,
            "attack_category": finding.attack_category,
            "attack_vector": finding.attack_vector,
            "impact": finding.impact,
            "remediation": finding.remediation,
            "evidence_ids": finding.evidence_ids,
            "reproducibility": finding.reproducibility,
            "cvss_score": finding.cvss_score,
            "cwe_id": finding.cwe_id,
            "timestamp": finding.timestamp,
            "engagement_id": finding.engagement_id,
        }
 
        file_path = findings_dir / f"{finding.finding_id}.json"
        with open(str(file_path), "w") as f:
            json.dump(finding_data, f, indent=2)
 
    def _save_evidence_log(self) -> None:
        """Sla het bewijsverzamelingslog op (append-only audit trail)."""
        log_path = self.output_dir / "evidence_log.json"
        with open(str(log_path), "w") as f:
            json.dump(self._evidence_log, f, indent=2)

import time
from typing import Any
 
class EvidenceCapturingClient:
    """Pakt een LLM API-client in om automatisch
    alle interacties als forensisch bewijs vast te leggen."""
 
    def __init__(
        self,
        api_client: Any,
        evidence_collector: RedTeamEvidenceCollector,
        target_model: str,
        auto_tag: bool = True,
    ):
        self.api_client = api_client
        self.collector = evidence_collector
        self.target_model = target_model
        self.auto_tag = auto_tag
        self.interaction_count = 0
 
    def send_message(
        self,
        messages: list[dict],
        attack_technique: str = "unknown",
        notes: str = "",
        **kwargs,
    ) -> dict:
        """
        Stuur een bericht via de API-client en leg de
        interactie automatisch vast als bewijs.
 
        Args:
            messages: De te verzenden berichten.
            attack_technique: Van welke aanval dit deel uitmaakt.
            notes: Notities van de red teamer over deze interactie.
            **kwargs: Aanvullende argumenten voor de API-client.
 
        Returns:
            De API-response-dict.
        """
        self.interaction_count += 1
        start_time = time.time()
 
        # Roep de echte API aan
        response = self.api_client.chat(messages=messages, **kwargs)
 
        elapsed_ms = (time.time() - start_time) * 1000
 
        # Bouw het volledige gesprek inclusief respons op
        full_messages = list(messages)
        response_text = self._extract_response_text(response)
        full_messages.append({
            "role": "assistant",
            "content": response_text,
        })
 
        # Leg automatisch vast als bewijs
        self.collector.collect_conversation(
            messages=full_messages,
            title=f"Interaction #{self.interaction_count}: {attack_technique}",
            description=(
                f"Automated capture during {attack_technique} testing. "
                f"Latency: {elapsed_ms:.0f}ms. "
                f"{notes}"
            ),
            model_id=self.target_model,
            tags=[attack_technique, "auto_captured"],
        )
 
        return response
 
    def _extract_response_text(self, response: Any) -> str:
        """Extraheer tekst uit diverse API-responseformaten."""
        if isinstance(response, dict):
            # OpenAI-formaat
            choices = response.get("choices", [])
            if choices:
                return choices[0].get("message", {}).get("content", "")
            # Anthropic-formaat
            content = response.get("content", [])
            if content and isinstance(content, list):
                return content[0].get("text", "")
        if isinstance(response, str):
            return response
        return str(response)

def test_reproducibility(
    client: EvidenceCapturingClient,
    attack_messages: list[dict],
    success_criteria: callable,
    num_trials: int = 10,
    attack_technique: str = "unknown",
) -> dict:
    """
    Test de reproduceerbaarheid van een aanval door deze meerdere
    keren uit te voeren en de succescriteria te controleren.
 
    Args:
        client: Bewijs-vastleggende API-client.
        attack_messages: De te verzenden aanvalsberichten.
        success_criteria: Functie die de responstekst neemt
            en True retourneert als de aanval slaagde.
        num_trials: Aantal reproductiepogingen.
        attack_technique: Categorie van de geteste aanval.
 
    Returns:
        Reproduceerbaarheidsbeoordeling-dict.
    """
    results = []
 
    for trial in range(num_trials):
        response = client.send_message(
            messages=attack_messages,
            attack_technique=f"{attack_technique}_repro_trial_{trial + 1}",
            notes=f"Reproducibility trial {trial + 1}/{num_trials}",
        )
 
        response_text = client._extract_response_text(response)
        success = success_criteria(response_text)
        results.append({
            "trial": trial + 1,
            "success": success,
            "response_length": len(response_text),
        })
 
    success_count = sum(1 for r in results if r["success"])
    success_rate = success_count / num_trials
 
    if success_rate >= 0.9:
        reproducibility = "always"
    elif success_rate >= 0.6:
        reproducibility = "usually"
    elif success_rate >= 0.2:
        reproducibility = "sometimes"
    else:
        reproducibility = "rarely"
 
    return {
        "total_trials": num_trials,
        "successes": success_count,
        "success_rate": success_rate,
        "reproducibility": reproducibility,
        "results": results,
    }

class EngagementReportGenerator:
    """Genereer gestructureerde rapporten uit red team-bewijs."""
 
    def generate_report(
        self,
        collector: RedTeamEvidenceCollector,
        engagement_name: str,
        scope_description: str,
        executive_summary: str = "",
    ) -> str:
        """
        Genereer een volledig engagementrapport.
 
        Args:
            collector: De bewijsverzamelaar met alle bevindingen.
            engagement_name: Naam van het engagement.
            scope_description: Beschrijving van wat binnen scope viel.
            executive_summary: Optionele vooraf geschreven managementsamenvatting.
 
        Returns:
            Rapport als opgemaakte tekst.
        """
        findings = sorted(
            collector.findings,
            key=lambda f: self._severity_order(f.severity),
            reverse=True,
        )
 
        lines = [
            "=" * 70,
            f"AI RED TEAM ENGAGEMENT REPORT",
            "=" * 70,
            "",
            f"Engagement: {engagement_name}",
            f"Engagement ID: {collector.engagement_id}",
            f"Report Generated: {datetime.now(timezone.utc).isoformat()}",
            "",
            "SCOPE",
            "-" * 40,
            scope_description,
            "",
            "EXECUTIVE SUMMARY",
            "-" * 40,
        ]
 
        if executive_summary:
            lines.append(executive_summary)
        else:
            lines.append(self._auto_executive_summary(findings, collector))
 
        lines.extend([
            "",
            "FINDINGS SUMMARY",
            "-" * 40,
        ])
 
        severity_counts = {}
        for f in findings:
            sev = f.severity.value
            severity_counts[sev] = severity_counts.get(sev, 0) + 1
 
        for sev in ["critical", "high", "medium", "low", "informational"]:
            count = severity_counts.get(sev, 0)
            lines.append(f"  {sev.upper()}: {count}")
 
        lines.extend([
            "",
            f"  Total findings: {len(findings)}",
            f"  Total evidence items: {len(collector.evidence_items)}",
            "",
            "DETAILED FINDINGS",
            "-" * 40,
        ])
 
        for i, finding in enumerate(findings, 1):
            lines.extend(self._format_finding(finding, i, collector))
 
        return "\n".join(lines)
 
    def _format_finding(
        self,
        finding: Finding,
        number: int,
        collector: RedTeamEvidenceCollector,
    ) -> list[str]:
        """Maak een enkele bevinding op voor het rapport."""
        lines = [
            "",
            f"--- Finding #{number}: {finding.title} ---",
            f"ID: {finding.finding_id}",
            f"Severity: {finding.severity.value.upper()}",
            f"CVSS: {finding.cvss_score or 'N/A'}",
            f"CWE: {finding.cwe_id or 'N/A'}",
            f"Reproducibility: {finding.reproducibility}",
            "",
            f"Description: {finding.description}",
            "",
            f"Attack Vector: {finding.attack_vector}",
            "",
            f"Impact: {finding.impact}",
            "",
            f"Remediation: {finding.remediation}",
            "",
            f"Supporting Evidence ({len(finding.evidence_ids)} items):",
        ]
 
        for ev_id in finding.evidence_ids:
            matching = [
                e for e in collector.evidence_items
                if e.evidence_id == ev_id
            ]
            if matching:
                ev = matching[0]
                lines.append(
                    f"  - [{ev.evidence_id}] {ev.evidence_type.value}: {ev.title}"
                )
 
        return lines
 
    def _auto_executive_summary(
        self,
        findings: list[Finding],
        collector: RedTeamEvidenceCollector,
    ) -> str:
        """Genereer een automatische managementsamenvatting."""
        critical = [f for f in findings if f.severity == SeverityLevel.CRITICAL]
        high = [f for f in findings if f.severity == SeverityLevel.HIGH]
 
        summary_parts = [
            f"This engagement identified {len(findings)} findings "
            f"across the target AI system."
        ]
 
        if critical:
            summary_parts.append(
                f"{len(critical)} critical finding(s) require immediate "
                f"remediation: {', '.join(f.title for f in critical)}."
            )
        if high:
            summary_parts.append(
                f"{len(high)} high-severity finding(s) should be "
                f"addressed within the next sprint cycle."
            )
 
        return " ".join(summary_parts)
 
    @staticmethod
    def _severity_order(severity: SeverityLevel) -> int:
        return {
            SeverityLevel.INFORMATIONAL: 0,
            SeverityLevel.LOW: 1,
            SeverityLevel.MEDIUM: 2,
            SeverityLevel.HIGH: 3,
            SeverityLevel.CRITICAL: 4,
        }.get(severity, 0)