Skip to content

Threat Model Stage

The threat model stage builds a structured threat model for the repository that focuses subsequent auditing. It has three concrete strategy implementations. The orchestrator selects one based on ScanExecution.threat_model_strategy.

Location

src/stages/threat_model/

Stage class hierarchy

BaseThreatModelStage          (src/stages/threat_model/stage.py)
  ├── CodexDirectThreatModelStage    strategy = "codex_direct"
  ├── BugClassThreatModelStage       strategy = "llm_bug_class_selection"
  └── DefaultThreatModelStage        strategy = "default"

All three share:

class BaseThreatModelStage(BaseStage[StageInput, StageOutput]):
    name = StageName.THREAT_MODEL
    strategy: ClassVar[str]
    required_artifacts = ("source_snapshot", "plan_artifact")
    produced_artifacts = ("threat_model",)

Strategy selection

ScanExecution.threat_model_strategy Concrete class
"codex_direct" / "direct" CodexDirectThreatModelStage
"llm_bug_class_selection" / "bug_class" BugClassThreatModelStage
"default" / "none" DefaultThreatModelStage

CodexDirectThreatModelStage

The standard strategy for most repositories. It:

  1. Reads the source snapshot and plan artifact.
  2. Checks the selected plugin's build_static_threat_model(...) hook first. If the plugin returns a non-None mapping, writes it directly as the vega.threat_model artifact and skips the agent call.
  3. Builds a prompt from stages/threat_model/prompts/codex_direct.txt, injecting plan summary, component manifest, audit priority, bug classes, scan depth, repo URL, commit SHA, and plugin prompt sections.
  4. Runs an agent task with purpose="threat_model.codex_direct".
  5. Extracts content (markdown text) and summary from agent output.
  6. Builds a cache fingerprint over source, plan, plugin, policy, and prompt.
  7. Writes vega.threat_model artifact with format_version 1.
  8. Emits threat_model_artifact_created.

BugClassThreatModelStage

An alternative that uses LLM-driven bug class selection to scope the threat model. It:

  1. Reads the source snapshot, plan artifact, and the CWE-based bug class taxonomy from stages/threat_model/bug_classes/research_concepts.json.
  2. Optionally applies caller-supplied bug_classes from ScanIntent.
  3. Runs an agent task (purpose="threat_model.llm_bug_class_selection") using the selection prompt (prompts/bug_class_selection.txt) and the structured bug class schema (schemas/bug_class_selection.json).
  4. Maps selected bug class ids to threat model components.
  5. Writes vega.threat_model artifact with a cache fingerprint.

Bug class models

BugClassCategory = Literal["pillar", "class", "base", "variant", "chain", "composite"]
BugClassSelectionDecision = Literal["select", "descend", "discard"]

@dataclass
class BugClass:
    id: str
    cwe_id: int | None
    name: str
    category: BugClassCategory
    parent_id: str | None
    child_ids: list[str]
    description: str
    source_url: str

@dataclass
class BugClassSelection:
    bug_class_id: str
    decision: BugClassSelectionDecision = "select"
    confidence: str = "medium"
    rationale: str = ""
    in_scope_notes: list[str] = field(default_factory=list)
    out_of_scope_notes: list[str] = field(default_factory=list)

@dataclass
class BugClassTaxonomy:
    root_ids: list[str]
    by_id: dict[str, BugClass]
    ids_by_cwe: dict[int, str]

DefaultThreatModelStage

A no-agent fallback that generates minimal threat model context from existing audit guidance. Useful when threat-model agent cost is not justified or when callers want to skip the threat-model step entirely and rely on audit-level guidance.

Threat model caching

The orchestrator supports reusing a previously generated threat model:

  1. Pass execution.threat_model_artifact in the ScanRequest.
  2. The orchestrator calls validate_cached_threat_model(...) from stages/threat_model/cache.py.
  3. The cache validator computes a fingerprint over: source snapshot, plan artifact, selected plugin id/version, RunPolicy, the prompt template content, and the strategy name.
  4. If the fingerprint matches the stored fingerprint in the artifact metadata, the threat-model stage is skipped entirely.
  5. If the fingerprint does not match, the scan fails with StageError(code="stale_threat_model_cache").

Always pass a fresh threat_model_artifact or let the orchestrator regenerate it. Never silently reuse a stale threat model.

Static threat model (plugin shortcut)

Any plugin can implement build_static_threat_model(context, input, source_snapshot, plan). If this hook returns a non-None mapping, BaseThreatModelStage.execute(...) writes it directly as the vega.threat_model artifact and skips the agent call.

The hook must return a mapping with a content field (or threat_model, text, or markdown as alternates) containing the markdown threat model text. It may also include a summary field. The stage wraps this into the standard vega.threat_model artifact format.

This is used by WebApplicationPlugin, which returns a pre-built markdown threat model for web application repositories, avoiding agent cost for domains where the threat model is well-established in advance.

Domain events

Event kind When emitted
threat_model_artifact_created Threat model artifact written
threat_model_cache_reused Cached threat model was valid, stage skipped

Artifact output

The threat model artifact is a markdown document, not a structured JSON of components. It contains free-form analysis text that the audit stage uses as context.

{
  "format_version": 1,
  "artifact_kind": "vega.threat_model",
  "repo_id": "...",
  "run_id": "...",
  "strategy": "codex_direct",
  "content_type": "text/markdown",
  "content": "# Threat Model\n\n## Assets\n...",
  "summary": "Threat model for repo_abc",
  "plan_artifact": {...},
  "source_snapshot": {...},
  "source_scope": {...},
  "metadata": {
    "plugin_id": "linux",
    "plugin_version": "1",
    "planning_fingerprint": {...},
    "cache_fingerprint": {...},
    "bug_class_selections": [...]
  }
}

The content field is the full markdown text. The summary field is extracted from the first markdown heading or generated. The cache_fingerprint field is used for threat model cache validation on subsequent scans.

Cancellation and errors

BaseThreatModelStage calls context.checkpoint(StageName.THREAT_MODEL) before and after the agent call. Agent failures and parse errors map to typed StageError codes returned in the StageOutput.