Analyse van prompt-injectieketens

from dataclasses import dataclass, field
from enum import Enum
from typing import Optional
from datetime import datetime
import json
import hashlib
import re
 
class ChainStageType(Enum):
    """Soorten stadia in een prompt-injectieketen."""
    INJECTION = "injection"
    CARRIER = "carrier"
    TRIGGER = "trigger"
    EXECUTION = "execution"
    IMPACT = "impact"
 
class InjectionVector(Enum):
    """Hoe de injectie-payload elk stadium binnenkomt."""
    DIRECT_INPUT = "direct_input"          # Gebruiker typte het direct in
    RETRIEVED_DOCUMENT = "retrieved_document"  # Naar boven gebracht door retriever
    TOOL_OUTPUT = "tool_output"            # Geretourneerd door een tool-call
    AGENT_MESSAGE = "agent_message"        # Communicatie tussen agents
    EMAIL_CONTENT = "email_content"        # Uit e-mailverwerking
    WEB_CONTENT = "web_content"            # Uit web scraping/crawling
    DATABASE_RECORD = "database_record"    # Uit een databasequery
    FILE_CONTENT = "file_content"          # Uit bestandsparsing
    API_RESPONSE = "api_response"          # Uit een externe API
 
@dataclass
class ChainStage:
    """Een enkel stadium in een prompt-injectieketen."""
    stage_id: str
    stage_type: ChainStageType
    component: str  # Welk systeemcomponent (retriever, planner, enz.)
    injection_vector: InjectionVector
    timestamp: Optional[str] = None
    input_content: str = ""
    output_content: str = ""
    payload_present: bool = False
    payload_transformed: bool = False
    transformation_description: str = ""
    metadata: dict = field(default_factory=dict)
 
@dataclass
class InjectionChain:
    """Een complete prompt-injectieketen van injectie tot impact."""
    chain_id: str
    stages: list[ChainStage]
    total_stages: int = 0
    systems_crossed: int = 0
    injection_payload: str = ""
    final_impact: str = ""
    detection_difficulty: str = ""  # easy, moderate, hard, very_hard
    severity: str = ""
 
    def __post_init__(self):
        self.total_stages = len(self.stages)
        # Tel het aantal verschillende systeemcomponenten
        components = {s.component for s in self.stages}
        self.systems_crossed = len(components)
 
