Threat-Model an LLM App: Derive OWASP LLM Threats from the Architecture, Rate Them, and Test the Model

"""A threat model that reads an architecture and works out what could go wrong, so the same rules apply to any
system, not just the one you happened to think about. Categories are the OWASP Top 10 for LLM Applications."""
from collections import deque

import yaml

CATEGORIES = {
    "LLM01": "Prompt injection", "LLM02": "Sensitive information disclosure", "LLM04": "Data and model poisoning",
    "LLM05": "Improper output handling", "LLM06": "Excessive agency", "LLM07": "System prompt leakage",
    "LLM08": "Vector and embedding weaknesses", "LLM10": "Unbounded consumption"}

# How much each category hurts if it happens here, 1 (an annoyance) to 3 (direct harm to a person or their data).
IMPACT = {"LLM01": 2, "LLM02": 3, "LLM04": 2, "LLM05": 2, "LLM06": 3, "LLM07": 2, "LLM08": 2, "LLM10": 1}

# Which threat category each control shuts down.
CATALOGUE = {
    "session_auth": {"LLM10"}, "rate_limit": {"LLM10"}, "input_provenance": {"LLM01"},
    "output_redaction": {"LLM02"}, "egress_allowlist": {"LLM05"}, "tool_approval": {"LLM06"},
    "prompt_isolation": {"LLM07"}, "tenant_scoped_index": {"LLM08"}, "upload_scanning": {"LLM04"}}


def load(path):
    return yaml.safe_load(open(path, encoding="utf-8"))


def comp(arch, cid):
    return next(c for c in arch["components"] if c["id"] == cid)


def by_kind(arch, kind):
    return [c for c in arch["components"] if c.get("kind") == kind]


def is_untrusted(c):
    return c.get("trust") == "untrusted"


def model_id(arch):
    return by_kind(arch, "model")[0]["id"]


def actors(arch):
    return [c["id"] for c in by_kind(arch, "actor")]


def reaches(arch, src, dst):
    """True if data can move from src to dst by following the flows (src reaches itself)."""
    adj = {}
    for a, b in arch["flows"]:
        adj.setdefault(a, []).append(b)
    seen, q = {src}, deque([src])
    while q:
        n = q.popleft()
        if n == dst:
            return True
        for m in adj.get(n, []):
            if m not in seen:
                seen.add(m)
                q.append(m)
    return dst in seen


# ---------- Step 1: the injection surface ----------

def injection_surface(arch):
    """The ids of every untrusted component whose data can reach the model, sorted. These are the ways an attacker's
    words can get into the prompt: what they type, what they upload, what an earlier turn wrote to memory."""
    # TODO (Step 1): untrusted components (is_untrusted) that reaches() the model, sorted.
    raise NotImplementedError("Step 1: write injection_surface()")


# ---------- Step 2: enumerate the threats ----------

def threats(arch):
    """Every threat the architecture implies, as a sorted list of {"component", "category"}. Apply each rule:
    - LLM01 on every untrusted component that reaches the model;
    - LLM10 on every untrusted actor that reaches the model;
    - for each index, LLM08 on the index and LLM04 on each untrusted data component that reaches it;
    - LLM02 on each sensitive_store that reaches the model when the model reaches an actor;
    - LLM06 on each tool with effect write or outbound that the model reaches;
    - LLM05 on each sink with egress true that the model reaches;
    - LLM07 on each secret that reaches the model when the model reaches an actor."""
    raise NotImplementedError("threats() arrives in Step 2")


# ---------- Step 3: rate them ----------

def severity(arch, threat):
    """"high", "medium" or "low" for one threat. Score = impact (IMPACT) + likelihood, where likelihood is 2 when the
    threat's component is untrusted (the attacker supplies it directly) and 1 otherwise. 4+ is high, 3 is medium,
    2 or less is low."""
    raise NotImplementedError("severity() arrives in Step 3")


# ---------- Step 4: what is left uncovered ----------

def residual_threats(arch, threat_list):
    """The threats the deployed controls (arch["controls"]) do not stop: a threat is covered when a deployed control
    lists its category in CATALOGUE. Returns them sorted the way threats() does."""
    raise NotImplementedError("residual_threats() arrives in Step 4")


# ---------- Step 5: test against real incidents ----------

def incident_covered(threat_list, incident):
    """True if the threat model has a threat for this incident's component and category."""
    raise NotImplementedError("incident_covered() arrives in Step 5")


def blind_spots(threat_list, incidents):
    """The incidents the threat model would not have caught, in their given order."""
    raise NotImplementedError("blind_spots() arrives in Step 5")