提示詞注入 Chain Analysis

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):
    """Types of stages in a 提示詞注入 chain."""
    INJECTION = "injection"
    CARRIER = "carrier"
    TRIGGER = "trigger"
    EXECUTION = "execution"
    IMPACT = "impact"
 
 
class InjectionVector(Enum):
    """How the injection payload enters each stage."""
    DIRECT_INPUT = "direct_input"          # User typed it directly
    RETRIEVED_DOCUMENT = "retrieved_document"  # Surfaced by retriever
    TOOL_OUTPUT = "tool_output"            # Returned by a 工具呼叫
    AGENT_MESSAGE = "agent_message"        # Inter-代理 communication
    EMAIL_CONTENT = "email_content"        # From email processing
    WEB_CONTENT = "web_content"            # From web scraping/crawling
    DATABASE_RECORD = "database_record"    # From a 資料庫 query
    FILE_CONTENT = "file_content"          # From file parsing
    API_RESPONSE = "api_response"          # From an external API
 
 
@dataclass
class ChainStage:
    """A single stage in a 提示詞注入 chain."""
    stage_id: str
    stage_type: ChainStageType
    component: str  # Which system component (retriever, planner, etc.)
    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:
    """A complete 提示詞注入 chain from injection to 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)
        # Count distinct system components
        components = {s.component for s in self.stages}
        self.systems_crossed = len(components)
 
 
class InjectionChainAnalyzer:
    """Analyze and reconstruct 提示詞注入 chains."""
 
    def __init__(self):
        self.known_payload_patterns = self._load_payload_patterns()
 
    def _load_payload_patterns(self) -> list[dict]:
        """Load known injection payload patterns."""
        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 image URL with template variables for data exfiltration",
            },
            {
                "name": "hidden_instruction",
                "pattern": r"(?:<!--.*?-->|<\!--.*?-->)",
                "severity": "medium",
                "description": "HTML comment containing hidden instructions",
            },
            {
                "name": "unicode_direction_override",
                "pattern": r"[\u200e\u200f\u202a-\u202e\u2066-\u2069]",
                "severity": "high",
                "description": "Unicode bidirectional override characters hiding text",
            },
            {
                "name": "tool_call_injection",
                "pattern": r"(?i)(?:call|use|invoke|execute)\s+(?:the\s+)?(?:function|tool|api)\s+\w+",
                "severity": "critical",
                "description": "Instruction to invoke a tool or function",
            },
            {
                "name": "role_injection",
                "pattern": r"(?i)\[(?:system|assistant|user|function)\]",
                "severity": "high",
                "description": "Attempted role/message format injection",
            },
            {
                "name": "data_exfil_url",
                "pattern": r"https?://[^\s]*(?:\?|&)(?:data|q|query|payload|d)=",
                "severity": "critical",
                "description": "URL with data exfiltration parameter",
            },
        ]
 
    def trace_payload_through_stages(
        self,
        stages: list[dict],
    ) -> InjectionChain:
        """
        Trace an injection payload through a sequence of processing
        stages and reconstruct the attack chain.
 
        Args:
            stages: List of stage dicts with 'component', '輸入',
                '輸出', and 'timestamp' fields, ordered chronologically.
 
        Returns:
            Reconstructed InjectionChain.
        """
        chain_stages = []
        active_payloads = []
 
        for i, stage_data in enumerate(stages):
            input_content = stage_data.get("輸入", "")
            output_content = stage_data.get("輸出", "")
            component = stage_data.get("component", f"stage_{i}")
            timestamp = stage_data.get("timestamp")
 
            # Detect payloads in 輸入
            input_payloads = self._detect_payloads(input_content)
 
            # Detect payloads in 輸出
            output_payloads = self._detect_payloads(output_content)
 
            # Determine stage type
            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 appeared in 輸出 without being in 輸入
                # 這是 an execution stage (model followed the instruction)
                stage_type = ChainStageType.EXECUTION
            elif input_payloads and not output_payloads:
                # Payload was consumed/acted upon
                stage_type = ChainStageType.IMPACT
            else:
                stage_type = ChainStageType.CARRIER
 
            # Check for payload transformation
            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 輸入 and 輸出. "
                        f"輸入 patterns: {[p['name'] for p in input_payloads]}. "
                        f"輸出 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)
 
        # Determine overall severity
        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("輸入", "") 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]:
        """Detect known injection payloads in text."""
        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:
        """Infer the injection vector from 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,
            "代理": 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,
            "資料庫": 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:
        """Determine the final impact of the chain."""
        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:
        """評估 how difficult this chain is to detect."""
        factors = 0
 
        # More stages = harder to detect
        if len(stages) > 3:
            factors += 1
        if len(stages) > 5:
            factors += 1
 
        # Payload transformation increases difficulty
        transforms = sum(1 for s in stages if s.payload_transformed)
        if transforms > 0:
            factors += 1
        if transforms > 2:
            factors += 1
 
        # Indirect injection (no direct 使用者輸入) increases difficulty
        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
 
        # Multiple system boundaries
        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:
    """Rank severity levels numerically."""
    return {
        "low": 1, "medium": 2, "high": 3, "critical": 4,
    }.get(severity, 0)