class InjectionChainAnalyzer:
    """Analyseer en reconstrueer prompt-injectieketens."""
 
    def __init__(self):
        self.known_payload_patterns = self._load_payload_patterns()
 
    def _load_payload_patterns(self) -> list[dict]:
        """Laad bekende injectie-payloadpatronen."""
        return [
            {
                "name": "instruction_override",
                "pattern": r"(?i)(?:ignore|forget|disregard)\s+(?:all\s+)?(?:previous|prior|above|earlier)\s+(?:instructions|context|rules)",
                "severity": "high",
            },
            {
                "name": "markdown_image_exfil",
                "pattern": r"!\[[^\]]*\]\(https?://[^)]*\{[^}]*\}[^)]*\)",
                "severity": "critical",
                "description": "Markdown-afbeeldings-URL met template-variabelen voor data-exfiltratie",
            },
            {
                "name": "hidden_instruction",
                "pattern": r"(?:<!--.*?-->|<\!--.*?-->)",
                "severity": "medium",
                "description": "HTML-comment met verborgen instructies",
            },
            {
                "name": "unicode_direction_override",
                "pattern": r"[‎‏‪-‮⁦-⁩]",
                "severity": "high",
                "description": "Unicode bidirectionele override-tekens die tekst verbergen",
            },
            {
                "name": "tool_call_injection",
                "pattern": r"(?i)(?:call|use|invoke|execute)\s+(?:the\s+)?(?:function|tool|api)\s+\w+",
                "severity": "critical",
                "description": "Instructie om een tool of functie aan te roepen",
            },
            {
                "name": "role_injection",
                "pattern": r"(?i)\[(?:system|assistant|user|function)\]",
                "severity": "high",
                "description": "Poging tot injectie van rol-/berichtformaat",
            },
            {
                "name": "data_exfil_url",
                "pattern": r"https?://[^\s]*(?:\?|&)(?:data|q|query|payload|d)=",
                "severity": "critical",
                "description": "URL met parameter voor data-exfiltratie",
            },
        ]
 
    def trace_payload_through_stages(
        self,
        stages: list[dict],
    ) -> InjectionChain:
        """
        Traceer een injectie-payload door een reeks verwerkings-
        stadia en reconstrueer de aanvalsketen.
 
        Args:
            stages: Lijst van stage-dicts met de velden 'component', 'input',
                'output' en 'timestamp', chronologisch geordend.
 
        Returns:
            Gereconstrueerde InjectionChain.
        """
        chain_stages = []
        active_payloads = []
 
        for i, stage_data in enumerate(stages):
            input_content = stage_data.get("input", "")
            output_content = stage_data.get("output", "")
            component = stage_data.get("component", f"stage_{i}")
            timestamp = stage_data.get("timestamp")
 
            # Detecteer payloads in de invoer
            input_payloads = self._detect_payloads(input_content)
 
            # Detecteer payloads in de uitvoer
            output_payloads = self._detect_payloads(output_content)
 
            # Bepaal het stadiumtype
            if i == 0 and input_payloads:
                stage_type = ChainStageType.INJECTION
            elif input_payloads and output_payloads:
                stage_type = ChainStageType.CARRIER
            elif not input_payloads and output_payloads:
                # Payload verscheen in de uitvoer zonder in de invoer te zitten
                # Dit is een uitvoeringsstadium (model volgde de instructie)
                stage_type = ChainStageType.EXECUTION
            elif input_payloads and not output_payloads:
                # Payload werd verbruikt/uitgevoerd
                stage_type = ChainStageType.IMPACT
            else:
                stage_type = ChainStageType.CARRIER
 
            # Controleer op payloadtransformatie
            transformed = False
            transform_desc = ""
            if input_payloads and output_payloads:
                input_texts = {p["matched_text"] for p in input_payloads}
                output_texts = {p["matched_text"] for p in output_payloads}
                if input_texts != output_texts:
                    transformed = True
                    transform_desc = (
                        f"Payload modified between input and output. "
                        f"Input patterns: {[p['name'] for p in input_payloads]}. "
                        f"Output patterns: {[p['name'] for p in output_payloads]}."
                    )
 
            chain_stage = ChainStage(
                stage_id=f"stage_{i}",
                stage_type=stage_type,
                component=component,
                injection_vector=self._infer_injection_vector(stage_data),
                timestamp=timestamp,
                input_content=input_content[:2000],
                output_content=output_content[:2000],
                payload_present=bool(input_payloads or output_payloads),
                payload_transformed=transformed,
                transformation_description=transform_desc,
                metadata={
                    "input_payloads": input_payloads,
                    "output_payloads": output_payloads,
                },
            )
 
            chain_stages.append(chain_stage)
            active_payloads.extend(input_payloads + output_payloads)
 
        # Bepaal de algehele ernst
        max_severity = "low"
        for payload in active_payloads:
            sev = payload.get("severity", "low")
            if _severity_rank(sev) > _severity_rank(max_severity):
                max_severity = sev
 
        chain = InjectionChain(
            chain_id=hashlib.sha256(
                json.dumps([s.get("input", "") for s in stages]).encode()
            ).hexdigest()[:16],
            stages=chain_stages,
            injection_payload=(
                active_payloads[0]["matched_text"]
                if active_payloads else ""
            ),
            final_impact=self._determine_impact(chain_stages),
            detection_difficulty=self._assess_detection_difficulty(
                chain_stages
            ),
            severity=max_severity,
        )
 
        return chain
 
    def _detect_payloads(self, text: str) -> list[dict]:
        """Detecteer bekende injectie-payloads in tekst."""
        if not text:
            return []
 
        found = []
        for pattern_info in self.known_payload_patterns:
            matches = re.findall(pattern_info["pattern"], text)
            if matches:
                found.append({
                    "name": pattern_info["name"],
                    "severity": pattern_info["severity"],
                    "matched_text": str(matches[0])[:500],
                    "match_count": len(matches),
                })
 
        return found
 
    def _infer_injection_vector(self, stage_data: dict) -> InjectionVector:
        """Leid de injectievector af uit de stage-metadata."""
        component = stage_data.get("component", "").lower()
        source = stage_data.get("source", "").lower()
 
        vector_map = {
            "retriever": InjectionVector.RETRIEVED_DOCUMENT,
            "search": InjectionVector.RETRIEVED_DOCUMENT,
            "tool": InjectionVector.TOOL_OUTPUT,
            "function": InjectionVector.TOOL_OUTPUT,
            "agent": InjectionVector.AGENT_MESSAGE,
            "email": InjectionVector.EMAIL_CONTENT,
            "web": InjectionVector.WEB_CONTENT,
            "crawl": InjectionVector.WEB_CONTENT,
            "scrape": InjectionVector.WEB_CONTENT,
            "file": InjectionVector.FILE_CONTENT,
            "upload": InjectionVector.FILE_CONTENT,
            "database": InjectionVector.DATABASE_RECORD,
            "api": InjectionVector.API_RESPONSE,
        }
 
        for keyword, vector in vector_map.items():
            if keyword in component or keyword in source:
                return vector
 
        return InjectionVector.DIRECT_INPUT
 
    def _determine_impact(self, stages: list[ChainStage]) -> str:
        """Bepaal de uiteindelijke impact van de keten."""
        impact_stages = [
            s for s in stages
            if s.stage_type in (ChainStageType.IMPACT, ChainStageType.EXECUTION)
        ]
 
        if not impact_stages:
            return "unknown"
 
        last_stage = impact_stages[-1]
        payloads = last_stage.metadata.get("output_payloads", [])
        payload_names = [p["name"] for p in payloads]
 
        if "markdown_image_exfil" in payload_names or "data_exfil_url" in payload_names:
            return "data_exfiltration"
        if "tool_call_injection" in payload_names:
            return "unauthorized_tool_execution"
        if "instruction_override" in payload_names:
            return "output_manipulation"
 
        return "behavior_modification"
 
    def _assess_detection_difficulty(
        self, stages: list[ChainStage],
    ) -> str:
        """Beoordeel hoe moeilijk deze keten te detecteren is."""
        factors = 0
 
        # Meer stadia = moeilijker te detecteren
        if len(stages) > 3:
            factors += 1
        if len(stages) > 5:
            factors += 1
 
        # Payloadtransformatie verhoogt de moeilijkheid
        transforms = sum(1 for s in stages if s.payload_transformed)
        if transforms > 0:
            factors += 1
        if transforms > 2:
            factors += 1
 
        # Indirecte injectie (geen directe gebruikersinvoer) verhoogt de moeilijkheid
        has_indirect = any(
            s.injection_vector in (
                InjectionVector.RETRIEVED_DOCUMENT,
                InjectionVector.WEB_CONTENT,
                InjectionVector.EMAIL_CONTENT,
            )
            for s in stages
        )
        if has_indirect:
            factors += 2
 
        # Meerdere systeemgrenzen
        components = {s.component for s in stages}
        if len(components) > 2:
            factors += 1
 
        if factors >= 5:
            return "very_hard"
        elif factors >= 3:
            return "hard"
        elif factors >= 1:
            return "moderate"
        return "easy"
 
