From aa66ac2fe8fc0d1c3d4059dd87238f1aa85bc35d Mon Sep 17 00:00:00 2001 From: Matt Adams Date: Tue, 21 Jul 2026 17:18:19 +0100 Subject: [PATCH] feat: control-overlay modifier + campaign-based scenarios MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two features from LinkedIn feedback on the tool: 1. Defensive control overlay (core/controls.py) — a per-page text input where users describe their own security controls; the scenario then assesses the kill chain against them (blocked / detected / missed, and the gaps to close). Mirrors the ai_uplift modifier: an optional prompt fragment threaded through the build_*_messages builders, a control_overlay LangSmith tag, and no change to core/scenario_page.py. Wired into pages 1 & 2 and the threat-group/campaign/custom MCP generate tools. 2. Campaign-based scenarios — build a tabletop around a real, documented ATT&CK campaign (Enterprise/ICS) instead of a generic group. A "Build from" selector on page 1 switches source; campaigns replay the FULL observed chain (no per-phase sampling), like ATLAS case studies. Adds resolve_campaign_kill_chain / list_campaigns (sharing a private helper with the group resolver), a CAMPAIGN template + build_campaign_ messages, scripts/generate_campaigns.py + data/campaigns*.json, and MCP list_campaigns / generate_campaign_scenario tools. Tests: new tests/test_controls.py; campaign resolver + list tests; campaign prompt + controls tests; MCP campaign/controls tests. Full suite green (191). CLAUDE.md updated for the new module, campaign path, and MCP tools. Co-Authored-By: Claude Opus 4.8 (1M context) --- CLAUDE.md | 12 +- core/attack_data.py | 79 +++++- core/controls.py | 105 ++++++++ core/prompts.py | 52 ++++ data/campaigns.json | 226 ++++++++++++++++++ data/campaigns_ics.json | 34 +++ mcp_server.py | 92 ++++++- ...241\357\270\217_Threat_Group_Scenarios.py" | 101 ++++++-- ...7\233\240\357\270\217_Custom_Scenarios.py" | 15 +- scripts/generate_campaigns.py | 67 ++++++ tests/test_attack_data.py | 52 ++++ tests/test_controls.py | 71 ++++++ tests/test_mcp_server.py | 47 +++- tests/test_prompts.py | 52 ++++ 14 files changed, 971 insertions(+), 34 deletions(-) create mode 100644 core/controls.py create mode 100644 data/campaigns.json create mode 100644 data/campaigns_ics.json create mode 100644 scripts/generate_campaigns.py create mode 100644 tests/test_controls.py diff --git a/CLAUDE.md b/CLAUDE.md index 59d97e2..24ff017 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -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 @@ -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 diff --git a/core/attack_data.py b/core/attack_data.py index ae183a2..5b3def5 100644 --- a/core/attack_data.py +++ b/core/attack_data.py @@ -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) @@ -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.""" @@ -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() @@ -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( @@ -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, diff --git a/core/controls.py b/core/controls.py new file mode 100644 index 0000000..adf8ddc --- /dev/null +++ b/core/controls.py @@ -0,0 +1,105 @@ +"""Defensive control overlay for conventional attack scenarios. + +A tabletop is more useful when it is measured against the defences the +organisation actually has. When the user describes their key security controls +(for example "EDR on all endpoints, network segmentation between IT and OT, MFA +on remote access, no egress logging"), this overlay asks the model to assess the +kill chain *against those controls* — calling out, stage by stage, whether each +technique would likely be **blocked, detected, or missed**, and where the gaps +are that the exercise should close. + +Like ``core.ai_uplift`` it is an optional per-page modifier appended to the user +message; it does not change the selected kill chain, only adds a defensive lens. +The pure ``append_controls`` core takes the control description as an explicit +string so headless callers (the MCP server, ``core.prompts``) reuse the exact +framing without touching ``st.session_state``. An empty / whitespace-only +description is a no-op. + +The AI Insider Threat page (page 3) does not use this overlay — it has its own +detection-strategy and NIST CSF control framing built into its prompt. +""" + +from __future__ import annotations + +import streamlit as st + +# Label and help text for the per-page control input. +CONTROLS_LABEL = "🏛️ Your security controls (optional)" +CONTROLS_HELP = ( + "Describe the key security controls in your environment (for example EDR, " + "network segmentation, MFA, email filtering, egress logging gaps). The " + "scenario will assess the attack chain against them — flagging, stage by " + "stage, what would likely be blocked, detected or missed, and which gaps to " + "close. Leave blank to skip." +) +CONTROLS_PLACEHOLDER = ( + "e.g. EDR on all endpoints, network segmentation between corporate and " + "production, MFA on all remote access, SIEM with 90-day retention, no egress " + "logging on the OT segment" +) + +# Appended to the user message when a control description is supplied. Written to +# slot cleanly after the existing "Your task" block in the scenario templates. +# The only ``.format`` field is ``{controls}`` — keep other braces out. +CONTROLS_PROMPT_TEMPLATE = """ +**Defensive control overlay:** +The organisation reports the following existing security controls: +{controls} + +Assess the kill chain above against these stated controls. This does **not** change the techniques in play — it adds a defensive lens. Throughout the scenario: +- For each stage of the attack, judge whether the relevant technique would most likely be **blocked**, **detected**, or **missed** given the controls above, and briefly say why. +- Be realistic and specific: tie each judgement to a named control (or the absence of one). Do not assume controls that were not listed. +- Call out the **coverage gaps** — the stages that would go undetected or unmitigated — as the priorities the exercise should help the team close. +- Where a control would generate a detection, note the **evidence** the response team should expect to find (log source, alert, artefact) so the tabletop can verify whether they would actually have seen it. +""" + + +def render_controls_input(page_id: str, *, label_visibility: str = "visible") -> str: + """Render the security-controls text area for a scenario page. + + The widget is keyed by ``page_id`` so the two scenario pages keep independent + state. Pass ``label_visibility="collapsed"`` when the input already sits inside + an expander titled with ``CONTROLS_LABEL``. Returns the current text value. + """ + return st.text_area( + CONTROLS_LABEL, + key=f"{page_id}_controls", + help=CONTROLS_HELP, + placeholder=CONTROLS_PLACEHOLDER, + label_visibility=label_visibility, + ) + + +def get_controls(page_id: str) -> str: + """Read the control description from session state without rendering. + + Streamlit persists keyed-widget state across reruns, so message assembly can + read the value before the widget itself is instantiated further down the + script. Defaults to ``""`` before the input has ever been rendered. + """ + return str(st.session_state.get(f"{page_id}_controls", "") or "") + + +def append_controls(user_content: str, controls: str) -> str: + """Append the control overlay to ``user_content`` when ``controls`` is non-empty. + + The pure, Streamlit-free core of ``apply_controls`` — takes an explicit + string so headless callers (the MCP server, ``core.prompts``) can reuse the + exact framing. Whitespace-only descriptions are treated as empty. + """ + controls = controls.strip() + if not controls: + return user_content + return user_content + CONTROLS_PROMPT_TEMPLATE.format(controls=controls) + + +def apply_controls(user_content: str, page_id: str) -> str: + """Append the control overlay when a description is present in session state.""" + return append_controls(user_content, get_controls(page_id)) + + +def controls_trace_tags(base_tags: tuple[str, ...], page_id: str) -> tuple[str, ...]: + """Add a ``control_overlay`` LangSmith tag when a control description is set.""" + if get_controls(page_id).strip(): + return base_tags + ("control_overlay",) + return base_tags diff --git a/core/prompts.py b/core/prompts.py index af8963e..49edbdb 100644 --- a/core/prompts.py +++ b/core/prompts.py @@ -19,6 +19,7 @@ from __future__ import annotations from core.ai_uplift import append_ai_uplift +from core.controls import append_controls from data.ai_insider_threats import ( AGENT_CAPABILITIES, CERT_DIMENSIONS, @@ -124,6 +125,7 @@ def build_threat_group_messages( industry: str, company_size: str, ai_uplift: bool = False, + controls: str = "", ) -> list[Message]: """Build the [system, user] messages for a Threat Group / ATLAS scenario.""" template = THREAT_GROUP_ATLAS_TEMPLATE if matrix == "ATLAS" else THREAT_GROUP_ATTACK_TEMPLATE @@ -135,6 +137,54 @@ def build_threat_group_messages( matrix=matrix, ) user_content = append_ai_uplift(user_content, ai_uplift) + user_content = append_controls(user_content, controls) + return [ + {"role": "system", "content": SCENARIO_SYSTEM_PROMPT}, + {"role": "user", "content": user_content}, + ] + + +# --- Campaign (page 1, "Campaign" source) human template --- +# Campaigns are documented, real-world intrusions. They exist only in the +# Enterprise and ICS matrices (not ATLAS), so a single ATT&CK-style template +# suffices. Unlike a threat group — whose techniques are sampled — a campaign +# replays the *full* set of techniques observed in the actual intrusion. + +CAMPAIGN_ATTACK_TEMPLATE = """ +**Background information:** +The company operates in the '{industry}' industry and is of size '{company_size}'. + +**Campaign information:** +This scenario is based on the real-world, documented MITRE ATT&CK campaign '{campaign_name}'. The following techniques were observed during the actual intrusion, drawn from the MITRE ATT&CK {matrix} Matrix: +{kill_chain_string} + +**Your task:** +Create an incident response testing scenario grounded in this documented campaign. The goal of the scenario is to test the company's incident response capabilities against the techniques that were actually used in this real intrusion, focusing on the {matrix} environment. Keep the scenario faithful to how the campaign is known to have operated while adapting it to the company's industry and size. + +Your response should be well structured and formatted using Markdown. Write in British English. +""" + + +def build_campaign_messages( + *, + matrix: str, + campaign_name: str, + kill_chain_string: str, + industry: str, + company_size: str, + ai_uplift: bool = False, + controls: str = "", +) -> list[Message]: + """Build the [system, user] messages for a Campaign scenario (Enterprise/ICS).""" + user_content = CAMPAIGN_ATTACK_TEMPLATE.format( + industry=industry, + company_size=company_size, + campaign_name=campaign_name, + kill_chain_string=kill_chain_string, + matrix=matrix, + ) + user_content = append_ai_uplift(user_content, ai_uplift) + user_content = append_controls(user_content, controls) return [ {"role": "system", "content": SCENARIO_SYSTEM_PROMPT}, {"role": "user", "content": user_content}, @@ -149,6 +199,7 @@ def build_custom_messages( industry: str, company_size: str, ai_uplift: bool = False, + controls: str = "", ) -> list[Message]: """Build the [system, user] messages for a Custom scenario.""" template = CUSTOM_ATLAS_TEMPLATE if matrix == "ATLAS" else CUSTOM_ATTACK_TEMPLATE @@ -160,6 +211,7 @@ def build_custom_messages( matrix=matrix, ) user_content = append_ai_uplift(user_content, ai_uplift) + user_content = append_controls(user_content, controls) return [ {"role": "system", "content": SCENARIO_SYSTEM_PROMPT}, {"role": "user", "content": user_content}, diff --git a/data/campaigns.json b/data/campaigns.json new file mode 100644 index 0000000..59d0677 --- /dev/null +++ b/data/campaigns.json @@ -0,0 +1,226 @@ +[ + { + "group": "2015 Ukraine Electric Power Attack", + "url": "https://attack.mitre.org/campaigns/C0028" + }, + { + "group": "2016 Ukraine Electric Power Attack", + "url": "https://attack.mitre.org/campaigns/C0025" + }, + { + "group": "2022 Ukraine Electric Power Attack", + "url": "https://attack.mitre.org/campaigns/C0034" + }, + { + "group": "2025 Poland Wiper Attacks", + "url": "https://attack.mitre.org/campaigns/C0063" + }, + { + "group": "3CX Supply Chain Attack", + "url": "https://attack.mitre.org/campaigns/C0057" + }, + { + "group": "Anthropic AI-orchestrated Campaign", + "url": "https://attack.mitre.org/campaigns/C0062" + }, + { + "group": "APT28 Nearest Neighbor Campaign", + "url": "https://attack.mitre.org/campaigns/C0051" + }, + { + "group": "APT41 DUST", + "url": "https://attack.mitre.org/campaigns/C0040" + }, + { + "group": "ArcaneDoor", + "url": "https://attack.mitre.org/campaigns/C0046" + }, + { + "group": "C0010", + "url": "https://attack.mitre.org/campaigns/C0010" + }, + { + "group": "C0011", + "url": "https://attack.mitre.org/campaigns/C0011" + }, + { + "group": "C0015", + "url": "https://attack.mitre.org/campaigns/C0015" + }, + { + "group": "C0017", + "url": "https://attack.mitre.org/campaigns/C0017" + }, + { + "group": "C0018", + "url": "https://attack.mitre.org/campaigns/C0018" + }, + { + "group": "C0021", + "url": "https://attack.mitre.org/campaigns/C0021" + }, + { + "group": "C0026", + "url": "https://attack.mitre.org/campaigns/C0026" + }, + { + "group": "C0027", + "url": "https://attack.mitre.org/campaigns/C0027" + }, + { + "group": "C0032", + "url": "https://attack.mitre.org/campaigns/C0032" + }, + { + "group": "C0033", + "url": "https://attack.mitre.org/campaigns/C0033" + }, + { + "group": "CostaRicto", + "url": "https://attack.mitre.org/campaigns/C0004" + }, + { + "group": "Cutting Edge", + "url": "https://attack.mitre.org/campaigns/C0029" + }, + { + "group": "FLORAHOX Activity", + "url": "https://attack.mitre.org/campaigns/C0053" + }, + { + "group": "Frankenstein", + "url": "https://attack.mitre.org/campaigns/C0001" + }, + { + "group": "FrostyGoop Incident", + "url": "https://attack.mitre.org/campaigns/C0041" + }, + { + "group": "FunnyDream", + "url": "https://attack.mitre.org/campaigns/C0007" + }, + { + "group": "HomeLand Justice", + "url": "https://attack.mitre.org/campaigns/C0038" + }, + { + "group": "Indian Critical Infrastructure Intrusions", + "url": "https://attack.mitre.org/campaigns/C0043" + }, + { + "group": "J-magic Campaign", + "url": "https://attack.mitre.org/campaigns/C0050" + }, + { + "group": "Juicy Mix", + "url": "https://attack.mitre.org/campaigns/C0044" + }, + { + "group": "KV Botnet Activity", + "url": "https://attack.mitre.org/campaigns/C0035" + }, + { + "group": "Leviathan Australian Intrusions", + "url": "https://attack.mitre.org/campaigns/C0049" + }, + { + "group": "Night Dragon", + "url": "https://attack.mitre.org/campaigns/C0002" + }, + { + "group": "Operation AkaiRyū", + "url": "https://attack.mitre.org/campaigns/C0060" + }, + { + "group": "Operation CuckooBees", + "url": "https://attack.mitre.org/campaigns/C0012" + }, + { + "group": "Operation Digital Eye", + "url": "https://attack.mitre.org/campaigns/C0061" + }, + { + "group": "Operation Dream Job", + "url": "https://attack.mitre.org/campaigns/C0022" + }, + { + "group": "Operation Dust Storm", + "url": "https://attack.mitre.org/campaigns/C0016" + }, + { + "group": "Operation Ghost", + "url": "https://attack.mitre.org/campaigns/C0023" + }, + { + "group": "Operation Honeybee", + "url": "https://attack.mitre.org/campaigns/C0006" + }, + { + "group": "Operation MidnightEclipse", + "url": "https://attack.mitre.org/campaigns/C0048" + }, + { + "group": "Operation Sharpshooter", + "url": "https://attack.mitre.org/campaigns/C0013" + }, + { + "group": "Operation Spalax", + "url": "https://attack.mitre.org/campaigns/C0005" + }, + { + "group": "Operation Wocao", + "url": "https://attack.mitre.org/campaigns/C0014" + }, + { + "group": "Outer Space", + "url": "https://attack.mitre.org/campaigns/C0042" + }, + { + "group": "Pikabot Distribution February 2024", + "url": "https://attack.mitre.org/campaigns/C0036" + }, + { + "group": "Quad7 Activity", + "url": "https://attack.mitre.org/campaigns/C0055" + }, + { + "group": "RedDelta Modified PlugX Infection Chain Operations", + "url": "https://attack.mitre.org/campaigns/C0047" + }, + { + "group": "RedPenguin", + "url": "https://attack.mitre.org/campaigns/C0056" + }, + { + "group": "Salesforce Data Exfiltration", + "url": "https://attack.mitre.org/campaigns/C0059" + }, + { + "group": "ShadowRay", + "url": "https://attack.mitre.org/campaigns/C0045" + }, + { + "group": "SharePoint ToolShell Exploitation", + "url": "https://attack.mitre.org/campaigns/C0058" + }, + { + "group": "SolarWinds Compromise", + "url": "https://attack.mitre.org/campaigns/C0024" + }, + { + "group": "SPACEHOP Activity", + "url": "https://attack.mitre.org/campaigns/C0052" + }, + { + "group": "Triton Safety Instrumented System Attack", + "url": "https://attack.mitre.org/campaigns/C0030" + }, + { + "group": "Versa Director Zero Day Exploitation", + "url": "https://attack.mitre.org/campaigns/C0039" + }, + { + "group": "Water Curupira Pikabot Distribution", + "url": "https://attack.mitre.org/campaigns/C0037" + } +] diff --git a/data/campaigns_ics.json b/data/campaigns_ics.json new file mode 100644 index 0000000..3179527 --- /dev/null +++ b/data/campaigns_ics.json @@ -0,0 +1,34 @@ +[ + { + "group": "2015 Ukraine Electric Power Attack", + "url": "https://attack.mitre.org/campaigns/C0028" + }, + { + "group": "2016 Ukraine Electric Power Attack", + "url": "https://attack.mitre.org/campaigns/C0025" + }, + { + "group": "2022 Ukraine Electric Power Attack", + "url": "https://attack.mitre.org/campaigns/C0034" + }, + { + "group": "2025 Poland Wiper Attacks", + "url": "https://attack.mitre.org/campaigns/C0063" + }, + { + "group": "FrostyGoop Incident", + "url": "https://attack.mitre.org/campaigns/C0041" + }, + { + "group": "Maroochy Water Breach", + "url": "https://attack.mitre.org/campaigns/C0020" + }, + { + "group": "Triton Safety Instrumented System Attack", + "url": "https://attack.mitre.org/campaigns/C0030" + }, + { + "group": "Unitronics Defacement Campaign", + "url": "https://attack.mitre.org/campaigns/C0031" + } +] diff --git a/mcp_server.py b/mcp_server.py index 9cd1f9d..339ea1f 100644 --- a/mcp_server.py +++ b/mcp_server.py @@ -38,6 +38,7 @@ from core.navigator import build_layer, dumps, parse_technique_id from core.prompts import ( build_ai_insider_messages, + build_campaign_messages, build_custom_messages, build_threat_group_messages, ) @@ -129,6 +130,11 @@ def _generate(config: LLMConfig, messages: list[dict]) -> str: return cleaned +def _maybe_controls(base_tags: tuple[str, ...], controls: str) -> tuple[str, ...]: + """Add a ``control_overlay`` trace tag when a control description is supplied.""" + return base_tags + ("control_overlay",) if controls.strip() else base_tags + + # =========================================================================== # Tier A — data tools (no LLM, no API key) # =========================================================================== @@ -155,6 +161,17 @@ def list_case_studies() -> list[dict]: return list(ad.list_case_studies()) +@mcp.tool() +def list_campaigns(matrix: Matrix = "Enterprise") -> list[dict]: + """List documented MITRE ATT&CK campaigns (Enterprise/ICS only). + + Returns ``[{"group", "url"}]`` — real-world, documented intrusions. Use a + ``group`` value from here with ``generate_campaign_scenario``. Campaigns do + not exist in the ATLAS matrix. + """ + return list(ad.list_campaigns(matrix)) + + @mcp.tool() def get_kill_chain( matrix: Matrix, @@ -162,6 +179,7 @@ def get_kill_chain( seed: int | None = None, industry: str | None = None, company_size: str | None = None, + controls: str = "", ) -> dict: """Resolve a threat group / case study to its kill chain (no LLM call). @@ -169,8 +187,10 @@ def get_kill_chain( ``seed`` for a deterministic draw); ATLAS uses the case study's full documented procedure. When both ``industry`` and ``company_size`` are given, a ready-to-run ``messages`` prompt is included so a client model can generate - the scenario directly. Returns ``techniques``, ``all_techniques``, - ``kill_chain_string`` and (optionally) ``messages``. + the scenario directly; pass ``controls`` (a description of the org's security + controls) to have that prompt assess the chain against them. Returns + ``techniques``, ``all_techniques``, ``kill_chain_string`` and (optionally) + ``messages``. """ kc = _resolve_kill_chain(matrix, group, seed) result = asdict(kc) @@ -181,6 +201,7 @@ def get_kill_chain( kill_chain_string=kc.kill_chain_string, industry=industry, company_size=company_size, + controls=controls, ) else: result["messages"] = None @@ -282,6 +303,7 @@ def generate_threat_group_scenario( api_key: str | None = None, api_base: str | None = None, ai_uplift: bool = False, + controls: str = "", seed: int | None = None, include_detection: bool = False, ) -> str: @@ -289,7 +311,9 @@ def generate_threat_group_scenario( Resolves the kill chain, calls the model (BYO-key), and returns finished Markdown. ``api_key`` omitted → the provider's env var is used. ``ai_uplift`` - adds the AI-enhanced-adversary framing; ``include_detection`` appends the + adds the AI-enhanced-adversary framing; ``controls`` (a description of the + org's security controls) makes the scenario assess the chain against them — + what would be blocked, detected or missed; ``include_detection`` appends the purple-team Detection & Response section. """ kc = _resolve_kill_chain(matrix, group, seed) @@ -302,10 +326,60 @@ def generate_threat_group_scenario( industry=industry, company_size=company_size, ai_uplift=ai_uplift, + controls=controls, ) config = _make_config( provider, model, api_key=api_key, api_base=api_base, - trace_name="Threat Group Scenario (MCP)", trace_tags=("threat_group_scenario", "mcp"), + trace_name="Threat Group Scenario (MCP)", + trace_tags=_maybe_controls(("threat_group_scenario", "mcp"), controls), + ) + scenario = _generate(config, messages) + if include_detection: + detection = _defense_markdown(matrix, [t["ATT&CK ID"] for t in kc.techniques]) + if detection: + scenario = f"{scenario}\n\n---\n\n{detection}" + return scenario + + +@mcp.tool() +def generate_campaign_scenario( + matrix: Literal["Enterprise", "ICS"], + campaign: str, + industry: str, + company_size: str, + provider: str = _DEFAULT_PROVIDER, + model: str = _DEFAULT_MODEL, + api_key: str | None = None, + api_base: str | None = None, + ai_uplift: bool = False, + controls: str = "", + include_detection: bool = False, +) -> str: + """Generate a scenario built around a real, documented ATT&CK campaign. + + Unlike a threat group (whose techniques are sampled), a campaign replays the + *full* set of techniques observed in the actual intrusion. Campaigns exist + only in the Enterprise and ICS matrices. Use a ``campaign`` value from + ``list_campaigns``. Calls the model (BYO-key) and returns finished Markdown. + ``controls`` makes the scenario assess the chain against your stated defences; + ``include_detection`` appends the purple-team Detection & Response section. + """ + kc = ad.resolve_campaign_kill_chain(matrix, campaign) + if not kc.techniques: + raise ValueError(f"No techniques found for campaign '{campaign}' in the {matrix} matrix.") + messages = build_campaign_messages( + matrix=matrix, + campaign_name=campaign, + kill_chain_string=kc.kill_chain_string, + industry=industry, + company_size=company_size, + ai_uplift=ai_uplift, + controls=controls, + ) + config = _make_config( + provider, model, api_key=api_key, api_base=api_base, + trace_name="Campaign Scenario (MCP)", + trace_tags=_maybe_controls(("campaign_scenario", "mcp"), controls), ) scenario = _generate(config, messages) if include_detection: @@ -327,13 +401,15 @@ def generate_custom_scenario( api_base: str | None = None, template_info: str = "", ai_uplift: bool = False, + controls: str = "", include_detection: bool = False, ) -> str: """Generate a custom scenario from a chosen set of ATT&CK / ATLAS techniques. ``techniques`` may be bare IDs or ``"Name (ID)"`` labels. Calls the model - (BYO-key) and returns finished Markdown. ``include_detection`` appends the - purple-team Detection & Response section. + (BYO-key) and returns finished Markdown. ``controls`` (a description of the + org's security controls) makes the scenario assess the chain against them; + ``include_detection`` appends the purple-team Detection & Response section. """ if not techniques: raise ValueError("At least one technique is required.") @@ -345,10 +421,12 @@ def generate_custom_scenario( industry=industry, company_size=company_size, ai_uplift=ai_uplift, + controls=controls, ) config = _make_config( provider, model, api_key=api_key, api_base=api_base, - trace_name="Custom Scenario (MCP)", trace_tags=("custom_scenario", "mcp"), + trace_name="Custom Scenario (MCP)", + trace_tags=_maybe_controls(("custom_scenario", "mcp"), controls), ) scenario = _generate(config, messages) if include_detection: diff --git "a/pages/1_\360\237\233\241\357\270\217_Threat_Group_Scenarios.py" "b/pages/1_\360\237\233\241\357\270\217_Threat_Group_Scenarios.py" index 5bfec4d..a1f5533 100644 --- "a/pages/1_\360\237\233\241\357\270\217_Threat_Group_Scenarios.py" +++ "b/pages/1_\360\237\233\241\357\270\217_Threat_Group_Scenarios.py" @@ -4,10 +4,17 @@ from core.ai_uplift import is_ai_uplift_on, render_ai_uplift_toggle, uplift_trace_tags from core.attack_data import ( load_attack_data, + resolve_campaign_kill_chain, resolve_case_study_kill_chain, resolve_threat_group_kill_chain, ) -from core.prompts import build_threat_group_messages +from core.controls import ( + CONTROLS_LABEL, + controls_trace_tags, + get_controls, + render_controls_input, +) +from core.prompts import build_campaign_messages, build_threat_group_messages from core.detections import ( build_defense_report, is_defense_narrative_on, @@ -49,20 +56,30 @@ def load_groups(matrix): return pd.read_json("./data/atlas-case-studies.json") +@st.cache_resource +def load_campaigns(matrix): + if matrix == "ICS": + return pd.read_json("./data/campaigns_ics.json") + return pd.read_json("./data/campaigns.json") + + # ------------------ Prompt Construction ------------------ # # Prompt text lives in core/prompts.py (shared with the MCP server). This page # only threads its own inputs + the AI-uplift toggle into the shared builder. -def build_messages(matrix, selected_group_alias, kill_chain_string): - return build_threat_group_messages( +def build_messages(matrix, source, selected_group_alias, kill_chain_string): + common = dict( matrix=matrix, - selected_group_alias=selected_group_alias, kill_chain_string=kill_chain_string, industry=industry, company_size=company_size, ai_uplift=is_ai_uplift_on("threat_group"), + controls=get_controls("threat_group"), ) + if source == "Campaign": + return build_campaign_messages(campaign_name=selected_group_alias, **common) + return build_threat_group_messages(selected_group_alias=selected_group_alias, **common) def build_layer_payload(): @@ -124,9 +141,29 @@ def _inline_controls(): st.markdown("# Generate Threat Group Scenario🛡️", unsafe_allow_html=True) matrix = st.session_state.get("matrix", "Enterprise") -groups = load_groups(matrix) +# For Enterprise/ICS the user can build from a threat actor *group* (techniques +# sampled) or a documented *campaign* (full observed chain). ATLAS has neither — +# it always uses a documented case study. if matrix == "ATLAS": + source = "ATLAS" +else: + choice = st.radio( + "Build the scenario from:", + ["Threat actor group", "Campaign"], + horizontal=True, + key="threat_group_source", + help=( + "A threat actor group samples one technique per kill-chain phase. A " + "campaign replays the full set of techniques observed in a real, " + "documented intrusion." + ), + ) + source = "Campaign" if choice == "Campaign" else "Group" + +groups = load_campaigns(matrix) if source == "Campaign" else load_groups(matrix) + +if source == "ATLAS": st.markdown( """ ### Select a Case Study @@ -138,6 +175,18 @@ def _inline_controls(): ) entity_label = "case study" select_placeholder = "Select Case Study" +elif source == "Campaign": + st.markdown( + f""" + ### Select a Campaign + + Use the drop-down selector below to select a documented campaign from the MITRE ATT&CK framework. + + You can then optionally view all of the {matrix} ATT&CK techniques observed in the campaign and/or the campaign's page on the MITRE ATT&CK site. + """ + ) + entity_label = "campaign" + select_placeholder = "Select Campaign" else: st.markdown( f""" @@ -169,22 +218,26 @@ def _inline_controls(): try: if selected_group_alias != select_placeholder: group_url = groups[groups['group'] == selected_group_alias]['url'].values[0] - if matrix == "ATLAS": + if source == "ATLAS": st.markdown(f"[View case study on atlas.mitre.org]({group_url})") + elif source == "Campaign": + st.markdown(f"[View the {selected_group_alias} campaign on attack.mitre.org]({group_url})") else: st.markdown(f"[View {selected_group_alias}'s page on attack.mitre.org]({group_url})") - # Kill-chain resolution (incl. the per-phase sampling for ATT&CK) lives in - # core.attack_data, shared with the MCP server. This page just renders it. - if matrix == "ATLAS": + # Kill-chain resolution (incl. the per-phase sampling for groups, and the + # full documented chain for campaigns) lives in core.attack_data, shared + # with the MCP server. This page just renders it. + if source == "ATLAS": kill_chain = resolve_case_study_kill_chain(selected_group_alias) + elif source == "Campaign": + kill_chain = resolve_campaign_kill_chain(matrix, selected_group_alias) else: kill_chain = resolve_threat_group_kill_chain(matrix, selected_group_alias) if not kill_chain.all_techniques: - entity = "case study" if matrix == "ATLAS" else "threat group" st.warning( - f"There are no {matrix} techniques associated with the {entity}: {selected_group_alias}" + f"There are no {matrix} techniques associated with the {entity_label}: {selected_group_alias}" ) st.stop() @@ -198,20 +251,30 @@ def _inline_controls(): st.dataframe(data=techniques_df, height=200, width='stretch', hide_index=True) kill_chain_string = kill_chain.kill_chain_string - messages = build_messages(matrix, selected_group_alias, kill_chain_string) + messages = build_messages(matrix, source, selected_group_alias, kill_chain_string) except Exception as e: st.error("An error occurred: " + str(e)) st.markdown("") -if matrix == "ATLAS": +if source == "ATLAS": st.markdown( """ ### Generate a Scenario Click the button below to generate a scenario based on the selected case study. The documented attack procedure from the case study will be used to generate the scenario. + It normally takes between 30-50 seconds to generate a scenario, although for local models this is highly dependent on your hardware and the selected model. ⏱️ + """ + ) +elif source == "Campaign": + st.markdown( + """ + ### Generate a Scenario + + Click the button below to generate a scenario based on the selected campaign. The full set of techniques observed in the documented campaign will be used to generate the scenario. + It normally takes between 30-50 seconds to generate a scenario, although for local models this is highly dependent on your hardware and the selected model. ⏱️ """ ) @@ -226,6 +289,10 @@ def _inline_controls(): """ ) +# Optional: describe your own controls so the scenario is measured against them. +with st.expander(CONTROLS_LABEL): + render_controls_input("threat_group", label_visibility="collapsed") + def _ready() -> bool: if model_provider != "Custom" and not st.session_state.get("llm_api_key"): @@ -248,13 +315,17 @@ def _ready() -> bool: return True +_base_tags = ("campaign_scenario",) if source == "Campaign" else ("threat_group_scenario",) + run_scenario_page( page_id="threat_group", build_messages=lambda: messages, is_ready=_ready, download_name=f"AttackGen {selected_group_alias} {matrix}.md", - trace_name="Threat Group Scenario", - trace_tags=uplift_trace_tags(("threat_group_scenario",), page_id="threat_group"), + trace_name="Campaign Scenario" if source == "Campaign" else "Threat Group Scenario", + trace_tags=controls_trace_tags( + uplift_trace_tags(_base_tags, page_id="threat_group"), page_id="threat_group" + ), inline_control=_inline_controls, build_layer=build_layer_payload, build_defense=build_defense_payload, diff --git "a/pages/2_\360\237\233\240\357\270\217_Custom_Scenarios.py" "b/pages/2_\360\237\233\240\357\270\217_Custom_Scenarios.py" index 48d8488..82ddfcf 100644 --- "a/pages/2_\360\237\233\240\357\270\217_Custom_Scenarios.py" +++ "b/pages/2_\360\237\233\240\357\270\217_Custom_Scenarios.py" @@ -3,6 +3,12 @@ from core.ai_uplift import is_ai_uplift_on, render_ai_uplift_toggle, uplift_trace_tags from core.attack_data import list_technique_options, load_attack_data +from core.controls import ( + CONTROLS_LABEL, + controls_trace_tags, + get_controls, + render_controls_input, +) from core.prompts import build_custom_messages from core.detections import ( build_defense_report, @@ -90,6 +96,7 @@ def build_messages(selected_techniques_string, template_info): industry=industry, company_size=company_size, ai_uplift=is_ai_uplift_on("custom"), + controls=get_controls("custom"), ) @@ -245,6 +252,10 @@ def template_selection(template, current_matrix): """ ) +# Optional: describe your own controls so the scenario is measured against them. +with st.expander(CONTROLS_LABEL): + render_controls_input("custom", label_visibility="collapsed") + def _ready() -> bool: if model_provider != "Custom" and not st.session_state.get("llm_api_key"): @@ -273,7 +284,9 @@ def _ready() -> bool: is_ready=_ready, download_name=f"AttackGen Custom {selected_template} {matrix}.md", trace_name="Custom Scenario", - trace_tags=uplift_trace_tags(("custom_scenario",), page_id="custom"), + trace_tags=controls_trace_tags( + uplift_trace_tags(("custom_scenario",), page_id="custom"), page_id="custom" + ), inline_control=_inline_controls, build_layer=build_layer_payload, build_defense=build_defense_payload, diff --git a/scripts/generate_campaigns.py b/scripts/generate_campaigns.py new file mode 100644 index 0000000..3af0073 --- /dev/null +++ b/scripts/generate_campaigns.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python3 +"""Regenerate data/campaigns.json and data/campaigns_ics.json from the bundled STIX. + +Page 1 offers a "Campaign" source that presents a dropdown built from these files +(name -> attack.mitre.org URL) and then looks the campaign's techniques up in the +STIX bundle at runtime. Pre-generating the dropdown data keeps it cheap — the +53 MB Enterprise bundle stays lazily loaded and is only read when a scenario is +actually generated (the same reason data/groups.json exists). + +So whenever the ATT&CK bundles in data/ are refreshed, rerun this to keep the +campaign dropdowns in sync with the shipped version — otherwise campaigns added +in a new ATT&CK release never appear. + +Usage: + python scripts/generate_campaigns.py +""" + +import json +import sys +from pathlib import Path + +from mitreattack.stix20 import MitreAttackData + +DATA_DIR = Path(__file__).parent.parent / "data" + +# (STIX bundle in data/, output mapping file) +SOURCES = [ + ("enterprise-attack.json", "campaigns.json"), + ("ics-attack.json", "campaigns_ics.json"), +] + + +def _attack_url(campaign) -> str | None: + """The campaign's canonical attack.mitre.org URL from its ATT&CK reference.""" + for ref in campaign.get("external_references", []): + if ref.get("source_name") == "mitre-attack": + return ref.get("url") + return None + + +def build_rows(bundle_path: Path) -> list[dict]: + data = MitreAttackData(str(bundle_path)) + rows = [] + for campaign in data.get_campaigns(remove_revoked_deprecated=True): + url = _attack_url(campaign) + if url: + # Reuse the "group" key so the page/loader code stays uniform with + # the threat-group dropdown (data/groups.json). + rows.append({"group": campaign["name"], "url": url}) + # Case-insensitive sort to match the existing dropdown ordering. + rows.sort(key=lambda row: row["group"].lower()) + return rows + + +def main() -> int: + for src, out in SOURCES: + rows = build_rows(DATA_DIR / src) + (DATA_DIR / out).write_text( + json.dumps(rows, indent=4, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + print(f"Wrote {len(rows)} campaigns to data/{out}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_attack_data.py b/tests/test_attack_data.py index d0e54be..13fecee 100644 --- a/tests/test_attack_data.py +++ b/tests/test_attack_data.py @@ -37,6 +37,17 @@ def get_groups_by_alias(self, alias): return [] def get_techniques_used_by_group(self, stix_id): + return self._relationship_entries() + + def get_campaigns_by_alias(self, alias): + if alias == "TESTCAMPAIGN": + return [SimpleNamespace(id="campaign--x")] + return [] + + def get_techniques_used_by_campaign(self, stix_id): + return self._relationship_entries() + + def _relationship_entries(self): return [ {"object": {"id": sid, "name": name, "kill_chain_phases": [{"phase_name": phase}]}} for sid, name, _ext, phase in self._TECHS @@ -123,6 +134,47 @@ def test_sampling_covers_every_phase_present(fake_enterprise, seed): assert len(kc.techniques) == 2 +class TestResolveCampaign: + """Campaigns replay the full documented chain — no per-phase sampling.""" + + def test_unknown_campaign_returns_empty(self, fake_enterprise): + kc = ad.resolve_campaign_kill_chain("Enterprise", "NOPE") + assert kc.techniques == [] and kc.all_techniques == [] and kc.kill_chain_string == "" + + def test_full_chain_no_sampling(self, fake_enterprise): + kc = ad.resolve_campaign_kill_chain("Enterprise", "TESTCAMPAIGN") + # No sampling: the sampled set is the full deduped display set (3 techniques, + # both Initial Access ones kept — unlike the group resolver which samples one). + assert kc.techniques == kc.all_techniques + assert len(kc.techniques) == 3 + phases = [t["Phase Name"] for t in kc.techniques] + # Phases ordered per PHASE_ORDER_ATTACK; hyphenated raw phase normalised. + assert phases == ["Initial Access", "Initial Access", "Command and Control"] + + def test_output_is_json_native(self, fake_enterprise): + kc = ad.resolve_campaign_kill_chain("Enterprise", "TESTCAMPAIGN") + json.dumps(kc.techniques) + for t in kc.techniques: + assert set(t) == {"Technique Name", "ATT&CK ID", "Phase Name"} + assert all(isinstance(v, str) for v in t.values()) + + def test_group_resolver_unaffected_by_shared_helper(self, fake_enterprise): + # The refactor that extracted the shared helper must leave group sampling intact. + kc = ad.resolve_threat_group_kill_chain("Enterprise", "TESTGROUP", seed=0) + assert len(kc.techniques) == 2 # one per phase + assert len(kc.all_techniques) == 3 + + +class TestListCampaigns: + def test_enterprise_shape(self): + campaigns = ad.list_campaigns("Enterprise") + assert campaigns and set(campaigns[0]) == {"group", "url"} + + def test_atlas_raises(self): + with pytest.raises(ValueError): + ad.list_campaigns("ATLAS") + + class TestResolveCaseStudyAtlas: """ATLAS resolution against the real bundled data (deterministic, no sampling).""" diff --git a/tests/test_controls.py b/tests/test_controls.py new file mode 100644 index 0000000..1853dc0 --- /dev/null +++ b/tests/test_controls.py @@ -0,0 +1,71 @@ +"""Tests for `core.controls` — the defensive control overlay. + +The overlay appends a control-assessment fragment to a scenario's prompt and +adds a `control_overlay` LangSmith tag when the user supplies a control +description. It is a no-op when the description is empty or whitespace, and is +namespaced per page so two scenario pages keep independent state. +""" + +from __future__ import annotations + +from core.controls import ( + apply_controls, + append_controls, + controls_trace_tags, + get_controls, +) + + +def test_append_is_noop_when_empty() -> None: + assert append_controls("base content", "") == "base content" + assert append_controls("base content", " \n ") == "base content" + + +def test_append_adds_fragment_and_interpolates_description() -> None: + result = append_controls("base content", " EDR everywhere, no egress logging ") + + assert result != "base content" + assert "Defensive control overlay" in result + # The description is stripped and interpolated verbatim. + assert "EDR everywhere, no egress logging" in result + assert result.startswith("base content") + + +def test_defaults_to_empty_before_input_rendered(fake_session_state) -> None: + assert get_controls("threat_group") == "" + assert apply_controls("base content", "threat_group") == "base content" + assert controls_trace_tags(("threat_group_scenario",), "threat_group") == ( + "threat_group_scenario", + ) + + +def test_apply_appends_from_session_state(fake_session_state) -> None: + fake_session_state["threat_group_controls"] = "MFA on all remote access" + + result = apply_controls("base content", "threat_group") + + assert "Defensive control overlay" in result + assert "MFA on all remote access" in result + + +def test_trace_tags_gains_control_overlay_when_set(fake_session_state) -> None: + fake_session_state["custom_controls"] = "segmentation" + + tags = controls_trace_tags(("custom_scenario",), "custom") + + assert tags == ("custom_scenario", "control_overlay") + + +def test_whitespace_only_description_is_off(fake_session_state) -> None: + fake_session_state["custom_controls"] = " " + + assert apply_controls("x", "custom") == "x" + assert controls_trace_tags(("custom_scenario",), "custom") == ("custom_scenario",) + + +def test_state_is_namespaced_per_page(fake_session_state) -> None: + fake_session_state["threat_group_controls"] = "EDR" + + assert get_controls("threat_group") == "EDR" + assert get_controls("custom") == "" + assert apply_controls("x", "custom") == "x" diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 446c758..1141219 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -40,13 +40,15 @@ def canned_kill_chain(monkeypatch): ) monkeypatch.setattr(s.ad, "resolve_threat_group_kill_chain", lambda *a, **k: kc) monkeypatch.setattr(s.ad, "resolve_case_study_kill_chain", lambda *a, **k: kc) + monkeypatch.setattr(s.ad, "resolve_campaign_kill_chain", lambda *a, **k: kc) return kc EXPECTED_TOOLS = { - "list_threat_groups", "list_case_studies", "get_kill_chain", "get_detection_report", - "get_navigator_layer", "list_ai_insider_options", "get_ai_insider_prompt", - "generate_threat_group_scenario", "generate_custom_scenario", "generate_ai_insider_scenario", + "list_threat_groups", "list_case_studies", "list_campaigns", "get_kill_chain", + "get_detection_report", "get_navigator_layer", "list_ai_insider_options", + "get_ai_insider_prompt", "generate_threat_group_scenario", "generate_campaign_scenario", + "generate_custom_scenario", "generate_ai_insider_scenario", } @@ -118,6 +120,45 @@ def test_generate_custom_scenario(self, no_langsmith, mock_litellm_completion): _args, kwargs = mock_litellm_completion.calls[0] assert "T1566" in kwargs["messages"][1]["content"] + def test_generate_campaign_scenario(self, no_langsmith, canned_kill_chain, mock_litellm_completion): + mock_litellm_completion.content = "# Campaign Scenario" + out = s.generate_campaign_scenario( + "Enterprise", "Operation Wocao", "Finance", "Large", + provider="Anthropic API", model="claude-sonnet-5", api_key="k", + ) + assert out == "# Campaign Scenario" + _args, kwargs = mock_litellm_completion.calls[0] + assert kwargs["model"] == "anthropic/claude-sonnet-5" + # Framed as a documented campaign, with the campaign name interpolated. + user = kwargs["messages"][1]["content"] + assert "Operation Wocao" in user and "real-world" in user + + def test_campaign_empty_kill_chain_raises(self, monkeypatch): + empty = KillChain("Enterprise", "C0027", [], "", []) + monkeypatch.setattr(s.ad, "resolve_campaign_kill_chain", lambda *a, **k: empty) + with pytest.raises(ValueError): + s.generate_campaign_scenario( + "Enterprise", "C0027", "Finance", "Large", + provider="Anthropic API", model="claude-sonnet-5", api_key="k", + ) + + def test_controls_flow_into_prompt_and_tags( + self, no_langsmith, canned_kill_chain, mock_litellm_completion + ): + mock_litellm_completion.content = "scenario" + s.generate_threat_group_scenario( + "Enterprise", "APT29", "Finance", "Large", + provider="Anthropic API", model="claude-sonnet-5", api_key="k", + controls="EDR everywhere, no egress logging", + ) + _args, kwargs = mock_litellm_completion.calls[0] + user = kwargs["messages"][1]["content"] + assert "Defensive control overlay" in user + assert "EDR everywhere, no egress logging" in user + # The control_overlay trace tag is attached only when controls are supplied. + assert s._maybe_controls(("t",), "x") == ("t", "control_overlay") + assert s._maybe_controls(("t",), " ") == ("t",) + def test_bad_provider_raises(self, canned_kill_chain): with pytest.raises(ValueError): s.generate_threat_group_scenario( diff --git a/tests/test_prompts.py b/tests/test_prompts.py index 94bda53..b947580 100644 --- a/tests/test_prompts.py +++ b/tests/test_prompts.py @@ -13,6 +13,7 @@ AI_INSIDER_SYSTEM_PROMPT, SCENARIO_SYSTEM_PROMPT, build_ai_insider_messages, + build_campaign_messages, build_custom_messages, build_threat_group_messages, ) @@ -76,6 +77,57 @@ def test_ai_uplift_appends_only_when_true(self): assert not base.endswith(AI_UPLIFT_PROMPT) assert upl.endswith(AI_UPLIFT_PROMPT) + def test_controls_append_only_when_supplied(self): + base = build_threat_group_messages( + matrix="Enterprise", selected_group_alias="APT29", kill_chain_string="x", + industry="F", company_size="L", + )[1]["content"] + withc = build_threat_group_messages( + matrix="Enterprise", selected_group_alias="APT29", kill_chain_string="x", + industry="F", company_size="L", controls="EDR, no egress logging", + )[1]["content"] + assert "Defensive control overlay" not in base + assert "Defensive control overlay" in withc + assert "EDR, no egress logging" in withc + + +class TestCampaignMessages: + def test_shape_and_content(self): + msgs = build_campaign_messages( + matrix="Enterprise", + campaign_name="Operation Wocao", + kill_chain_string="Initial Access: Phishing (T1566)", + industry="Finance", + company_size="Large", + ) + assert len(msgs) == 2 + assert msgs[0] == {"role": "system", "content": SCENARIO_SYSTEM_PROMPT} + user = msgs[1]["content"] + assert "Operation Wocao" in user + assert "real-world" in user # framed as a documented campaign + assert "Finance" in user and "Large" in user + assert "Enterprise" in user # {matrix} interpolated + assert "Initial Access: Phishing (T1566)" in user + assert BRITISH in user + + def test_ai_uplift_and_controls_append(self): + user = build_campaign_messages( + matrix="ICS", campaign_name="C0028", kill_chain_string="x", + industry="Energy", company_size="Large", + ai_uplift=True, controls="OT segmentation", + )[1]["content"] + assert "OT segmentation" in user # controls description interpolated + assert "AI-enhanced adversary framing" in user # ai_uplift present + assert "Defensive control overlay" in user + + def test_no_modifiers_by_default(self): + user = build_campaign_messages( + matrix="Enterprise", campaign_name="C0027", kill_chain_string="x", + industry="F", company_size="L", + )[1]["content"] + assert "AI-enhanced adversary framing" not in user + assert "Defensive control overlay" not in user + class TestCustomMessages: def test_attack_template_omits_british_english(self):