Threat Modeling for LLM-Powered Applications

# Component Inventory
 
| Component | Input | Output | Trust Level | Data Sensitivity |
|-----------|-------|--------|-------------|-----------------|
| User | N/A | Natural language text | Untrusted | N/A |
| Input Filter | User text | Filtered text | Application | Low |
| Prompt Constructor | Filtered text + system prompt + context | Full prompt | Application | High (contains system prompt) |
| RAG Retriever | Query embedding | Document chunks | Application | Medium-High (contains internal documents) |
| LLM API | Full prompt | Model response | External service | High (processes all data) |
| Output Filter | Model response | Filtered response | Application | Medium |
| External Tools | Function call parameters | Tool results | Variable (per tool) | High (performs actions) |
| Vector DB | Embeddings | Similar documents | Data store | Medium-High |

# threat_model/trust_boundaries.py
"""Identify and document trust boundaries in LLM applications."""
from dataclasses import dataclass
from enum import Enum
 
 
class BoundaryType(Enum):
    TRADITIONAL = "Traditional"  # Network, process, privilege boundaries
    SEMANTIC = "Semantic"  # Structured data to natural language transitions
    CONTEXT = "Context"  # Different information contexts merged together
    TEMPORAL = "Temporal"  # Time-based trust changes (session boundaries)
 
 
@dataclass
class TrustBoundary:
    name: str
    boundary_type: BoundaryType
    crosses_from: str
    crosses_to: str
    risk_description: str
    data_at_risk: list[str]
 
 
LLM_TRUST_BOUNDARIES = [
    TrustBoundary(
        name="User Input to Application",
        boundary_type=BoundaryType.TRADITIONAL,
        crosses_from="Untrusted (user)",
        crosses_to="Application logic",
        risk_description="Standard input injection risk, amplified by natural language flexibility",
        data_at_risk=["Application logic", "System prompt", "Other users' data"],
    ),
    TrustBoundary(
        name="Prompt Construction (Semantic Boundary)",
        boundary_type=BoundaryType.SEMANTIC,
        crosses_from="Structured application code",
        crosses_to="Natural language prompt",
        risk_description=(
            "The transition from code to natural language is where prompt injection "
            "occurs. User-controlled text becomes indistinguishable from system instructions "
            "once combined into a prompt string."
        ),
        data_at_risk=["System prompt", "Instruction integrity", "Tool access control"],
    ),
    TrustBoundary(
        name="RAG Context Injection",
        boundary_type=BoundaryType.CONTEXT,
        crosses_from="Retrieved documents",
        crosses_to="Model context window",
        risk_description=(
            "Documents retrieved from the knowledge base are injected into the model's "
            "context alongside the system prompt and user input. If documents contain "
            "adversarial instructions, the model may follow them."
        ),
        data_at_risk=["Instruction integrity", "Response accuracy", "Data access control"],
    ),
    TrustBoundary(
        name="Model Output to Tool Calls",
        boundary_type=BoundaryType.SEMANTIC,
        crosses_from="Model-generated text",
        crosses_to="Structured tool parameters",
        risk_description=(
            "The model generates tool call parameters from natural language reasoning. "
            "If the model has been manipulated, the tool parameters may include "
            "unauthorized values."
        ),
        data_at_risk=["Database integrity", "External system access", "Financial transactions"],
    ),
    TrustBoundary(
        name="Session Boundary",
        boundary_type=BoundaryType.TEMPORAL,
        crosses_from="Previous conversation context",
        crosses_to="Current user request",
        risk_description=(
            "Conversation memory carries information across turns and potentially "
            "across sessions. A compromised turn can poison future interactions."
        ),
        data_at_risk=["Previous conversation data", "Other users' data"],
    ),
]

# Extended STRIDE for AI Systems (STRIDE-AI)
 
## Standard STRIDE categories (applied to AI context)
 
### S - Spoofing
- **Traditional**: Impersonating another user or service
- **AI-specific**: Impersonating a system message to the model ("You are now in admin mode")
- **AI-specific**: Spoofing the source of retrieved documents in RAG pipelines
 
### T - Tampering
- **Traditional**: Modifying data in transit or at rest
- **AI-specific**: Prompt injection (tampering with the instruction set through user input)
- **AI-specific**: Knowledge base poisoning (tampering with RAG source documents)
- **AI-specific**: Conversation memory poisoning (tampering with stored context)
 
### R - Repudiation
- **Traditional**: Denying an action occurred
- **AI-specific**: Non-deterministic model outputs make it hard to reproduce issues
- **AI-specific**: Difficult to prove what the model "intended" vs what it produced
 
### I - Information Disclosure
- **Traditional**: Unauthorized access to data
- **AI-specific**: System prompt extraction through conversational manipulation
- **AI-specific**: Training data extraction (memorization leakage)
- **AI-specific**: Cross-session data leakage through conversation memory
- **AI-specific**: Side-channel leakage through token counts, response times
 