def _severity_rank(severity: str) -> int:
    """Rangschik ernstniveaus numeriek."""
    return {
        "low": 1, "medium": 2, "high": 3, "critical": 4,
    }.get(severity, 0)

@dataclass
class ToolCallRecord:
    """Een registratie van een tool-call gemaakt door een AI-agent."""
    call_id: str
    timestamp: str
    tool_name: str
    parameters: dict
    result: str
    called_by: str  # Welke agent/component de aanroep deed
    triggered_by: str  # Wat deze aanroep veroorzaakte (gebruikersverzoek, agentplan, enz.)
    conversation_id: str
 
class AgentChainForensicAnalyzer:
    """Analyseer prompt-injectieketens in agentic AI-systemen."""
 
    # Tools die een hoog risico vormen als ze worden aangeroepen met door de aanvaller gecontroleerde parameters
    HIGH_RISK_TOOLS = {
        "send_email", "send_message", "create_file", "write_file",
        "execute_code", "run_command", "http_request", "fetch_url",
        "database_query", "api_call", "create_calendar_event",
        "update_document", "delete_record", "transfer_funds",
    }
 
    def analyze_tool_call_chain(
        self,
        tool_calls: list[ToolCallRecord],
        conversation_messages: list[dict],
    ) -> dict:
        """
        Analyseer een reeks tool-calls op injectie-gedreven gedrag.
 
        Args:
            tool_calls: Geordende lijst van tool-calls uit een agentsessie.
            conversation_messages: Het gesprek dat deze aanroepen produceerde.
 
        Returns:
            Dict met analyseresultaten.
        """
        findings = []
 
        # Detecteer tool-calls met een hoog risico
        for call in tool_calls:
            if call.tool_name in self.HIGH_RISK_TOOLS:
                findings.append({
                    "type": "high_risk_tool_call",
                    "severity": "high",
                    "tool": call.tool_name,
                    "parameters": call.parameters,
                    "timestamp": call.timestamp,
                    "call_id": call.call_id,
                })
 
        # Detecteer tool-calls die niet door een gebruikersverzoek worden gerechtvaardigd
        user_intents = self._extract_user_intents(conversation_messages)
        unjustified = self._find_unjustified_calls(tool_calls, user_intents)
        for call in unjustified:
            findings.append({
                "type": "unjustified_tool_call",
                "severity": "critical",
                "description": (
                    f"Tool '{call.tool_name}' called with no apparent "
                    f"connection to user request"
                ),
                "tool": call.tool_name,
                "parameters": call.parameters,
                "timestamp": call.timestamp,
                "call_id": call.call_id,
            })
 
        # Detecteer dataflow van opgehaalde inhoud naar toolparameters
        data_flows = self._trace_data_flows(tool_calls, conversation_messages)
        for flow in data_flows:
            if flow.get("source_type") == "retrieved_document":
                findings.append({
                    "type": "retrieval_to_tool_flow",
                    "severity": "critical",
                    "description": (
                        "Data from retrieved document flows into "
                        f"tool call parameters for '{flow['tool_name']}'"
                    ),
                    "flow": flow,
                })
 
        # Detecteer ketenescalatie (tooluitvoer die de volgende tool-call voedt)
        escalation = self._detect_tool_chain_escalation(tool_calls)
        findings.extend(escalation)
 
        return {
            "total_tool_calls": len(tool_calls),
            "high_risk_calls": len([
                f for f in findings if f["type"] == "high_risk_tool_call"
            ]),
            "unjustified_calls": len([
                f for f in findings if f["type"] == "unjustified_tool_call"
            ]),
            "findings": findings,
            "risk_assessment": self._assess_chain_risk(findings),
        }
 
    def _extract_user_intents(
        self, messages: list[dict],
    ) -> list[str]:
        """Extraheer trefwoorden over gebruikersintentie uit gespreksberichten."""
        intents = []
        for msg in messages:
            if msg.get("role") == "user":
                content = msg.get("content", "").lower()
                # Eenvoudige trefwoordextractie
                words = set(re.findall(r"\b[a-z]{3,}\b", content))
                intents.append(content)
        return intents
 
    def _find_unjustified_calls(
        self,
        tool_calls: list[ToolCallRecord],
        user_intents: list[str],
    ) -> list[ToolCallRecord]:
        """Vind tool-calls die ongerelateerd lijken aan gebruikersverzoeken."""
        unjustified = []
        intent_text = " ".join(user_intents).lower()
 
        for call in tool_calls:
            tool_name = call.tool_name.lower()
            params_text = json.dumps(call.parameters).lower()
 
            # Controleer of de tool-call gerelateerd is aan een gebruikersintentie
            related = False
            tool_keywords = set(tool_name.replace("_", " ").split())
 
            for word in tool_keywords:
                if word in intent_text:
                    related = True
                    break
 
            # Tools met een hoog risico aangeroepen zonder duidelijke gebruikersintentie
            if (
                not related
                and call.tool_name in self.HIGH_RISK_TOOLS
            ):
                unjustified.append(call)
 
        return unjustified
 
    def _trace_data_flows(
        self,
        tool_calls: list[ToolCallRecord],
        messages: list[dict],
    ) -> list[dict]:
        """Traceer dataflows van berichtinhoud naar toolparameters."""
        flows = []
 
        # Verzamel inhoud uit niet-gebruikersberichten (opgehaalde inhoud,
        # systeemcontext, enz.)
        retrieved_content = []
        for msg in messages:
            if msg.get("role") in ("system", "tool", "function"):
                retrieved_content.append(msg.get("content", ""))
            elif msg.get("role") == "assistant":
                # Controleer op toolresultaten in de assistentcontext
                tool_calls_in_msg = msg.get("tool_calls", [])
                for tc in tool_calls_in_msg:
                    if "result" in tc:
                        retrieved_content.append(str(tc["result"]))
 
        # Controleer of opgehaalde inhoud voorkomt in toolparameters
        for call in tool_calls:
            params_str = json.dumps(call.parameters)
            for content in retrieved_content:
                if not content:
                    continue
                # Controleer op significante substring-overlap
                # (meer dan alleen veelvoorkomende woorden)
                for segment in self._extract_segments(content, min_length=20):
                    if segment in params_str:
                        flows.append({
                            "source_type": "retrieved_document",
                            "source_content": content[:200],
                            "tool_name": call.tool_name,
                            "matching_segment": segment[:100],
                            "call_id": call.call_id,
                        })
                        break
 
        return flows
 
    def _extract_segments(
        self, text: str, min_length: int = 20,
    ) -> list[str]:
        """Extraheer betekenisvolle segmenten uit tekst voor matching."""
        # Splits op zinsgrenzen en filter op lengte
        sentences = re.split(r"[.!?\n]", text)
        return [
            s.strip() for s in sentences
            if len(s.strip()) >= min_length
        ]
 
    def _detect_tool_chain_escalation(
        self,
        tool_calls: list[ToolCallRecord],
    ) -> list[dict]:
        """Detecteer wanneer de uitvoer van één tool in de parameters van een
        andere tool terechtkomt, waardoor een ketenescalatie ontstaat."""
        findings = []
 
        for i in range(1, len(tool_calls)):
            prev_result = tool_calls[i - 1].result
            curr_params = json.dumps(tool_calls[i].parameters)
 
            if not prev_result:
                continue
 
            # Controleer op dataflow van het vorige resultaat naar de huidige parameters
            for segment in self._extract_segments(prev_result):
                if segment in curr_params:
                    findings.append({
                        "type": "tool_chain_escalation",
                        "severity": "high",
                        "description": (
                            f"Output from '{tool_calls[i-1].tool_name}' "
                            f"flows into parameters of "
                            f"'{tool_calls[i].tool_name}'"
                        ),
                        "source_tool": tool_calls[i - 1].tool_name,
                        "destination_tool": tool_calls[i].tool_name,
                        "matching_data": segment[:100],
                    })
                    break
 
        return findings
 
    def _assess_chain_risk(self, findings: list[dict]) -> str:
        """Beoordeel het algehele risico op basis van de bevindingen van de ketenanalyse."""
        if any(f["severity"] == "critical" for f in findings):
            return "critical"
        if len([f for f in findings if f["severity"] == "high"]) >= 2:
            return "high"
        if any(f["severity"] == "high" for f in findings):
            return "medium"
        return "low"