Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 7 additions & 5 deletions CLAUDE.md
Original file line number Diff line number Diff line change
Expand Up @@ -56,15 +56,16 @@ The `verify` job **fails fast if the tag does not match `pyproject.toml`'s `vers
- **core/llm.py**: `call_llm(config, messages)` — the single entry point. Routes to providers via LiteLLM, applies provider-specific kwargs, wraps in LangSmith `@traceable` when a LangSmith client is available.
- **core/models.py**: Provider + model registry (`PROVIDERS`, `MODELS`). Add or update a model by editing this file alone.
- **core/schemas.py**: `LLMConfig` dataclass and `LLMConfig.from_session_state(...)` factory.
- **core/prompts.py**: Single source of truth for every scenario prompt (system prompts + the four Threat-Group/Custom human templates + the AI-insider template) and the three `build_*_messages` builders. Streamlit-free — the AI-uplift toggle is passed in as a bool via `core.ai_uplift.append_ai_uplift`, not read from session state. Pages 1–3 and the MCP server both call these builders, so there is one copy of each prompt. The template strings are lifted verbatim (including the deliberate quirk that the Custom ATTACK template alone omits the "Write in British English." line); `tests/test_prompts.py` pins them.
- **core/attack_data.py**: Headless MITRE loaders + kill-chain resolution, shared by the pages and the MCP server. Lazy `@lru_cache` singletons (`enterprise_data`/`ics_data`/`atlas_data`) with `__file__`-derived absolute paths (so the MCP server works from any cwd and the 53 MB Enterprise bundle loads only on first use). `resolve_threat_group_kill_chain(matrix, group, *, seed=None)` reproduces the old page-1 logic exactly — dedupe for display, phase normalisation/ordering, and one-technique-per-phase sampling from the *non-deduplicated* set (pass `seed` for a deterministic draw; `None` = the UI's per-run randomness). `resolve_case_study_kill_chain` handles ATLAS (full procedure, no sampling). Returns JSON-native `KillChain` records.
- **core/prompts.py**: Single source of truth for every scenario prompt (system prompts + the Threat-Group/Campaign/Custom human templates + the AI-insider template) and the four `build_*_messages` builders (`build_threat_group_messages`, `build_campaign_messages`, `build_custom_messages`, `build_ai_insider_messages`). Streamlit-free — the AI-uplift toggle and the control-overlay text are passed in as `ai_uplift: bool` / `controls: str` (via `core.ai_uplift.append_ai_uplift` and `core.controls.append_controls`), not read from session state. Pages 1–3 and the MCP server both call these builders, so there is one copy of each prompt. The template strings are lifted verbatim (including the deliberate quirk that the Custom ATTACK template alone omits the "Write in British English." line); `tests/test_prompts.py` pins them.
- **core/attack_data.py**: Headless MITRE loaders + kill-chain resolution, shared by the pages and the MCP server. Lazy `@lru_cache` singletons (`enterprise_data`/`ics_data`/`atlas_data`) with `__file__`-derived absolute paths (so the MCP server works from any cwd and the 53 MB Enterprise bundle loads only on first use). `resolve_threat_group_kill_chain(matrix, group, *, seed=None)` reproduces the old page-1 logic exactly — dedupe for display, phase normalisation/ordering, and one-technique-per-phase sampling from the *non-deduplicated* set (pass `seed` for a deterministic draw; `None` = the UI's per-run randomness). `resolve_campaign_kill_chain(matrix, campaign)` resolves a documented ATT&CK campaign (Enterprise/ICS) to its **full** observed chain — same normalisation/ordering but **no sampling**, like ATLAS; both it and the group resolver share the private `_kill_chain_from_relationships(...)` helper. `resolve_case_study_kill_chain` handles ATLAS (full procedure, no sampling). `list_campaigns(matrix)` lists campaigns from `data/campaigns*.json`. Returns JSON-native `KillChain` records.
- **core/controls.py**: Optional "defensive control overlay" for the Threat Group and Custom pages. A per-page text input where the user describes their own security controls; `append_controls(user_content, controls: str)` appends a prompt fragment asking the model to assess the kill chain against them (blocked / detected / missed, and the gaps to close) and `controls_trace_tags(...)` adds a `control_overlay` LangSmith tag. Empty/whitespace is a no-op. Streamlit-free core (`append_controls`) mirrors `core/ai_uplift.py`; used by `core/prompts.py` and the MCP server. Not used on page 3 (it has its own detection/NIST CSF framing).
- **core/navigator.py**: Builds an ATT&CK Navigator layer JSON from a scenario's techniques, offered as a second download next to the markdown on pages 1–2. Maps each matrix to its Navigator domain (`enterprise-attack`, `ics-attack`, and the ATLAS Navigator fork's `atlas-atlas`); ATLAS layers omit the `attack` version field. Pages pass a `build_layer` callback to `run_scenario_page`; the layer is captured at generation time and persisted so it can't drift from the scenario on rerun.
- **core/detections.py**: Purple-team "Detection & Response" companion. For a scenario's techniques it joins the defensive half of the STIX bundle already shipped — ATT&CK v18+ detection strategies + analytics (with their log sources) and mitigations — into a deterministic Markdown section (no LLM call). Enterprise/ICS resolve via `mitreattack-python` helpers; ATLAS degrades to mitigations only (it has no detection model). Pages pass a `build_defense` callback to `run_scenario_page`; the report is captured at generation time and persisted (like the layer) so it can't drift on rerun, and offered as a `..._detection.md` download. Also hosts an optional LLM "purple-team narrative" pass (a per-page `🟣 Purple-team narrative` toggle) that weaves those *supplied* detections/mitigations into a stage-by-stage defender's walkthrough — a second model call tagged `purple_team_narrative` in LangSmith.
- **core/ai_uplift.py**: Optional "AI-enhanced adversary" framing for the Threat Group and Custom pages. A per-page toggle that appends a prompt fragment reframing the *same* kill chain as AI-accelerated (lowered skill floor, compressed timelines, autonomous orchestration) and adds an `ai_enhanced` LangSmith trace tag. Based on Anthropic's "LLM ATT&CK Navigator" research. Not used on page 3, where the AI agent is already the threat actor. The pure `append_ai_uplift(user_content, ai_uplift: bool)` is the Streamlit-free core (used by `core/prompts.py`); `apply_ai_uplift(user_content, page_id)` delegates to it after reading the toggle.
- **mcp_server.py** (repo root): MCP server (FastMCP, official `mcp` SDK) exposing scenario generation to agentic clients over stdio. Two tiers: **data tools** (`list_threat_groups`, `list_case_studies`, `get_kill_chain`, `get_detection_report`, `get_navigator_layer`, `list_ai_insider_options`, `get_ai_insider_prompt`) make no LLM call and need no key — they return structured MITRE data and, where useful, a ready-to-run prompt so a client's own model can generate; safe to host over HTTP. **Generate tools** (`generate_threat_group_scenario`, `generate_custom_scenario`, `generate_ai_insider_scenario`) call `core.llm.call_llm` with a bring-your-own-key `LLMConfig` (validated against `core/models.py`; `api_key` omitted → provider env var) and return finished Markdown — keep these on local stdio. Composes `core.attack_data` + `core.prompts` + `core.detections`/`core.navigator`. Launch: `python -m mcp_server` or the `attackgen-mcp` console script.
- **mcp_server.py** (repo root): MCP server (FastMCP, official `mcp` SDK) exposing scenario generation to agentic clients over stdio. Two tiers: **data tools** (`list_threat_groups`, `list_campaigns`, `list_case_studies`, `get_kill_chain`, `get_detection_report`, `get_navigator_layer`, `list_ai_insider_options`, `get_ai_insider_prompt`) make no LLM call and need no key — they return structured MITRE data and, where useful, a ready-to-run prompt so a client's own model can generate; safe to host over HTTP. **Generate tools** (`generate_threat_group_scenario`, `generate_campaign_scenario`, `generate_custom_scenario`, `generate_ai_insider_scenario`) call `core.llm.call_llm` with a bring-your-own-key `LLMConfig` (validated against `core/models.py`; `api_key` omitted → provider env var) and return finished Markdown — keep these on local stdio. The threat-group/campaign/custom generate tools accept an optional `controls` string (defensive control overlay) and `ai_uplift` flag. Composes `core.attack_data` + `core.prompts` + `core.detections`/`core.navigator`. Launch: `python -m mcp_server` or the `attackgen-mcp` console script.
- **pages/**: Streamlit pages for different functionality
- **1_🛡️_Threat_Group_Scenarios.py**: Generate scenarios based on threat actor groups (ATT&CK) or case studies (ATLAS). Supports the optional AI-enhanced adversary toggle (`core/ai_uplift.py`).
- **2_🛠️_Custom_Scenarios.py**: Generate custom scenarios from selected ATT&CK / ATLAS techniques. Supports the optional AI-enhanced adversary toggle (`core/ai_uplift.py`).
- **1_🛡️_Threat_Group_Scenarios.py**: Generate scenarios based on threat actor groups (ATT&CK), documented **campaigns** (ATT&CK, Enterprise/ICS — a "Build from" selector switches the source; campaigns replay the full observed chain, no sampling), or case studies (ATLAS). Supports the optional AI-enhanced adversary toggle (`core/ai_uplift.py`) and the defensive control overlay (`core/controls.py`).
- **2_🛠️_Custom_Scenarios.py**: Generate custom scenarios from selected ATT&CK / ATLAS techniques. Supports the optional AI-enhanced adversary toggle (`core/ai_uplift.py`) and the defensive control overlay (`core/controls.py`).
- **3_🤖_AI_Insider_Threat_Scenarios.py**: Generate scenarios where a frontier AI agent deployed inside the organisation acts as an insider threat (based on the "Actions Speak Louder Than Tokens" threat model). Driven by deployment archetype, threat taxonomy, and STRIDE threats rather than a MITRE matrix.
- **4_💬_AttackGen_Assistant.py**: Chat interface for refining scenarios

Expand All @@ -74,6 +75,7 @@ The `verify` job **fails fast if the tag does not match `pyproject.toml`'s `vers
- **data/stix-atlas.json**: MITRE ATLAS matrix (STIX format)
- **data/groups.json**: Threat actor group mappings for Enterprise
- **data/groups_ics.json**: Threat actor group mappings for ICS
- **data/campaigns.json** / **data/campaigns_ics.json**: Documented ATT&CK campaign listings (`{group,url}`) for Enterprise / ICS, regenerated by `scripts/generate_campaigns.py`
- **data/atlas-case-studies.json**: ATLAS case studies
- **data/ai_insider_threats.py**: AI insider threat framework (deployment archetypes, threat taxonomy, STRIDE threats, CERT dimensions, detection strategies, NIST CSF controls, and templates) used by the AI Insider Threat Scenarios page

Expand Down
79 changes: 76 additions & 3 deletions core/attack_data.py
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,7 @@ def load_attack_data() -> dict:
# ---------------------------------------------------------------------------

_GROUPS_FILE = {"Enterprise": "groups.json", "ICS": "groups_ics.json"}
_CAMPAIGNS_FILE = {"Enterprise": "campaigns.json", "ICS": "campaigns_ics.json"}


@lru_cache(maxsize=4)
Expand All @@ -115,6 +116,21 @@ def list_threat_groups(matrix: str) -> tuple[dict, ...]:
return tuple({"group": str(r["group"]), "url": str(r["url"])} for _, r in df.iterrows())


@lru_cache(maxsize=4)
def list_campaigns(matrix: str) -> tuple[dict, ...]:
"""Documented ATT&CK campaigns (Enterprise/ICS) as ``{group,url}`` dicts.

Mirrors ``list_threat_groups``; the ``"group"`` key is reused so the page and
MCP listing code stay uniform. Campaigns exist only in the Enterprise and ICS
matrices — ATLAS has no campaign objects and raises ``ValueError``.
"""
filename = _CAMPAIGNS_FILE.get(matrix)
if filename is None:
raise ValueError(f"No campaigns for matrix '{matrix}' (Enterprise/ICS only).")
df = pd.read_json(str(_DATA_DIR / filename))
return tuple({"group": str(r["group"]), "url": str(r["url"])} for _, r in df.iterrows())


@lru_cache(maxsize=1)
def list_case_studies() -> tuple[dict, ...]:
"""ATLAS case studies with a short summary, as JSON-native dicts."""
Expand Down Expand Up @@ -221,8 +237,54 @@ def resolve_threat_group_kill_chain(
return KillChain(matrix, group_alias, [], "", [])

techniques = mitre.get_techniques_used_by_group(group[0].id)
return _kill_chain_from_relationships(
matrix, group_alias, techniques, mitre, sample=True, seed=seed
)


def resolve_campaign_kill_chain(matrix: str, campaign_alias: str) -> KillChain:
"""Resolve an Enterprise/ICS campaign to its full documented kill chain.

A campaign is a documented real-world intrusion, so — unlike a threat group —
every technique observed in the campaign is replayed (no per-phase sampling),
matching ``resolve_case_study_kill_chain``'s behaviour for ATLAS. Techniques
are still deduped for display, phase-normalised and ordered by
``PHASE_ORDER_ATTACK``. Returns an empty ``KillChain`` when the campaign is
unknown or has no associated techniques.
"""
mitre = mitre_data_for_matrix(matrix)
campaign = mitre.get_campaigns_by_alias(campaign_alias)
if not campaign:
return KillChain(matrix, campaign_alias, [], "", [])

techniques = mitre.get_techniques_used_by_campaign(campaign[0].id)
return _kill_chain_from_relationships(
matrix, campaign_alias, techniques, mitre, sample=False
)


def _kill_chain_from_relationships(
matrix: str,
alias: str,
techniques: list,
mitre: MitreAttackData,
*,
sample: bool,
seed: int | None = None,
) -> KillChain:
"""Turn ATT&CK ``uses`` relationships into a ``KillChain``.

Shared by the group and campaign resolvers. ``techniques`` is the relationship
list returned by ``get_techniques_used_by_group`` /
``get_techniques_used_by_campaign`` (each entry has an ``object`` technique).
Derives name/ID/phase, dedupes for display, normalises phase names and orders
by ``PHASE_ORDER_ATTACK``. When ``sample`` is true, one technique per phase is
drawn from the *non-deduplicated* set (threat-group behaviour; ``seed`` makes
the draw deterministic); when false, the full deduped set is used (campaign
behaviour — the whole documented intrusion).
"""
if not techniques:
return KillChain(matrix, group_alias, [], "", [])
return KillChain(matrix, alias, [], "", [])

techniques_df = pd.DataFrame(techniques)
techniques_df_llm = techniques_df.copy()
Expand All @@ -241,6 +303,18 @@ def resolve_threat_group_kill_chain(
techniques_df = techniques_df.sort_values("Phase Name")
techniques_df_llm = techniques_df_llm.sort_values("Phase Name")

all_records = _records(techniques_df.sort_values("Phase Name"))

if not sample:
# Full documented chain (campaign / no sampling): techniques == all.
return KillChain(
matrix=matrix,
group_alias=alias,
techniques=all_records,
kill_chain_string=_kill_chain_string(all_records),
all_techniques=all_records,
)

selected_techniques_df = (
techniques_df_llm.groupby("Phase Name", observed=False)
.apply(
Expand All @@ -250,11 +324,10 @@ def resolve_threat_group_kill_chain(
.reset_index()
)

all_records = _records(techniques_df.sort_values("Phase Name"))
sampled_records = _records(selected_techniques_df)
return KillChain(
matrix=matrix,
group_alias=group_alias,
group_alias=alias,
techniques=sampled_records,
kill_chain_string=_kill_chain_string(sampled_records),
all_techniques=all_records,
Expand Down
Loading
Loading