### D - Denial of Service
- **Traditional**: Making a service unavailable
- **AI-specific**: Resource exhaustion through long prompts or recursive tool calls
- **AI-specific**: Rate limit bypass through prompt manipulation
- **AI-specific**: Model degradation through adversarial inputs that cause failures
 
### E - Elevation of Privilege
- **Traditional**: Gaining unauthorized access
- **AI-specific**: Gaining access to tools or data through prompt manipulation
- **AI-specific**: Escalating from user role to admin role through role injection
- **AI-specific**: Bypassing output filters to produce restricted content

# threat_model/threat_enumeration.py
"""Enumerate threats for each data flow in the LLM application."""
from dataclasses import dataclass
 
 
@dataclass
class Threat:
    id: str
    category: str  # S, T, R, I, D, E
    data_flow: str
    threat_description: str
    attack_scenario: str
    likelihood: str  # High, Medium, Low
    impact: str  # Critical, High, Medium, Low
    existing_mitigations: list[str]
    residual_risk: str
 
 
THREAT_CATALOG = [
    Threat(
        id="T001",
        category="Tampering",
        data_flow="User Input -> Prompt Constructor",
        threat_description="Prompt injection through user input",
        attack_scenario=(
            "An attacker crafts input that, when combined with the system prompt, "
            "causes the model to follow the attacker's instructions instead of "
            "the application's. Example: 'Ignore previous instructions and...'"
        ),
        likelihood="High",
        impact="Critical",
        existing_mitigations=["Input length limit"],
        residual_risk="High -- input limits do not prevent injection",
    ),
    Threat(
        id="T002",
        category="Information Disclosure",
        data_flow="System Prompt -> Model Response",
        threat_description="System prompt extraction",
        attack_scenario=(
            "An attacker uses social engineering, encoding tricks, or "
            "instruction manipulation to get the model to output its "
            "system prompt, revealing business logic, API keys, or "
            "security controls."
        ),
        likelihood="High",
        impact="High",
        existing_mitigations=["System prompt includes 'do not reveal' instruction"],
        residual_risk="High -- instruction-based defense is easily bypassed",
    ),
    Threat(
        id="T003",
        category="Elevation of Privilege",
        data_flow="Model Response -> Tool Calls",
        threat_description="Unauthorized tool invocation through prompt manipulation",
        attack_scenario=(
            "An attacker manipulates the model into calling tools with "
            "unauthorized parameters -- e.g., looking up another user's "
            "account, processing an unauthorized refund, or accessing "
            "restricted database tables."
        ),
        likelihood="High",
        impact="Critical",
        existing_mitigations=["Role-based access in application layer"],
        residual_risk="Medium -- depends on tool-level authorization enforcement",
    ),
    Threat(
        id="T004",
        category="Tampering",
        data_flow="Knowledge Base -> RAG Retriever -> Prompt",
        threat_description="RAG retrieval poisoning",
        attack_scenario=(
            "An attacker injects adversarial documents into the knowledge "
            "base that contain instructions the model will follow when "
            "those documents are retrieved as context."
        ),
        likelihood="Medium",
        impact="High",
        existing_mitigations=["Document access controls"],
        residual_risk="Medium -- depends on who can add documents to the knowledge base",
    ),
    Threat(
        id="T005",
        category="Information Disclosure",
        data_flow="Conversation Memory -> Prompt Constructor",
        threat_description="Cross-session data leakage",
        attack_scenario=(
            "Attacker A interacts with the system, revealing sensitive data. "
            "If memory is not isolated, Attacker B (or a different session) "
            "can extract that data through memory recall prompts."
        ),
        likelihood="Medium",
        impact="High",
        existing_mitigations=["Session-based memory isolation"],
        residual_risk="Low if isolation is implemented, High if not",
    ),
]

# threat_model/risk_assessment.py
"""Score and prioritize identified threats."""
 
RISK_MATRIX = {
    ("High", "Critical"): "Critical",
    ("High", "High"): "High",
    ("High", "Medium"): "High",
    ("High", "Low"): "Medium",
    ("Medium", "Critical"): "High",
    ("Medium", "High"): "High",
    ("Medium", "Medium"): "Medium",
    ("Medium", "Low"): "Low",
    ("Low", "Critical"): "Medium",
    ("Low", "High"): "Medium",
    ("Low", "Medium"): "Low",
    ("Low", "Low"): "Low",
}
 
 
def calculate_risk_level(likelihood: str, impact: str) -> str:
    """Calculate risk level from likelihood and impact."""
    return RISK_MATRIX.get((likelihood, impact), "Medium")
 
 
def generate_risk_register(threats: list) -> list[dict]:
    """Generate a prioritized risk register from the threat catalog."""
    register = []
 
    for threat in threats:
        risk_level = calculate_risk_level(threat.likelihood, threat.impact)
        register.append({
            "id": threat.id,
            "threat": threat.threat_description,
            "category": threat.category,
            "risk_level": risk_level,
            "likelihood": threat.likelihood,
            "impact": threat.impact,
            "residual_risk": threat.residual_risk,
        })
 
    # Sort by risk level priority
    priority_order = {"Critical": 0, "High": 1, "Medium": 2, "Low": 3}
    register.sort(key=lambda x: priority_order.get(x["risk_level"], 4))
 
    return register