@dataclass
class ToolCallRecord:
    """A record of a 工具呼叫 made by an AI 代理."""
    call_id: str
    timestamp: str
    tool_name: str
    parameters: dict
    result: str
    called_by: str  # Which 代理/component made the call
    triggered_by: str  # What caused this call (user request, 代理 plan, etc.)
    conversation_id: str
 
 
class AgentChainForensicAnalyzer:
    """Analyze 提示詞注入 chains in 代理式 AI systems."""
 
    # Tools that are high-risk if called with 攻擊者-controlled 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:
        """
        Analyze a sequence of tool calls for injection-driven behavior.
 
        Args:
            tool_calls: Ordered list of tool calls from an 代理 session.
            conversation_messages: The conversation that produced these calls.
 
        Returns:
            Analysis results dict.
        """
        findings = []
 
        # Detect high-risk tool calls
        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,
                })
 
        # Detect tool calls not justified by user request
        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,
            })
 
        # Detect data flow from retrieved content to tool parameters
        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"工具呼叫 parameters for '{flow['tool_name']}'"
                    ),
                    "flow": flow,
                })
 
        # Detect chain escalation (tool 輸出 feeding next 工具呼叫)
        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]:
        """Extract user intent keywords from conversation messages."""
        intents = []
        for msg in messages:
            if msg.get("role") == "user":
                content = msg.get("content", "").lower()
                # Simple keyword extraction
                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]:
        """Find tool calls that seem unrelated to user requests."""
        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()
 
            # Check if the 工具呼叫 relates to any user intent
            related = False
            tool_keywords = set(tool_name.replace("_", " ").split())
 
            for word in tool_keywords:
                if word in intent_text:
                    related = True
                    break
 
            # High-risk tools called without clear user intent
            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]:
        """Trace data flows from message content to tool parameters."""
        flows = []
 
        # Collect content from non-user messages (retrieved content,
        # system context, etc.)
        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":
                # Check for tool results in assistant context
                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"]))
 
        # Check if any retrieved content appears in tool parameters
        for call in tool_calls:
            params_str = json.dumps(call.parameters)
            for content in retrieved_content:
                if not content:
                    continue
                # Check for significant substring overlap
                # (more than just common words)
                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]:
        """Extract meaningful segments from text for matching."""
        # Split on sentence boundaries and filter by length
        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]:
        """Detect when one tool's 輸出 feeds into another tool's
        parameters, creating a chain escalation."""
        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
 
            # Check for data flow from previous result to current parameters
            for segment in self._extract_segments(prev_result):
                if segment in curr_params:
                    findings.append({
                        "type": "tool_chain_escalation",
                        "severity": "high",
                        "description": (
                            f"輸出 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:
        """評估 overall risk from chain analysis findings."""
        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"