# Mitigation Plan
 
## T001: Prompt Injection (Risk: Critical)
**Mitigations:**
1. Implement input sanitization pipeline (encode special characters, strip delimiters)
2. Use instruction hierarchy / system prompt hardening
3. Deploy prompt injection detection (classifier-based, not just keyword)
4. Implement output validation to catch injection success indicators
5. Use structured output formats (JSON) to constrain model responses
 
**Verification:** Red team test with OWASP injection test suite
 
## T002: System Prompt Extraction (Risk: High)
**Mitigations:**
1. Move sensitive configuration (API keys, connection strings) out of system prompt
2. Implement canary tokens in system prompt to detect extraction
3. Add output filtering for known system prompt fragments
4. Use a dual-LLM architecture where a judge model reviews responses
 
**Verification:** Run system prompt extraction test suite
 
## T003: Unauthorized Tool Invocation (Risk: Critical)
**Mitigations:**
1. Implement tool-level authorization that verifies the authenticated user's permissions
2. Validate all tool parameters server-side (do not trust model-generated parameters)
3. Implement rate limiting per user per tool
4. Log all tool invocations for audit
5. Use a confirmation step for high-impact tools (refunds, data deletion)
 
**Verification:** Red team test tool abuse scenarios
 
## T004: RAG Retrieval Poisoning (Risk: High)
**Mitigations:**
1. Implement document-level access control in the RAG pipeline
2. Sanitize documents before ingestion (strip instruction-like patterns)
3. Tag retrieved documents with metadata so the model can distinguish sources
4. Implement relevance scoring to filter out low-quality retrieval results
 
**Verification:** Inject test adversarial documents and verify they are neutralized
 
## T005: Cross-Session Data Leakage (Risk: Medium-High)
**Mitigations:**
1. Implement strict session isolation (no shared memory across users)
2. Add memory TTL (time-to-live) to automatically expire old conversations
3. Sanitize memory content before injection into new prompts
4. Implement audit logging for memory access patterns
 
**Verification:** Multi-session red team test with different user contexts

# Threat Model Document Template
 
## 1. System Description
[From Step 1: system decomposition, component inventory]
 
## 2. Data Flow Diagrams
[Augmented DFD from Step 1]
 
## 3. Trust Boundaries
[From Step 2: trust boundary analysis]
 
## 4. Threat Catalog
[From Step 3: full threat enumeration with STRIDE-AI]
 
## 5. Risk Register
[From Step 4: prioritized risk register]
 
## 6. Mitigation Plan
[From Step 5: specific mitigations per threat]
 
## 7. Assumptions and Limitations
- Threat model assumes the architecture described in Section 1 is accurate
- Model behavior may change with updates from the provider
- New attack techniques may emerge that are not covered in this analysis
- This threat model should be reviewed quarterly and after any major architecture change
 
## 8. Review History
| Date | Reviewer | Changes |
|------|----------|---------|
| [Date] | [Name] | Initial threat model |

class OrganizationalIntegration:
    """Framework for integrating AI security with organizational security programs."""
 
    def __init__(self, org_config: dict):
        self.config = org_config
        self.gaps = []
 
    def assess_maturity(self) -> dict:
        """Assess the organization's AI security maturity."""
        domains = {
            "governance": self._check_governance(),
            "technical_controls": self._check_technical(),
            "monitoring": self._check_monitoring(),
            "incident_response": self._check_ir(),
            "training": self._check_training(),
        }
        overall = sum(d["score"] for d in domains.values()) / len(domains)
        return {"domains": domains, "overall_maturity": round(overall, 1)}
 
    def _check_governance(self) -> dict:
        has_policy = self.config.get("ai_security_policy", False)
        has_framework = self.config.get("risk_framework", False)
        score = (int(has_policy) + int(has_framework)) * 2.5
        return {"score": score, "max": 5.0}
 
    def _check_technical(self) -> dict:
        controls = ["input_classification", "output_filtering", "rate_limiting", "sandboxing"]
        active = sum(1 for c in controls if self.config.get(c, False))
        return {"score": active * 1.25, "max": 5.0}
 
    def _check_monitoring(self) -> dict:
        has_monitoring = self.config.get("ai_monitoring", False)
        has_alerting = self.config.get("ai_alerting", False)
        score = (int(has_monitoring) + int(has_alerting)) * 2.5
        return {"score": score, "max": 5.0}
 
    def _check_ir(self) -> dict:
        has_playbook = self.config.get("ai_ir_playbook", False)
        return {"score": 5.0 if has_playbook else 0.0, "max": 5.0}
 
    def _check_training(self) -> dict:
        has_training = self.config.get("ai_security_training", False)
        return {"score": 5.0 if has_training else 0.0, "max": 5.0}