diff --git a/.claude/skills/secopstm-threat-modeling/SKILL.md b/.claude/skills/secopstm-threat-modeling/SKILL.md index e1de3d1..df0b93f 100644 --- a/.claude/skills/secopstm-threat-modeling/SKILL.md +++ b/.claude/skills/secopstm-threat-modeling/SKILL.md @@ -25,7 +25,7 @@ description: Perform comprehensive threat modeling for target systems using SecO ``` {system-name}/ -├── model.md # Threat model +├── model.md # System model ├── context/ │ └── {system-name}_context.yaml ├── BOM/ @@ -129,7 +129,7 @@ Output modeling strategy for user confirmation: system name, complexity (high/me - "Drone 350" → `drone-350` **Filename Consistency**: -- model.md title: `# Threat Model: DJI Power Inspection` +- model.md title: `# System Model: DJI Power Inspection` - Context file: `context/dji-power-inspection_context.yaml` - Ensure consistent use of hyphens/underscores diff --git a/.claude/skills/secopstm-threat-modeling/references/context-format.md b/.claude/skills/secopstm-threat-modeling/references/context-format.md index adeac86..7e1c054 100644 --- a/.claude/skills/secopstm-threat-modeling/references/context-format.md +++ b/.claude/skills/secopstm-threat-modeling/references/context-format.md @@ -2,7 +2,7 @@ ## Overview -The GDAF (Generic Threat Modeling and Attack Framework) context file defines the attack context for threat modeling. It includes attack objectives, threat actors, and risk criteria. +The GDAF (Goal-Driven Attack Flow) context file defines the attack context used for attack path analysis. It includes attack objectives, threat actors, and risk criteria. ## File Location diff --git a/.claude/skills/secopstm-threat-modeling/references/dsl-syntax.md b/.claude/skills/secopstm-threat-modeling/references/dsl-syntax.md index 45bd5cc..d295b1a 100644 --- a/.claude/skills/secopstm-threat-modeling/references/dsl-syntax.md +++ b/.claude/skills/secopstm-threat-modeling/references/dsl-syntax.md @@ -18,7 +18,7 @@ ## Overview -SecOpsTM threat model is a Markdown file parsed in three passes: +A SecOpsTM system model is a Markdown file parsed in three passes: 1. **Pass 0** — `## Context` (model-level configuration) 2. **First pass** — `## Boundaries`, `## Actors`, `## Servers`, `## Data` (element definitions) @@ -40,14 +40,16 @@ SecOpsTM threat model is a Markdown file parsed in three passes: **⚠️ CRITICAL**: model.md must use the following title format: ```markdown -# Threat Model: {System Name} +# System Model: {System Name} ``` **Format requirements**: -- Must start with `# Threat Model: ` (English colon `:` + space) +- Must start with `# System Model: ` (English colon `:` + space) - System name follows after the space -- **Correct**: `# Threat Model: DJI Power Inspection Drone System` -- **Incorrect**: `# DJI Power Inspection Drone System - Threat Model` +- **Correct**: `# System Model: DJI Power Inspection Drone System` +- **Incorrect**: `# DJI Power Inspection Drone System - System Model` +- The legacy `# Threat Model: {System Name}` form is still read correctly (with a + deprecation warning) — always write new models with `# System Model:`. **Impact**: Incorrect title format causes SecOpsTM to fail generating report filenames correctly. @@ -163,7 +165,7 @@ Actors represent people, external systems, or roles interacting with the system, ## Section: Servers -Servers represent assets, components, or systems in the threat model. +Servers represent assets, components, or systems in the system model. ### Server Attributes @@ -482,7 +484,7 @@ Use `config/asset_types_community.yaml` and `config/protocols_community.yaml` to **⚠️ CRITICAL FORMAT REQUIREMENT**: SecOpsTM parser requires EXACT DSL syntax. DO NOT use custom Markdown formats like `### Boundary:` or `**Name**:`. Use ONLY the list format shown below. ```markdown -# Threat Model: {System English Name} +# System Model: {System English Name} ## Description {High-level system description in English} diff --git a/.claude/skills/secopstm-threat-modeling/references/modeling-strategy.md b/.claude/skills/secopstm-threat-modeling/references/modeling-strategy.md index 8022196..5f73736 100644 --- a/.claude/skills/secopstm-threat-modeling/references/modeling-strategy.md +++ b/.claude/skills/secopstm-threat-modeling/references/modeling-strategy.md @@ -41,7 +41,7 @@ Use single-model when **ALL** of the following are true: ``` {system-name}/ -├── model.md # Single threat model +├── model.md # Single system model ├── context/ │ └── {system-name}_context.yaml ├── BOM/ diff --git a/.claude/skills/secopstm-threat-modeling/references/naming-conventions.md b/.claude/skills/secopstm-threat-modeling/references/naming-conventions.md index 7535bbd..cafee26 100644 --- a/.claude/skills/secopstm-threat-modeling/references/naming-conventions.md +++ b/.claude/skills/secopstm-threat-modeling/references/naming-conventions.md @@ -2,7 +2,7 @@ ## Overview -Proper naming conventions ensure consistency and readability in threat models. +Proper naming conventions ensure consistency and readability in system models. ## Asset Names diff --git a/.claude/skills/secopstm-threat-modeling/references/output-structure.md b/.claude/skills/secopstm-threat-modeling/references/output-structure.md index d5fd1bd..3469fe2 100644 --- a/.claude/skills/secopstm-threat-modeling/references/output-structure.md +++ b/.claude/skills/secopstm-threat-modeling/references/output-structure.md @@ -8,7 +8,7 @@ This document describes the output structure for SecOpsTM threat modeling. ``` {system-name}/ -├── model.md # Threat model +├── model.md # System model ├── context/ │ └── {system-name}_context.yaml ├── BOM/ @@ -22,7 +22,7 @@ This document describes the output structure for SecOpsTM threat modeling. | File | Required | Description | |------|----------|-------------| -| `model.md` | ✅ | Main threat model file in SecOpsTM DSL format | +| `model.md` | ✅ | Main system model file in SecOpsTM DSL format | | `context/{system-name}_context.yaml` | ✅ | GDAF attack context (objectives, threat actors, risk criteria) | | `BOM/{asset_name}.yaml` | ✅ | Bill of Materials - one file per key asset | | `config/asset_types_community.yaml` | ❌ | Custom asset type definitions (if needed) | @@ -56,7 +56,7 @@ This document describes the output structure for SecOpsTM threat modeling. |------|----------|-------------| | `main.md` | ✅ | Top-level architecture showing subsystem boundaries and data flows | | `context/{project-name}_context.yaml` | ✅ | Global GDAF context for the entire project | -| `{subsystem}/model.md` | ✅ | Subsystem-specific threat model | +| `{subsystem}/model.md` | ✅ | Subsystem-specific system model | | `{subsystem}/BOM/{asset}.yaml` | ✅ | Subsystem-specific BOM files | | `config/*` | ❌ | Custom asset types and protocols (shared across subsystems) | diff --git a/.claude/skills/secopstm-threat-modeling/references/staging-format.md b/.claude/skills/secopstm-threat-modeling/references/staging-format.md index d46e8cc..9fced72 100644 --- a/.claude/skills/secopstm-threat-modeling/references/staging-format.md +++ b/.claude/skills/secopstm-threat-modeling/references/staging-format.md @@ -2,7 +2,7 @@ ## Overview -Before generating the threat model, collect and structure information in the `{project-name}/.tm-raw/` directory. This enables incremental collection, source tracking, and conflict resolution. +Before generating the system model, collect and structure information in the `{project-name}/.tm-raw/` directory. This enables incremental collection, source tracking, and conflict resolution. ## Directory Structure diff --git a/.claude/skills/secopstm-threat-modeling/references/workflow.md b/.claude/skills/secopstm-threat-modeling/references/workflow.md index b732e03..37edff1 100644 --- a/.claude/skills/secopstm-threat-modeling/references/workflow.md +++ b/.claude/skills/secopstm-threat-modeling/references/workflow.md @@ -266,11 +266,11 @@ After information gathering and complexity assessment, **present the modeling st - Rationale: {Why this decomposition} Please confirm if the modeling strategy is appropriate: -1. Confirm - Generate threat model with this strategy +1. Confirm - Generate system model with this strategy 2. Adjust - User provides modification suggestions (e.g., merge some subsystems/change to single-model) ``` -**After user confirmation**, then generate the specific threat model files. +**After user confirmation**, then generate the specific system model files. --- ### Step 6: Generate Model File diff --git a/.gitignore b/.gitignore index 4ec6900..94c28e3 100644 --- a/.gitignore +++ b/.gitignore @@ -210,8 +210,15 @@ package-lock.json node_modules/ .claude/settings.local.json CLAUDE.md +decisions.md +tasks.md external_data.tar.gz.sha256 external_data.tar.gz # Git worktrees .worktrees/ + +# Subagent-driven-development scratch workspace (briefs, diffs, ledger) +.superpowers/ +*.docx +/docs/superpowers diff --git a/DOCKER_HUB.md b/DOCKER_HUB.md index 9291e9e..49b4c26 100644 --- a/DOCKER_HUB.md +++ b/DOCKER_HUB.md @@ -11,14 +11,18 @@ An [OWASP project](https://owasp.org/www-project-secopstm/) · Apache 2.0 ## Quick start — offline, no API key ```bash -docker run -p 5000:5000 \ +docker run -p 127.0.0.1:5000:5000 \ -v $(pwd)/output:/app/output \ ellipse2v/secopstm:latest ``` -Open **http://localhost:5000**, paste your threat model and click **Generate**. +Open **http://localhost:5000**, paste your system model and click **Generate**. Reports land in `$(pwd)/output//`. +> **Security note:** there is no authentication on any route (single-user +> tool by design). `-p 127.0.0.1:5000:5000` binds to localhost only. Only +> use `-p 5000:5000` (all interfaces) on a network you fully trust. + --- ## Quick start — with AI enrichment (LLM + RAG) @@ -37,8 +41,8 @@ docker run --rm \ **Step 2 — Run the server:** ```bash -docker run -p 5000:5000 \ - -e NVIDIA_API_KEY=your_key \ +docker run -p 127.0.0.1:5000:5000 \ + -e NVIDIA_NIM_API_KEY=your_key \ -v secopstm-rag:/app/rag \ -v $(pwd)/output:/app/output \ ellipse2v/secopstm:latest @@ -52,10 +56,12 @@ The named volume `secopstm-rag` persists across restarts and image upgrades. | Provider | Environment variable | Free tier | |---|---|---| -| NVIDIA NIM *(default)* | `NVIDIA_API_KEY` | ✅ [build.nvidia.com](https://build.nvidia.com/meta/llama-3_3-70b-instruct) | +| NVIDIA NIM *(default)* | `NVIDIA_NIM_API_KEY` | ✅ [build.nvidia.com](https://build.nvidia.com/meta/llama-3_3-70b-instruct) | | Google Gemini | `GEMINI_API_KEY` | ✅ | | OpenAI | `OPENAI_API_KEY` | ❌ | | Mistral | `MISTRAL_API_KEY` | ✅ | +| Groq | `GROQ_API_KEY` | ✅ [console.groq.com](https://console.groq.com) | +| xAI | `XAI_API_KEY` | ❌ [console.x.ai](https://console.x.ai) | | Ollama (local, fully offline) | — no key needed — | ✅ | --- @@ -65,7 +71,7 @@ The named volume `secopstm-rag` persists across restarts and image upgrades. | What | Docker flag | Notes | |---|---|---| | Output reports | `-v $(pwd)/output:/app/output` | Files land in `output//` on the host | -| Threat model files | `-v $(pwd)/models:/models` | Then pass `--model-file /models/model.md` | +| System model files | `-v $(pwd)/models:/models` | Then pass `--model-file /models/model.md` | | AI config | `-v $(pwd)/ai_config.yaml:/app/config/ai_config.yaml:ro` | Change provider or model without rebuilding | | LLM prompts | `-v $(pwd)/prompts.yaml:/app/config/prompts.yaml:ro` | Override system prompts | | CVE definitions | `-v $(pwd)/cve_definitions.yml:/app/cve_definitions.yml:ro` | Per-asset CVE list | @@ -74,8 +80,8 @@ The named volume `secopstm-rag` persists across restarts and image upgrades. ### Full example with all mounts ```bash -docker run -p 5000:5000 \ - -e NVIDIA_API_KEY=your_key \ +docker run -p 127.0.0.1:5000:5000 \ + -e NVIDIA_NIM_API_KEY=your_key \ -v $(pwd)/ai_config.yaml:/app/config/ai_config.yaml:ro \ -v secopstm-rag:/app/rag \ -v $(pwd)/models:/models \ @@ -88,7 +94,7 @@ docker run -p 5000:5000 \ ## Changing the port ```bash -docker run -p 8080:5000 ... # accessible on http://localhost:8080 +docker run -p 127.0.0.1:8080:5000 ... # accessible on http://localhost:8080 ``` --- diff --git a/README.md b/README.md index 3811a44..b1f7420 100644 --- a/README.md +++ b/README.md @@ -15,7 +15,7 @@ This project is a Python-based, end-to-end STRIDE threat modeling and analysis f - **Generate optimized Attack Flow diagrams** for key objectives (Tampering, Spoofing, Information Disclosure, Repudiation). - **Extend and customize** all mappings, calculations, and reporting logic. - **Run as a web-based editor** for live, interactive threat modeling. -- **AI-Enhanced Threat Analysis (Hybrid Mode)**: Threats from three independent engines — pytm rule engine, component-level LLM, and a cross-model RAG pipeline (ChromaDB + HuggingFace) — are automatically deduplicated and unified before reporting. Boundary objects are also analysed as AI targets. Supports Ollama (offline), Gemini, OpenAI, Mistral, and any LiteLLM-compatible provider. Configured in `config/ai_config.yaml`. +- **AI-Enhanced Threat Analysis (Hybrid Mode)**: Threats from three independent engines — pytm rule engine, component-level LLM, and a cross-model RAG pipeline (ChromaDB + HuggingFace) — are automatically deduplicated and unified before reporting. Boundary objects are also analysed as AI targets. Supports Ollama (offline), Gemini, OpenAI, Mistral, Groq, xAI, and any LiteLLM-compatible provider. Configured in `config/ai_config.yaml`. - **Pure CLI & CI integration**: A `secopstm` command ships after `pip install -e .`. Use `--output-format json --stdout` to pipe structured output to dashboards or SIEM without starting a server. - **Versioned JSON output**: Every JSON export is stamped `schema_version: "1.0"` and validated against `threat_analysis/schemas/v1/threat_model_report.schema.json`. @@ -23,13 +23,34 @@ This project is a Python-based, end-to-end STRIDE threat modeling and analysis f [![Python 3.8+](https://img.shields.io/badge/python-3.8+-blue.svg)](https://www.python.org/downloads/) [![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0) -[![Maintenance](https://img.shields.io/badge/Maintained%3F-yes-green.svg)](https://github.com/yourusername/cybersec-mapping/graphs/commit-activity) +[![Maintenance](https://img.shields.io/badge/Maintained%3F-yes-green.svg)](https://github.com/ellipse2v/SecOpsTM/graphs/commit-activity) + +--- + +## 🤖 AI Roles + +When AI is enabled, SecOpsTM runs several distinct LLM passes, each with its own system +prompt and purpose (all defined in `config/prompts.yaml`, editable without touching Python +code). All roles share the one provider configured in `config/ai_config.yaml` — there's no +per-role provider selection — and every pass degrades gracefully offline: if no provider is +reachable, the tool falls back to the pytm rule engine silently. + +| Role | Purpose | Where it shows up | +|---|---|---| +| **DSL Generation** | Turns a natural-language system description into a valid SecOpsTM Markdown model | "Generate from prompt" in the web editor | +| **STRIDE Component Analysis** | Generates STRIDE threats per actor/server/boundary, strictly scoped to the component's declared technology tags (no invented stacks) | Per-component AI enrichment, `(AI)`-sourced threats in the report | +| **RAG System-Level Analysis** | Cross-component/cross-boundary threats retrieved from a local ChromaDB knowledge base (CVE, CAPEC, ATT&CK, D3FEND) | System-level pass, `(LLM)`-sourced threats — skipped on small models below `rag.min_components` | +| **SOC Analyst** | Detection-engineering pass: Sigma/Splunk SPL/KQL rule suggestions and IOCs per threat | HTML report "SOC Analysis" section | +| **CISO Triage** | Board-level risk briefing summarizing the highest-priority threats | HTML report "CISO Briefing" section | +| **Red Team** | Attempts to advance a GDAF attack scenario, citing only facts from the grounding block (CVE, misconfig, exposed ports) | Red/Blue adversarial debate on top GDAF scenarios | +| **Blue Team** | Challenges the Red team's attempt with SIEM/EDR/IDS controls and calls out detection gaps | Red/Blue adversarial debate | +| **Attack Path Narrator** | Short grounded narrative + business impact for each automatically discovered attack path (already-fixed hops/techniques — explains them, never adds new ones). Instructed to never emit an ID (T-number, CVE, CAPEC, D3-); any response that does is discarded, not trusted | "Automatically Discovered Attack Paths" section | --- ## ✨ New Interactive Features -The framework's web interface has been significantly updated to provide a more powerful and seamless threat modeling experience. +The web interface now supports the following. ### Interactive Diagrams The generated diagrams (both in the live editor and in exported HTML reports) are not static images. They are fully interactive SVGs that allow you to: @@ -40,13 +61,12 @@ The generated diagrams (both in the live editor and in exported HTML reports) ar ### Interactive Legend - **Filter Connections**: The diagram legend is now interactive. Click on a protocol (e.g., HTTPS, TCP) to instantly show or hide all dataflows using that protocol, making it easy to analyze specific parts of your data flow. -### Robust Project Generation -- The **"Generate All"** feature is designed to handle complex projects with nested sub-models. -- If you have a project with multiple threat model files, it intelligently detects if any referenced sub-models are not currently open in the editor and prompts you to select your project's root directory. This ensures that all necessary files are found and included for a complete and accurate generation of all reports and diagrams. +### Project Generation +- The **"Generate All"** feature handles projects with nested sub-models. +- If you have a project with multiple system model files, it detects when a referenced sub-model is not currently open in the editor and prompts you to select your project's root directory, so all files are found before generating reports and diagrams. ### Simple Server Mode -- The integrated web server can be started with a path to a project directory (`--project path/to/your/project`). It will automatically find all `*.md` threat model files within that project and open them in tabs, ready for editing. -[![Maintenance](https://img.shields.io/badge/Maintained%3F-yes-green.svg)](https://github.com/ellipse2v/cybersec-mapping/graphs/commit-activity) +- The integrated web server can be started with a path to a project directory (`--project path/to/your/project`). It automatically finds all `*.md` system model files within that project and opens them in tabs, ready for editing. --- @@ -61,13 +81,18 @@ For detailed information on features, usage, and advanced customization, please ### Option A — Docker (no Python setup required) ```bash -docker run -p 5000:5000 \ +docker run -p 127.0.0.1:5000:5000 \ -v $(pwd)/output:/app/output \ ellipse2v/secopstm:latest ``` Open . Reports land in `$(pwd)/output//`. +> **Security note:** the server has no authentication on any route (it's a +> single-user tool). `-p 127.0.0.1:5000:5000` keeps it reachable only from +> this machine. Only bind it to all interfaces (`-p 5000:5000`) on a network +> you fully trust. + #### With AI enrichment (LLM + RAG) **Default provider: NVIDIA NIM** (Llama 3.3 70B) — free API key at https://build.nvidia.com/meta/llama-3_3-70b-instruct @@ -77,14 +102,14 @@ Open . Reports land in `$(pwd)/output//`. docker run --rm -v secopstm-rag:/app/rag ellipse2v/secopstm:latest --init-rag # Step 2 — Run -docker run -p 5000:5000 \ - -e NVIDIA_API_KEY=your_key \ +docker run -p 127.0.0.1:5000:5000 \ + -e NVIDIA_NIM_API_KEY=your_key \ -v secopstm-rag:/app/rag \ -v $(pwd)/output:/app/output \ ellipse2v/secopstm:latest ``` -Other supported providers: `GEMINI_API_KEY`, `OPENAI_API_KEY`, `MISTRAL_API_KEY`. Ollama works fully offline (no key needed). +Other supported providers: `GEMINI_API_KEY`, `OPENAI_API_KEY`, `MISTRAL_API_KEY`, `GROQ_API_KEY`, `XAI_API_KEY`. Ollama works fully offline (no key needed). --- diff --git a/ROADMAP_SUMMARY.md b/ROADMAP_SUMMARY.md index 52961c3..7434686 100644 --- a/ROADMAP_SUMMARY.md +++ b/ROADMAP_SUMMARY.md @@ -3,9 +3,9 @@ This roadmap outlines the SecOpsTM framework's current capabilities and future plans. ### Implemented Features: -- **Hierarchical & Interactive Threat Models**: Support for linked sub-models and navigable diagrams. +- **Hierarchical & Interactive System Models**: Support for linked sub-models and navigable diagrams. - **Automated Mitigations**: Context-aware suggestions from OWASP ASVS, NIST, and CIS Controls. -- **Ansible IaC Integration**: Automatic threat model generation from Ansible playbooks. +- **Ansible IaC Integration**: Automatic system model generation from Ansible playbooks. - **Web-Based GUI**: Interactive web interface for real-time model editing and preview. - **MITRE ATT&CK Integration**: Export to Navigator and STIX 2.1 formats for visualization and threat intelligence. diff --git a/SECURITY.md b/SECURITY.md index e21f508..caa8eaa 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -39,7 +39,7 @@ - Remote code execution via the web server (`--server` mode) - Unauthorized access to the filesystem through the API - - Injection vulnerabilities (YAML, template, command) in threat model parsing + - Injection vulnerabilities (YAML, template, command) in system model parsing - Authentication bypass (if authentication is added in future versions) The following are **out of scope**: diff --git a/architecture.md b/architecture.md index 0b39c85..ac533d5 100644 --- a/architecture.md +++ b/architecture.md @@ -13,9 +13,10 @@ threatModelBypyTm/ │ ├── threat_rules.py Additional threat rule definitions │ ├── mitigation_suggestions.py MitigationStixMapper, framework mitigations │ ├── severity_calculator_module.py SeverityCalculator + RiskContext VOC scoring -│ ├── update_config.py CLI helper to update ai_config.yaml +│ ├── update_config.py Dev script: regenerates static/js/config.js from +│ │ config_generator.py — run manually before commits │ ├── utils.py _validate_path_within_project, extract_json_from_llm_response -│ ├── data_loader.py Lazy loaders for external data files +│ ├── validate.py `secopstm validate` subcommand — offline DSL/config lint │ │ │ ├── core/ Domain model │ │ ├── models_module.py ThreatModel, ExtendedThreat, CustomThreat @@ -24,29 +25,39 @@ threatModelBypyTm/ │ │ ├── model_parser.py ModelParser — Markdown DSL → ThreatModel (2-pass) │ │ ├── model_factory.py create_threat_model() — wires parser + model │ │ ├── model_validator.py ModelValidator — pre-process validation +│ │ ├── model_completeness.py ModelCompletenessChecker — DSL quality score +│ │ ├── dsl_constants.py DSL_ENUMS — single source of truth for DSL field values │ │ ├── mitre_mapping_module.py MitreMapping — STRIDE→CAPEC→ATT&CK→D3FEND │ │ ├── mitre_static_maps.py Hard-coded ATTACK_D3FEND_MAPPING dict +│ │ ├── attack_id_validator.py AttackIdValidator — validates ATT&CK/CAPEC/D3FEND IDs +│ │ │ (incl. LLM-cited ones) against the committed corpora +│ │ ├── asset_technique_mapper.py AssetTechniqueMapper — asset type/attrs → ATT&CK techniques │ │ ├── cve_service.py CVEService — single-pass JSONL (CAPEC+CWE) + YAML definitions +│ │ ├── data_loader.py Lazy loaders for external_data/ files (ATT&CK, CAPEC, D3FEND, NIST…) │ │ ├── vex_loader.py VEXLoader — CycloneDX VEX document loader (standalone/dir/auto) │ │ ├── bom_loader.py BOMLoader — BOM YAML per asset; active_cves/fixed_cves via VEX state +│ │ ├── accepted_risks.py AcceptedRiskLoader — analyst risk-acceptance decisions per threat │ │ ├── ai_cache.py AIThreatCache — SHA-256-keyed cache (.secopstm_ai_cache.json) -│ │ ├── attack_chain.py AttackChainAnalyzer — graph traversal, chained threat paths +│ │ ├── attack_chain.py AttackChainAnalyzer — bottom-up: chains existing threats via dataflows +│ │ ├── gdaf_engine.py GDAFEngine — top-down: objectives + threat actors → attack scenarios +│ │ ├── debate_engine.py RedBlueDebateEngine — adversarial Red/Blue debate over GDAF scenarios │ │ ├── threat_consolidator.py ThreatConsolidator — Jaccard dedup, AI wins over pytm +│ │ ├── threat_ranker.py ThreatRanker — weighted composite rank + trim of consolidated threats +│ │ ├── stride_constants.py STRIDE_CATEGORIES — single source of truth for the 6 STRIDE names │ │ └── report_serializer.py ReportSerializer — stable versioned dict, IDs T-NNNN │ │ │ ├── ai_engine/ AI inference layer │ │ ├── embedding_factory.py get_embeddings() — provider-agnostic factory +│ │ ├── prompt_loader.py Loads config/prompts.yaml sections for each AI role │ │ ├── rag_service.py RAGThreatGenerator — ChromaDB + LangChain RAG chain │ │ ├── providers/ -│ │ │ ├── base_provider.py BaseLLMProvider (ABC): check_connection, -│ │ │ │ generate_threats, generate_attack_flow +│ │ │ ├── base_provider.py BaseLLMProvider (ABC): check_connection, generate_threats │ │ │ ├── litellm_client.py LiteLLMClient — low-level async LiteLLM wrapper │ │ │ │ (static factory create(), generate_content generator) │ │ │ ├── litellm_provider.py LiteLLMProvider(BaseLLMProvider) — orchestrates client │ │ │ └── ollama_provider.py OllamaProvider(BaseLLMProvider) — Ollama-specific │ │ └── prompts/ -│ │ ├── stride_prompts.py STRIDE_SYSTEM_PROMPT, build_component_prompt() -│ │ └── attack_flow_prompts.py ATTACK_FLOW_SYSTEM_PROMPT, build_attack_flow_prompt() +│ │ └── stride_prompts.py STRIDE_SYSTEM_PROMPT, build_component_prompt() │ │ │ ├── generation/ Output artifact generators │ │ ├── diagram_generator.py DiagramGenerator — DOT/SVG via Graphviz subprocess @@ -54,7 +65,9 @@ threatModelBypyTm/ │ │ ├── report_generator.py ReportGenerator — HTML report via Jinja2 │ │ ├── stix_generator.py StixGenerator — STIX 2.1 bundle JSON │ │ ├── attack_navigator_generator.py AttackNavigatorGenerator — Navigator layer JSON -│ │ ├── attack_flow_generator.py AttackFlowGenerator — Attack Flow STIX objects +│ │ ├── attack_flow_generator.py AttackFlowGenerator — STRIDE-technique attack paths +│ │ │ (pure graph traversal, no LLM — see decisions.md) +│ │ ├── attack_flow_builder.py AttackFlowBuilder — writes GDAF scenarios to .afb files │ │ ├── graphviz_to_json_metadata.py DOT → JSON with element metadata │ │ ├── graphviz_to_konva.py DOT → Konva.js canvas JSON (GUI editor) │ │ ├── tactic_logic.py Tactic ordering and filtering helpers @@ -149,7 +162,7 @@ Markdown file .boundaries{} .actors[] .servers[] .dataflows[] → ThreatModel.process_threats() → pytm.TM.process() (PyTM built-in rules) - → _expand_class_targets() (class → instances) + → reads tm.findings (condition-matched, already per-instance) → _apply_custom_threats() (custom_threats.py rules) → _group_threats() (by STRIDE category) → _perform_mitre_analysis() (MitreMapping.analyze_pytm_threats_list()) @@ -266,6 +279,72 @@ ReportGenerator._get_all_threats_with_mitre_info() → injected into HTML report as "⛓️ Attack Chain Analysis" section ``` +### 5b. GDAF Attack Path Generation + Red/Blue Debate + +``` +GDAFEngine.run(threat_model, context_yaml) (top-down, complementary to AttackChainAnalyzer) + → reads attack_objectives + threat_actors from context YAML + (or _auto_context() — a minimal context synthesized from the model's servers if no YAML given) + → _build_graph() (nodes = components, edges = dataflows + + sub-model bridging edges for submodel= servers) + → per (objective, threat_actor) pair: graph traversal from actor's entry point to objective + → hop_weight includes boundary.traversal_difficulty bonus (low=+0.3, medium=+0.1, high=+0.0) + → per-hop MITRE technique assignment (AssetTechniqueMapper) + → returns List[AttackScenario], stored as threat_model.gdaf_scenarios + → AttackFlowBuilder(gdaf_scenarios).generate_and_save() (.afb export, MITRE Attack Flow v3.0) + +ReportGenerator.generate_html_report() — after attack_id_validation, before CISO triage: + → if debate.enabled and gdaf_scenarios: + RedBlueDebateEngine.run(gdaf_scenarios) (one configured LLM provider, two personas) + → per top-N scenario, N rounds: + Red persona: advances the attack using only facts in the grounding block + Blue persona: blocks/detects, cites SIEM/EDR/IDS, lists detection_gaps + → mutates scenario.score / risk_level in place (does NOT create new threats) + → re-run AttackFlowBuilder(...).generate_and_save() (.afb files re-written so they + agree with debate-adjusted scores) +``` + +### 5c. SOC Analyst + CISO Triage + +``` +ReportGenerator.generate_html_report() — after the debate re-write, before the final template render: + → _build_threat_graph_data(threat_model, all_detailed_threats, gdaf_scenarios, debate_results) + → node/edge graph with GDAF path overlays, reflects debate-adjusted scores + + → if ai_provider and all_detailed_threats: + AIService._enrich_with_soc_analysis() (per-threat, soc_analyst prompt persona) + → Sigma / Splunk SPL / KQL rule suggestions + IOCs + → stored on threat.ai_details["soc_analysis"], rendered in the "SOC Analysis" section + + ReportGenerator._run_ciso_triage(all_detailed_threats, gdaf_scenarios, debate_results) + → LiteLLMProvider.generate_ciso_triage() (ciso_triage prompt persona) + → board-level risk summary, rendered in the "CISO Briefing" section +``` + +### 5d. Discovered Attack Paths + Narrative + +``` +ReportGenerator.generate_html_report() — after CISO triage, before the template render: + → if attack_flows.enabled: + AttackFlowGenerator(all_detailed_threats, allowed_categories=...).get_paths_summary() + → best (highest-severity) path per STRIDE category through the threats' own + MITRE techniques — pure graph traversal, no LLM call, no GDAF context needed + → each hop self-reports its own threat_id/threat_description (not positional + lookup into threat_ids[i] — that list silently skips threats with no id) + + → if ai_provider and attack_flows.include_narrative: + ReportGenerator._generate_path_narratives(discovered_attack_paths) (in place) + → per path (≤6 — one per STRIDE category): builds a grounding block from the + path's own hops (target, technique NAME — not ID, tactic, related threat + description) and calls LiteLLMProvider.generate_attack_path_narrative() + (attack_path_narrative prompt persona) + → persona is instructed to never emit an ID (T-number, CVE, CAPEC, D3-) — + _narrative_has_id_leakage() regex-checks the response regardless of + cooperation; any match discards the entire narrative (fail closed, not a + partial-trust patch), logged as a grounding violation + → on success, path["narrative"] / path["business_impact"] are set +``` + ### 6b. Diagram Generation + Trust Colors ``` @@ -319,26 +398,28 @@ tooling/build_vector_store.py 1. ~~**Duplicated JSON extraction logic**~~ — **Fixed**: `extract_json_from_llm_response()` is now a single shared function in `threat_analysis/utils.py`. -2. **pytm.Boundary monkey-patching** (`models_module.py:27-35`) — Adds `isTrusted`, `protocol`, - `port`, `data` attributes because pytm does not provide them. Labelled `# HACK` in code. - Breaks if pytm is updated. +2. ~~**pytm.Boundary monkey-patching**~~ — **Fixed**: replaced by the `SecOpsBoundary(Boundary)` + subclass in `models_module.py`, which adds `isTrusted`, `protocol`, `port`, `data` without + patching the pytm class at runtime. -3. **Sync wrapper over async generator** (`ai_service.py:generate_markdown_from_prompt_sync`) — - Uses `loop.run_until_complete(gen.__anext__())` per chunk inside a Flask thread. Fragile in - concurrent request scenarios. +3. ~~**Sync wrapper over async generator**~~ — **Fixed**: `generate_markdown_from_prompt_sync` + now submits to a persistent background event loop via `asyncio.run_coroutine_threadsafe()` + (see the "RAG init parallel to AI connection check" and cross-model RAG decisions in + `decisions.md`), instead of calling `loop.run_until_complete()` per chunk. 4. ~~**Hardcoded rate limit sleep**~~ — **Fixed**: `rate_limit_sleep` is now configurable via `ai_config.yaml → threat_generation.rate_limit_sleep`. 5. ~~**`_get_output_dir` defined twice**~~ — **Fixed**: duplicate removed from `ExportService`. -6. **`venv-py310/` in the repo tree** — virtualenv directory should be in `.gitignore`. +6. ~~**`venv-py310/` in the repo tree**~~ — **Fixed**: already listed in `.gitignore`. 7. **`requirements.txt` vs `pyproject.toml` diverge** — `pyproject.toml` lists minimal runtime deps; `requirements.txt` includes all AI/ML extras. Installing from `pyproject.toml` alone (`pip install -e .`) is insufficient for AI features; use `pip install -e ".[ai]"`. -8. **`package-lock.json` at root** — Untracked orphaned file (no `package.json` exists). +8. **`package-lock.json` at root** — Present on disk but already gitignored (no `package.json` + exists). Harmless local artifact, not tracked; safe to delete locally if it bothers you. 9. **`ThreatModel` requires `CVEService`** injected at construction — tight coupling; if CVE data files are missing the entire model fails to instantiate. @@ -360,3 +441,12 @@ tooling/build_vector_store.py from the server are injected with `innerHTML`. Component names in the DOT template are not HTML-escaped, which could allow stored XSS via crafted component names. Risk is limited to single-user local server but should be addressed before any multi-user deployment. + +14. ~~**`prompts.yaml`'s `attack_flow` persona was dead code**~~ — **Fixed**: removed. It was a + single-shot LLM attack-flow generator (`BaseLLMProvider.generate_attack_flow()` + the + `attack_flow_prompts.py` module + `prompts.yaml`'s `attack_flow:` section) that was only ever + called from unit tests, never from the application — see the "Attack Flow generation: graph + traversal, not a single-shot LLM persona" decision in `decisions.md` for why it stayed removed. + The real "Discovered Attack Paths" feature and `.afb` export (`generation/attack_flow_generator.py`, + `generation/attack_flow_builder.py`) are unaffected — pure STRIDE-technique graph traversal over + the fully consolidated pytm+AI+LLM threat list, no LLM call needed. diff --git a/config/ai_config.yaml b/config/ai_config.yaml index c279da5..b363fec 100644 --- a/config/ai_config.yaml +++ b/config/ai_config.yaml @@ -33,7 +33,10 @@ ai_providers: max_tokens: 4096 timeout: 30 ssl_verify: true # false | "/path/to/ca-bundle.crt" for self-signed certs - # enable_thinking: false # Anthropic extended thinking — set true only with a large max_tokens budget + # enable_thinking: false # vLLM/Qwen3-style hybrid thinking toggle — only set this on providers + # that actually accept an `enable_thinking` request field. The Anthropic API does NOT (it uses a + # nested `thinking: {type: ...}` object instead) — setting this under the anthropic: block causes + # every request to fail with "enable_thinking: Extra inputs are not permitted". gemini: enabled: false @@ -52,10 +55,29 @@ ai_providers: api_base: "" ssl_verify: true + # Anthropic API — https://console.anthropic.com (separate from a Claude Pro/Code + # subscription; needs its own pay-as-you-go API key) + # export ANTHROPIC_API_KEY=sk-ant-... + anthropic: + enabled: false + model: "claude-haiku-4-5" # fast/cheap for iteration; "claude-sonnet-5" for higher quality + api_key_env: "ANTHROPIC_API_KEY" + temperature: 0.2 + # 1024 truncated batch calls: the STRIDE schema asks for 3-5 threats per component + # with ~12 fields each (attack_scenario, business_impact object, capec/cwe/d3fend + # lists, ...), and this value also floors the batch ceiling (max_tokens x4 in + # generate_threats_batch) — 1024 capped a 5-component batch at 4096 tokens total + # and cut the JSON off mid-response. 8192 still truncated some 5-component batches + # on Haiku's verbose output — raised further. Haiku 4.5 supports up to 64K output, + # so this only raises the ceiling, not actual usage/cost. + max_tokens: 16000 + timeout: 30 + ssl_verify: true + # NVIDIA Integrate API (OpenAI-compatible) — https://build.nvidia.com/models # export NVIDIA_NIM_API_KEY=nvapi-... nvidia_nim: - enabled: true + enabled: false model: "meta/llama-3.3-70b-instruct" # any model from build.nvidia.com api_key_env: "NVIDIA_NIM_API_KEY" api_base: "https://integrate.api.nvidia.com/v1" @@ -65,6 +87,30 @@ ai_providers: timeout: 60 ssl_verify: true + # Groq (LPU inference — fast, generous free tier) — https://console.groq.com + # export GROQ_API_KEY=gsk_... + groq: + enabled: true + model: "llama-3.3-70b-versatile" # or "llama-3.1-8b-instant" (faster/cheaper) + api_key_env: "GROQ_API_KEY" + api_base: "" # defaults to https://api.groq.com/openai/v1 when empty + temperature: 0.3 + max_tokens: 4096 + timeout: 30 + ssl_verify: true + + # xAI Grok — https://console.x.ai (pay-as-you-go API key) + # export XAI_API_KEY=xai-... + xai: + enabled: false + model: "grok-3-latest" # or "grok-4-fast-reasoning", "grok-3-mini-fast-beta" (cheaper) + api_key_env: "XAI_API_KEY" + api_base: "" # defaults to https://api.x.ai/v1 when empty + temperature: 0.3 + max_tokens: 4096 + timeout: 30 + ssl_verify: true + ollama: enabled: false model: "llama3:8b-instruct-q4_K_M" @@ -73,7 +119,6 @@ ai_providers: num_ctx: 4096 num_predict: 4096 temperature: 0.3 - attack_flow_temperature: 0.2 # lower = more deterministic for structured JSON output timeout: 60 # aiohttp session timeout (seconds); connection-level ssl_verify: true # Note: Ollama retry count (3 attempts, exponential back-off) is fixed in @@ -85,12 +130,14 @@ ai_providers: threat_generation: mode: "hybrid" # "traditional" | "ai" | "hybrid" fallback_to_traditional: true - max_threats_per_component: 10 + # Per-component cap on AI-generated threats kept (0 = no limit). When more than this + # many are returned, the highest-confidence ones are kept. + max_threats_per_component: 7 # Rate limiting — set to 1.5 for Gemini free tier (15 RPM), 0.0 for paid/Ollama - rate_limit_sleep: 1.5 + rate_limit_sleep: 2 # Concurrency — 1 = sequential (safe for free tiers); 3–5 for paid/Ollama - max_concurrent_ai_requests: 3 + max_concurrent_ai_requests: 1 # Batch enrichment — components grouped per LLM call (5 = 5× fewer API round trips) # Set to 1 to disable batching and fall back to one call per component batch_size: 5 @@ -106,8 +153,9 @@ threat_generation: confidence: 0.3 risk_signals: 0.3 + # AI-generated threats below this confidence (0.0-1.0) are dropped before being + # attached to the model. confidence_threshold: 0.75 - include_attack_flows: true # LiteLLM retry policy — applies to all providers num_retries: 3 # automatic retries on rate-limit / transient errors @@ -137,6 +185,12 @@ rag: enabled: true vector_db: "chroma" top_k_examples: 5 + # Below this many components (actors + servers), RAG's system-level pass is + # skipped — cross-model/cross-boundary chaining needs several components to + # produce grounded output; on tiny models it mostly restates single-component + # threats already covered by the per-component AI pass. Component-level AI + # enrichment is unaffected. 0 disables the threshold (always run RAG). + min_components: 3 # Embeddings — default is local HuggingFace (no API key, works offline) # WARNING: changing the model requires rebuilding the vector store: @@ -148,13 +202,21 @@ embedding: api_base: "" ssl_verify: true -# Attack Flow generation +# "Automatically Discovered Attack Paths" — the STRIDE-technique-based attack path +# finder (AttackFlowGenerator), independent of GDAF. Scope: only the HTML report +# section built by generate_html_report(); the standalone .afb export flows (CLI +# --attack-flow, server's dedicated "Export Attack Flow" button) always use every +# STRIDE category regardless of generate_for_categories, since they're meant to be a +# complete, unfiltered dump of what pytm/AI found. attack_flows: enabled: true - format: "stix_2.1" - max_depth: 7 - include_detection_points: true generate_for_categories: ["Tampering", "Information Disclosure", "Elevation of Privilege"] + # Adds a short grounded narrative + business_impact to each discovered path in the + # HTML report (attack_path_narrative persona, config/prompts.yaml). One LLM call per + # path (naturally capped at one per STRIDE category, so ≤6). Requires AI enabled; + # skipped silently offline. The persona is instructed to never emit an ID (no + # T-numbers, CVE, CAPEC, D3-) — any response that does is discarded, not trusted. + include_narrative: true # GDAF engine defaults — override per model via gdaf_min_technique_score in ## Context # or via risk_criteria in the per-model GDAF context YAML @@ -165,6 +227,17 @@ gdaf: max_paths_per_objective: 3 gdaf_min_technique_score: 0.8 +# Red/Blue adversarial debate over the top-N GDAF attack scenarios (opt-in). +# See docs/gdaf.md. Runs after GDAF and before the HTML report is written. +debate: + enabled: true + top_n: 5 + min_viability_threshold: 0.5 + max_rounds: 3 + viability_delta_threshold: 0.1 + debate_factor_min: 0.5 + debate_factor_max: 1.5 + # Fine-tuning (Phase 3 — optional, not enabled by default) fine_tuning: enabled: false diff --git a/decisions.md b/decisions.md deleted file mode 100644 index 3593ca2..0000000 --- a/decisions.md +++ /dev/null @@ -1,497 +0,0 @@ -# Technical Decisions — SecOpsTM - ---- - -## Decision: pytm as the core threat rule engine - -**Choice:** Build on top of the `pytm` library rather than writing STRIDE rules from scratch. - -**Why it matters:** pytm provides a mature, tested rule set for STRIDE applied to Actor/Server/Dataflow -topologies. It handles threat enumeration automatically once the model is built. This means the -project gets rule coverage for free and can focus on enrichment (MITRE mapping, AI, reporting). - -**Trade-off accepted:** pytm's data model is rigid (Boundary, Actor, Server, Dataflow, Data) and its -threat targeting (class-level, not instance-level) required `_expand_class_targets()` to fix. -The `Boundary` class also lacks several attributes needed by the project, requiring runtime -monkey-patching. - ---- - -## Decision: Custom Markdown DSL instead of Python API - -**Choice:** Users define their threat models in a custom Markdown format (`## Actors`, `## Servers`, -`## Dataflows`, etc.) rather than writing Python code against the pytm API. - -**Why it matters:** Makes threat models accessible to security architects who are not Python -developers. The DSL files are human-readable, version-controllable, and can be edited in the -Monaco editor in the web UI. Templates in `threatModel_Template/` lower the barrier to entry. - -**Trade-off accepted:** The `ModelParser` is custom-built and fragile — section headers are -hard-coded strings, and the 2-pass design adds complexity. Breaking the DSL grammar breaks all -existing model files. - ---- - -## Decision: LiteLLM as the LLM abstraction layer - -**Choice:** All LLM calls go through LiteLLM (`litellm.acompletion`) rather than vendor-specific -SDKs (e.g., `openai`, `google-generativeai`). - -**Why it matters:** A single provider can be swapped by changing one line in `ai_config.yaml` -(`enabled: true/false`). Supports Gemini, OpenAI, Mistral, and Ollama with a unified interface. -Enterprise users can point `api_base` at an internal proxy (OpenAI-compatible) without code changes. - -**What was rejected:** Direct SDK calls per provider would require separate code paths and would not -support the "first enabled provider wins" selection logic cleanly. - ---- - -## Decision: RAG pipeline using ChromaDB + LangChain - -**Choice:** Retrieval-Augmented Generation using a local ChromaDB vector store pre-built from -the `external_data/` knowledge base (CAPEC, CVE, ATT&CK, D3FEND, CIS, NIST). - -**Why it matters:** Enables the LLM to ground its threat generation in real security knowledge -(CVE precedents, CAPEC attack patterns) without fine-tuning. The vector store is built once -offline (`tooling/build_vector_store.py`) and ships with the project. - -**What was rejected:** Fine-tuning (Phase 3 in the roadmap) was explicitly deferred — too costly -and less flexible than RAG for a knowledge base that updates frequently. - ---- - -## Decision: HuggingFace `all-MiniLM-L6-v2` as default embedding model - -**Choice:** Local sentence-transformer model run on CPU as the default for vector similarity search. - -**Why it matters:** Zero network dependency and zero API cost for embeddings. Supports fully -air-gapped/sovereign deployments. The model is small (80 MB) and fast enough for the query -volumes involved (single-user tool, not a SaaS). - -**What was rejected:** Cloud embedding APIs (Google, OpenAI) are supported via `embedding_factory.py` -but not the default, to avoid mandatory API key setup for the RAG feature. - ---- - -## Decision: Flask (not FastAPI) as the web framework - -**Choice:** Flask with `Flask[async]` extension. - -**Why it matters:** The project started without AI streaming and Flask was the natural choice for -a simple single-user tool. `Flask[async]` was added to support SSE streaming of AI generation -output without a full framework migration. - -**Trade-off accepted:** The sync/async boundary is awkward — `generate_markdown_from_prompt_sync` -uses `loop.run_until_complete(gen.__anext__())` per chunk, which is not safe under concurrent load. -FastAPI would handle this more cleanly but would require significant route refactoring. - ---- - -## Decision: Server-Sent Events (SSE) for AI progress - -**Choice:** Use a global `queue.Queue` (`ai_status_event_queue`) and SSE endpoint -(`/api/ai_status_stream`) to push real-time AI processing progress to the browser. - -**Why it matters:** AI enrichment (per-component LLM calls) can take 30–120 seconds. Without live -feedback the UI appears frozen. SSE is simpler than WebSockets for unidirectional server→client -streams and requires no additional dependencies. - -**What was rejected:** Polling (too many requests) and WebSockets (overkill for unidirectional flow). - ---- - -## Decision: Lazy imports throughout the codebase - -**Choice:** Heavy imports (`litellm`, `langchain`, `chromadb`, `sentence_transformers`) are done -inside methods or `__init__`, not at module top level. - -**Why it matters:** The Flask server starts fast even when AI features are disabled or dependencies -are not installed. `ThreatModelService` is also lazy-initialized on first request. - -**Consequence:** Import errors for optional AI deps surface at runtime, not at import time. This -is intentional — the tool works without AI deps if AI config is disabled. - ---- - -## Decision: `ExtendedThreat` subclassing `pytm.Threat` - -**Choice:** AI-generated and custom threats are represented as `ExtendedThreat(pytm.Threat)` with -an added `source` attribute (`"pytm"`, `"AI"`, `"LLM"`). - -**Why it matters:** Allows the rest of the pipeline (grouping, MITRE mapping, reporting) to treat -all threats uniformly regardless of origin. The `source` field enables the report to distinguish -traditional vs. AI-generated threats and apply different rendering. - -**What was rejected:** A parallel `AIThreat` class with no pytm inheritance would have required -forking all downstream threat processing logic. - ---- - -## Decision: Offline data pipeline in `tooling/` - -**Choice:** Scripts that download, parse, and transform security data (MITRE ATT&CK JSON, CAPEC XML, -CVE JSONL, NIST XLSX) are separate one-shot tools, not part of the runtime application. - -**Why it matters:** The transformed artifacts (`external_data/*.json`, `vector_store/`) are committed -to the repo so end users do not need network access or tool dependencies to use the threat model. -The pipeline only needs to run when upstream data sources update. - -**Trade-off accepted:** The repo is large (CVE JSONL files from 1999–2025 are significant). -The vector store binary files are also committed, which inflates repo size. - ---- - -## Decision: `config/ai_config.yaml` as the single AI configuration surface - -**Choice:** All AI provider selection, model parameters, RAG settings, and embedding config live -in one YAML file rather than environment variables or a `.env` file. - -**Why it matters:** A single file to edit for setup. Enterprise users can template this file -for different deployment environments. The `api_key_env` indirection keeps secrets in environment -variables while keeping non-sensitive config in the file. - -**Consequence:** The YAML file must be present at runtime even if AI is disabled. Path is resolved -relative to `PROJECT_ROOT` (detected via `Path(__file__).resolve().parents[N]`). - ---- - -## Decision: Jaccard-based deduplication — AI source wins - -**Choice:** When a pytm threat and an AI/LLM threat cover the same `(target, stride_category)` with -similar descriptions (Jaccard word-overlap ≥ 0.3 or substring containment), the AI version replaces -the pytm version. Implemented in `threat_analysis/core/threat_consolidator.py`. - -**Why it matters:** Without deduplication, the same risk (e.g. "SQL injection on DB server") appears -two or three times in the report — once from pytm rules, once from the component LLM, possibly again -from the RAG pipeline. This inflates threat counts and confuses prioritisation. - -**Why Jaccard and not embedding similarity:** Embedding similarity would require loading the -sentence-transformer model during report generation, breaking the offline-without-AI-deps path. -Jaccard on word sets is purely offline, instantaneous, and sufficient for the overlap patterns -observed in practice (same STRIDE category, same target, paraphrased description). - -**Why AI wins over pytm:** AI-generated threats include richer context (attack scenario, recommended -controls). Keeping the AI version provides more actionable output. The pytm threat is only removed -when similarity is confirmed — unrelated pytm threats are always preserved. - ---- - -## Decision: Versioned JSON schema — `schema_version: "1.0"` - -**Choice:** Every JSON export is serialised by `ReportSerializer` and stamped `schema_version: "1.0"`. -The schema is validated against `threat_analysis/schemas/v1/threat_model_report.schema.json` -(JSON Schema 2020-12). Threats get stable IDs in the format `T-0001`. - -**Why it matters:** SIEM integrations, dashboards, and CI gates that consume the JSON output break -silently when the structure changes. A versioned, validated schema gives consumers a stable contract -and allows the tool to signal breaking changes explicitly. - -**What was rejected:** Ad-hoc dict construction (the previous approach) offered no contract and -generated different key sets depending on which code path ran. Migration to a serialiser was -preferred over adding version guards to the ad-hoc dict. - ---- - -## Decision: RiskContext VOC scoring — CWE as CVSS proxy - -**Choice:** Instead of a CVSS score (absent from the CVE JSONL files), the `RiskContext` dataclass -encodes four binary context signals that adjust the STRIDE base score: -- CVE match for target/category → +0.5 (confirmed exploitability evidence) -- High-risk CWE class (injection, memory corruption, hardcoded creds…) → +0.3 -- Network-exposed without auth/encryption → +0.7 -- D3FEND defensive mitigations in place → −0.5 - -**Why CWE instead of CVSS:** Inspection of the CVE JSONL corpus revealed that entries only carry -`CWE`, `CAPEC`, and `TECHNIQUES` keys — no CVSS scores. CWE class is used as an exploitability -proxy: the `_HIGH_RISK_CWES` frozenset covers the 14 most commonly weaponised weakness classes. - -**Why additive deltas instead of a new scoring model:** The existing base-score system (by STRIDE -category, target multiplier, protocol, classification) already has calibrated values. Additive -deltas from `RiskContext` layer on top without invalidating historical scores or requiring -re-calibration of the full model. - ---- - -## Decision: CVEService single-pass JSONL loading - -**Choice:** `CVEService._ensure_maps_loaded()` populates both `_cve_to_capec_map` and -`_cve_to_cwe_map` in a single traversal of all JSONL files (~26 files, 1999–2025). - -**Why it matters:** The JSONL corpus is large. A separate pass to build the CWE map (naively -added alongside the existing CAPEC pass) would double I/O and parse time at startup. A single -lazy-loaded pass, gated by `if self._cve_to_capec_map is not None: return`, ensures both maps -are always consistent and the files are opened exactly once. - ---- - -## Decision: `secopstm` CLI entry point via pyproject.toml - -**Choice:** A `secopstm` command is installed by `pip install -e .` via -`[project.scripts] secopstm = "threat_analysis.__main__:main"`. New flags added: -`--output-format {all,html,json,stix}`, `--output-file PATH`, `--stdout`. - -**Why it matters:** CI pipelines need a stable, pip-installable command — not `python -m -threat_analysis` with an implicit working directory. `--stdout` enables direct piping to `jq` -or a SIEM without writing intermediate files. All legacy invocations (`python -m threat_analysis ---server`, `--gui`, `--project`, `--ansible-path`) remain 100% backward-compatible because new -flags all default to the previous behaviour. - ---- - -## Decision: AttackChainAnalyzer — pure graph traversal, no ML - -**Choice:** `threat_analysis/core/attack_chain.py` uses simple directed graph traversal (dataflows -as edges, threats indexed by component name) to identify multi-step attack paths. No NLP, no -ML, no network access. - -**Why it matters:** A common question in threat review is "how do these threats chain together -into a realistic attack scenario?" The attack chain section gives reviewers a pre-computed answer -without additional tooling. It reuses the already-computed `all_threats` list — no extra AI call, -no extra cost. - -**Why not ML/graph DB:** The system topology is small (typically <50 nodes). Simple iteration over -dataflows is O(E × max_threats_per_node). Graph DB (Neo4j) or ML clustering would add heavyweight -dependencies for zero additional accuracy at this scale. - -**Trade-off accepted:** Chains are identified by dataflow adjacency only — threats within the same -boundary but not connected by a modelled dataflow are not chained. This is intentional: if no -dataflow is modelled, no data movement is assumed. - ---- - -## Decision: Trust boundary colors baked into DOT template (B1) - -**Choice:** Trusted boundaries use `color="#2e7d32"` (dark green, solid), untrusted use -`color="#c62828"` (dark red, dashed). These values are in `threat_model.dot.j2` — they affect -the DOT output directly, which means both the Graphviz-generated SVG export and the HTML diagram -render consistent colors. - -**Why baked into DOT and not post-processed in SVG:** The user generates SVG files independently -via `DiagramGenerator.generate_custom_svg_export()`. Any coloring applied only at the HTML stage -would be absent from the exported SVG. Putting the color logic in the DOT template is the only -way to guarantee consistency across all output formats. - -**What was rejected:** A JavaScript overlay that colorizes boundaries in the HTML view only was -considered but rejected for the trust convention specifically (B1). It was however adopted for -the severity heat map (B2) because severity data requires threat processing, which is not -available during SVG generation. - ---- - -## Decision: Severity heat map as JS-only toggle (B2) - -**Choice:** The severity overlay in HTML diagrams is implemented entirely in JavaScript — it reads -`severity_map_json` injected at template render time and toggles SVG `fill` attributes on demand. -The base SVG is never modified. The toggle button hides itself when `severity_map` is empty. - -**Why JS-only:** Severity requires the full threat scoring pipeline to run (`process_threats()` → -`_compute_severity_map()`). In the live editor preview, this pipeline does not run (too slow for -real-time editing). A JS toggle that receives pre-computed data at render time is the only approach -that works in both the live preview (empty map → button hidden) and the exported HTML (data -present → button shown). - -**Why preserve original fills:** The overlay must be reversible so users can compare the -architectural diagram (original colors) against the severity view. The `originalNodeFills` map -saves each node's fill before the overlay is applied and restores it on toggle-off. - ---- - -## Decision: Cross-model RAG analysis via sync wrapper (A2) - -**Choice:** `AIService.generate_rag_threats_sync()` uses `asyncio.run_coroutine_threadsafe()` on -a persistent background event loop (`_get_sync_loop()`) rather than `asyncio.run()` or a new -`ThreadPoolExecutor`. It is called from `ReportGenerator.generate_project_reports()` after all -sub-models are processed and `main_threat_model.sub_models` is populated. - -**Why a persistent loop instead of `asyncio.run()`:** `generate_project_reports()` is called from -a Flask route handler which already runs inside an event loop. `asyncio.run()` would raise -`RuntimeError: cannot run nested event loop`. The persistent background loop pattern (already used -by `generate_markdown_from_prompt_sync`) avoids this issue cleanly. - -**Why after sub-model recursion:** The RAG query needs the full project markdown to generate -cross-boundary threats. The sub-model list is only complete after the recursive pass over -`all_processed_models`. Calling RAG before that would give it only the main model's context, -defeating the purpose of cross-model analysis. - ---- - -## Decision: Boundaries as AI threat targets (A3) - -**Choice:** `SecOpsBoundary` and `Boundary` objects are added to `all_elements` in -`AIService._enrich_with_ai_threats()`. Because boundaries lack `is_public` and `is_authenticated` -attributes (those belong to Servers/Actors), boundary-specific logic was added: `elem_type` is set -to `"Trust Boundary (Trusted|Untrusted)"` and the trust level is injected into the prompt. - -**Why boundaries merit their own AI threats:** Boundaries represent trust zone transitions — -exactly the places where privilege escalation, lateral movement, and data exfiltration paths are -most likely. pytm does not generate threats targeting boundaries directly (only the dataflows -crossing them). Adding boundaries as AI targets surfaces zone-level threats that the rule engine -misses. - -**What was rejected:** Generating boundary threats as a separate post-processing step (e.g. from -the ATT&CK mapping alone) was considered but would not leverage the contextual knowledge of the -LLM about the specific architecture described in the model. - ---- - -## Decision: GDAF top-down approach alongside bottom-up AttackChainAnalyzer - -**Choice:** Implement `GDAFEngine` as a separate, complementary attack path generator that works -top-down from attacker objectives, rather than extending `AttackChainAnalyzer` (which works -bottom-up from discovered threats). - -**Why it matters:** The two engines answer different questions: -- `AttackChainAnalyzer` answers: "Given the threats already found, which ones can be chained - across dataflows?" It starts from the threat inventory and finds adjacencies. -- `GDAFEngine` answers: "If an adversary with these capabilities wanted to reach this target, - what path through the architecture would they take?" It starts from intent and finds paths. - -The bottom-up approach surfaces unexpected chaining of individually low-risk threats. The -top-down approach guarantees that high-value objectives (domain compromise, data exfiltration) -are always evaluated, even if no pytm or AI threats have been generated yet — which matters for -architectures described in the DSL but not yet enriched by AI. - -**Why not replace `AttackChainAnalyzer`:** The two outputs are structurally different. -`AttackChainAnalyzer` chains existing `ExtendedThreat` objects with their STRIDE categories, -MITRE mappings, and severity scores — directly linkable to the threat report. `GDAFEngine` -produces new `AttackScenario` objects with per-hop MITRE technique assignments from a graph -traversal — suited for Attack Flow export and adversary simulation, not for augmenting the -existing threat table. - -**Trade-off accepted:** GDAF requires a context YAML file with objectives and actor profiles. -Without this file, `GDAFEngine.run()` returns an empty list. This is intentional: the engine -is opt-in and the context file is the mechanism for expressing organizational threat intelligence -that cannot be inferred from the architecture model alone. - ---- - -## Decision: Services collectés depuis les dataflows — pas de duplication DSL - -**Choice:** Les protocoles exposés par un asset sont inférés depuis les dataflows adjacents (entrants + sortants) plutôt qu'un champ `services=[]` explicite sur le serveur. - -**Why it matters:** Le DSL modélise déjà les protocoles sur chaque dataflow. Exiger une déclaration redondante `services=[SSH, RDP]` sur le serveur violerait DRY et introduirait des incohérences (services déclarés mais aucun dataflow correspondant). Le collecteur de `_build_graph()` fait un agrégat automatique au moment de la construction du graphe. - -**Trade-off accepted:** Les services locaux sans dataflow modélisé (ex. un port ouvert mais jamais utilisé dans le modèle) ne sont pas visibles. Cette limitation est intentionnelle : si un flux n'est pas modélisé, aucun chemin d'attaque ne l'emprunte de toute façon. - ---- - -## Decision: BOM (Bill of Materials) comme enrichissement optionnel hors-DSL - -**Choice:** Les données d'inventaire (version OS, CVEs, patch level, detection level) sont stockées dans des fichiers YAML séparés par asset (`BOM/{asset_name}.yaml`), pas dans le DSL. - -**Why it matters:** Le DSL décrit l'architecture logique (composants, flux, confiance). L'inventaire technique (versions, CVEs) est une donnée opérationnelle qui change fréquemment, doit pouvoir être générée depuis des outils tiers (Qualys, Tenable, CMDB), et ne doit pas polluer la lisibilité du modèle d'architecture. - -**Why YAML and not inline DSL:** Un export de scanner de vulnérabilités peut directement alimenter le répertoire `BOM/`. Le format YAML est compatible avec les exports Ansible inventory, Tenable, Qualys. L'intégration directe dans le DSL nécessiterait une migration de tous les modèles existants. - -**Auto-discovery:** `BOMLoader` cherche `{model_parent}/BOM/` automatiquement si `bom_directory` n'est pas déclaré dans `## Context`. Pas de configuration requise pour commencer. - ---- - -## Decision: `traversal_difficulty` sur les boundaries — coût de traversée dans GDAF - -**Choice:** Un attribut `traversal_difficulty=low|medium|high` sur les boundaries contrôle un bonus de `hop_weight` : `low` → +0.3 (zone facile à traverser = chemin plus exploitable), `medium` → +0.1, `high` → +0.0 (attaquant très capable requis = chemin moins probable mais toujours évalué). - -**Why it matters:** Le trust binaire (`isTrusted=true/false`) ne distingue pas une DMZ d'un réseau OT airgappé. Un attaquant peut traverser les deux, mais la probabilité et les techniques nécessaires sont radicalement différentes. Le bonus de hop_weight reflète l'accessibilité relative. - -**Why NOT reduce score for high difficulty:** Un chemin traversant un réseau OT (traversal_difficulty=high) doit toujours apparaître dans les résultats — c'est précisément ce chemin qu'un analyste doit évaluer. La difficulté réduit légèrement son score relatif mais ne l'élimine pas. - ---- - -## Decision: Sub-model drill-down — child IS the parent at higher detail - -**Choice:** A server in a threat model can reference a child model via `submodel=./path/to/model.md`. -The child is not a separate dependency or a called service — it IS the same component described -at a finer granularity. The parent diagram links to the child; the child diagram shows ghost -nodes for external connections from the parent. - -**Why it matters:** Large architectures become unreadable when everything is in one flat diagram. -Drill-down allows a threat model to be decomposed: the top-level model shows the overall -architecture with summarized components, and each component that merits deeper analysis gets its -own model file with the same full DSL support (actors, servers, dataflows, STRIDE analysis, AI -enrichment). - -**Why ghost nodes in the child diagram:** When a security reviewer drills into a sub-component, -they need to know what external systems communicate with it — this context is essential for -identifying trust boundary threats and data flow risks. Ghost nodes provide that context without -requiring the child model to re-declare all parent-level connections, which would create -maintenance duplication. - -**Why bridging edges in GDAF:** An attacker who compromises a server does not stop at that node -— they move into the server's internals. Without bridging edges, GDAF would treat the parent -server as a terminal node and miss all internal paths. Bridging reuses the already-parsed -sub-model graph, adding only two passes of edge injection (entry + exit bridges), at essentially -zero extra cost. - -**What was rejected:** A separate "project dependencies" section in the DSL (defining which -models communicate with each other) was considered but would require maintaining two separate -representations of the same relationship — the server definition and a separate link declaration. -The `submodel=` inline keyword keeps the link co-located with the server definition. - ---- - -## Decision: RAG init parallel to AI connection check — `run_in_executor` - -**Choice:** `AIService.init_ai()` starts `RAGThreatGenerator()` in `loop.run_in_executor(None, ...)` -immediately, before calling `provider.check_connection()`. The asyncio Future is awaited only -after the connection check completes. - -**Why it matters:** The two operations are completely independent: -- `check_connection()` is bottlenecked by `import litellm` (~64s on WSL2 cold cache, network - call to fetch model pricing/metadata). -- `RAGThreatGenerator.__init__()` is bottlenecked by `import langchain_chroma` (~26s cold). -By overlapping them, the 26s chroma import is hidden inside the 64s litellm import. -Net saving: ~26s on every cold start where RAG is enabled. - -**Why `run_in_executor` and not `asyncio.create_task`:** `RAGThreatGenerator.__init__()` is -synchronous and blocking (disk I/O, model loading). `create_task` would run in the same thread -and block the event loop. `run_in_executor` offloads to the default `ThreadPoolExecutor`, -keeping the event loop responsive. - -**Fallback:** If the pre-warm raises (vector store missing, config error), `init_ai()` retries -synchronously once before setting `rag_generator = None`. - -**What was rejected:** Starting RAG initialization from `ThreatModelService.__init__()` (before -`init_ai()` is even called) was considered but would require passing the AI config status back -and creates a race condition if the AI goes offline between warmup and use. - ---- - -## Decision: BOM known_cves augment CVEService in scoring pipeline - -**Choice:** `ReportGenerator._get_bom_loader(threat_model)` resolves a `BOMLoader` from the -model's `_model_file_path` (auto-discovers `{parent}/BOM/`) or `context_config['bom_directory']`. -In `_get_all_threats_with_mitre_info()`, BOM `known_cves` are appended to the list returned by -`cve_service.get_cves_for_equipment()` before the CAPEC intersection check — for both pytm and -AI-element threat scoring loops. - -**Why it matters:** `CVEService.get_cves_for_equipment()` only knows CVEs declared in -`cve_definitions.yml` (YAML definitions keyed by equipment name). BOM files carry the -authoritative list of CVEs for each deployed asset (from vulnerability scanners, patch managers, -or SBOM pipelines). Combining both sources ensures that CVE-based VOC scoring reflects reality -rather than only the manually maintained YAML. - -**What was rejected:** Feeding BOM CVEs directly into `CVEService` as a secondary lookup would -require mutating a shared service object or adding BOM awareness to `CVEService` itself — coupling -it to the BOM infrastructure. Augmenting the call site list is purely additive and keeps -`CVEService` focused on its existing JSONL/YAML data. - ---- - -## Decision: initial_model_file_path global for Flask server BOM/context auto-discovery - -**Choice:** `server.py` adds an `initial_model_file_path: Optional[str]` module-level global, -set in `run_server()` when a model file is successfully loaded. It is passed as `model_file_path` -to every `export_*_logic()` route call, which propagates it through `ThreatModelService` → -`ExportService` → `create_threat_model()` → `ThreatModel._model_file_path`. With `_model_file_path` -set, `ExportService._resolve_gdaf_context()` and `_resolve_bom_directory()` can auto-discover -the model's `context/` and `BOM/` sibling directories. - -**Why a global:** The Flask routes are stateless functions with no access to the `run_server()` -scope. The alternative — injecting the path through `ThreatModelService.__init__()` — would -require a new constructor parameter and rebinding of the service on model switch, which is more -invasive. A module-level global is consistent with the existing `initial_markdown_content` and -`initial_project_path` pattern already used by the server. - -**Scope:** When the server starts without a model file (empty editor mode), `initial_model_file_path` -stays `None` and exports proceed without BOM/context auto-discovery — matching the previous -behaviour exactly. diff --git a/developer.md b/developer.md index 09a5067..6c49b7f 100644 --- a/developer.md +++ b/developer.md @@ -48,13 +48,13 @@ This will download the necessary file to `threat_analysis/external_data/`. ## Future Development: Infrastructure as Code (IaC) Integration -This section outlines the planned integration of Infrastructure as Code (IaC) tools, starting with Ansible, into the STRIDE Threat Analysis Framework. The goal is to enhance the "Threat Modeling as Code" philosophy by enabling automated threat model generation and analysis directly from IaC configurations. +This section outlines the planned integration of Infrastructure as Code (IaC) tools, starting with Ansible, into the STRIDE Threat Analysis Framework. The goal is to enhance the "Threat Modeling as Code" philosophy by enabling automated system model generation and threat analysis directly from IaC configurations. ### Vision -The core idea is to bridge the gap between infrastructure definitions (e.g., Ansible playbooks, Terraform configurations) and the threat model. By parsing IaC files, we can automatically infer system components, data flows, and configurations, and then use this information to: +The core idea is to bridge the gap between infrastructure definitions (e.g., Ansible playbooks, Terraform configurations) and the system model. By parsing IaC files, we can automatically infer system components, data flows, and configurations, and then use this information to: -1. **Dynamically Generate/Update Threat Models**: Automatically create or update `threatModel_Template/threat_model.md` based on the deployed infrastructure. +1. **Dynamically Generate/Update System Models**: Automatically create or update `threatModel_Template/threat_model.md` based on the deployed infrastructure. 2. **Automate Threat Detection**: Identify potential STRIDE threats and map them to MITRE ATT&CK techniques based on the IaC configuration (e.g., open ports, insecure configurations). 3. **Detect New Threats**: Compare threat analysis reports before and after IaC changes to highlight newly introduced threats or changes in the attack surface. @@ -65,7 +65,7 @@ This initiative directly supports the "Threat Modeling as Code" workflow detaile The integration will involve the following steps: 1. **IaC Parser Development**: Create dedicated modules to read and interpret IaC configuration files (e.g., Ansible playbooks, Terraform `.tf` files). These parsers will extract relevant infrastructure details (e.g., servers, network configurations, deployed services). -2. **Mapping to Threat Model DSL**: The extracted IaC data will be translated into the components of our Markdown-based Threat Model DSL (Boundaries, Actors, Servers, Data, Dataflows). +2. **Mapping to System Model DSL**: The extracted IaC data will be translated into the components of our Markdown-based System Model DSL (Boundaries, Actors, Servers, Data, Dataflows). 3. **Dynamic `threatModel_Template/threat_model.md` Generation**: The translated data will be used to generate or update the `threatModel_Template/threat_model.md` file programmatically. 4. **Automated Analysis & Reporting**: The existing threat analysis framework (`threat_analysis/__main__.py`) will then process the generated `threatModel_Template/threat_model.md` to perform STRIDE analysis, MITRE mapping, and report generation. 5. **Change Detection (Future)**: Mechanisms will be explored to compare analysis results over time, identifying new or mitigated threats resulting from IaC changes. @@ -75,12 +75,12 @@ The integration will involve the following steps: We will begin by integrating with Ansible. This phase will focus on: 1. **Ansible Playbook/Inventory Parser**: A module will be developed to parse Ansible playbooks and inventory files to identify hosts, roles, tasks, and network configurations. -2. **Mapping Logic**: Define clear rules for how Ansible constructs (e.g., hosts, tasks that open ports, roles that install services) map to threat model elements. -3. **CLI Integration**: A new command-line option (e.g., `--from-ansible `) will be added to `threat_analysis/__main__.py` to trigger the Ansible-based threat model generation. +2. **Mapping Logic**: Define clear rules for how Ansible constructs (e.g., hosts, tasks that open ports, roles that install services) map to system model elements. +3. **CLI Integration**: A new command-line option (e.g., `--from-ansible `) will be added to `threat_analysis/__main__.py` to trigger the Ansible-based system model generation. #### Test Playbook -To facilitate development and testing, a sample Ansible playbook will be stored within the project. This playbook will define a simple infrastructure that can be used to validate the parsing and threat model generation logic. +To facilitate development and testing, a sample Ansible playbook will be stored within the project. This playbook will define a simple infrastructure that can be used to validate the parsing and system model generation logic. **Location**: `tests/ansible_playbooks/simple_web_server.yml` (or similar) diff --git a/docs/Roadmap.md b/docs/Roadmap.md index b1c4784..7b57801 100644 --- a/docs/Roadmap.md +++ b/docs/Roadmap.md @@ -8,12 +8,12 @@ This document outlines the development roadmap for the SecOpsTM framework, track This section highlights the core capabilities that are already integrated into the framework. -- **Hierarchical & Interactive Threat Models**: Decompose a large system into multiple, linked sub-models with navigable HTML diagrams. -- **Advanced Threat Model Validation**: A dedicated validation module (`model_validator.py`) checks for consistency, unique names, and valid references within the model before analysis. +- **Hierarchical & Interactive System Models**: Decompose a large system into multiple, linked sub-models with navigable HTML diagrams. +- **Advanced System Model Validation**: A dedicated validation module (`model_validator.py`) checks for consistency, unique names, and valid references within the model before analysis. - **Automated and Enriched Mitigation Suggestions**: The framework proposes context-aware mitigations based on recognized frameworks (OWASP ASVS, NIST, CIS Controls) for each identified MITRE ATT&CK technique. -- **Ansible IaC Integration**: Automatically generate a threat model directly from Ansible playbooks and inventories, using embedded metadata for a rich, accurate representation. -- **Pre-defined Architecture Templates**: A library of pre-built threat models for common architectural patterns is available in the `threatModel_Template/` directory to accelerate initial setup. -- **Web-Based GUI**: An interactive web interface for real-time editing of threat models with a live-updating diagram preview. +- **Ansible IaC Integration**: Automatically generate a system model directly from Ansible playbooks and inventories, using embedded metadata for a rich, accurate representation. +- **Pre-defined Architecture Templates**: A library of pre-built system models for common architectural patterns is available in the `threatModel_Template/` directory to speed up initial setup. +- **Web-Based GUI**: An interactive web interface for real-time editing of system models with a live-updating diagram preview. - **MITRE ATT&CK Navigator Export**: Generate JSON layer files for visualization and analysis in the MITRE ATT&CK Navigator. - **STIX 2.1 Reporting**: Export threat intelligence data in the standardized STIX 2.1 format for interoperability with other security tools. @@ -24,16 +24,16 @@ This section highlights the core capabilities that are already integrated into t This section outlines the strategic vision and planned features for future releases. - **Integration with Vulnerability Databases (CVE)**: Link identified MITRE ATT&CK techniques to known CVEs or common vulnerabilities (e.g., OWASP Top 10) to provide even deeper context. -- **Attack Path Enumeration and Simulation**: Implement algorithms to automatically identify and visualize potential attack paths through the threat model. +- **Attack Path Enumeration and Simulation**: Implement algorithms to automatically identify and visualize potential attack paths through the system model. - **Integration with Security Orchestration, Automation, and Response (SOAR) Platforms**: Develop connectors to push threat intelligence and mitigation recommendations directly into SOAR platforms. - **Machine Learning-Enhanced Threat Identification**: Train a machine learning model to predict potential threats that may not be covered by existing rules. - **Expanded IaC Tool Support (Terraform, CloudFormation)**: Create new plugins for other popular IaC tools. - **Enhanced User Interface (UI) and User Experience (UX)**: Redesign the web GUI to be more intuitive, with features like drag-and-drop model creation and a visual rule editor. - **Compliance Mapping (NIST, PCI-DSS, etc.)**: Map identified threats and suggested mitigations to specific controls in major compliance frameworks. - **Threat Intelligence Feed Integration**: Integrate with external threat intelligence feeds (e.g., from MISP, Anomali) to automatically update the threat model with the latest real-world attack techniques. -- **Collaborative Real-Time Editing**: Enhance the web GUI to support real-time, multi-user collaborative editing of threat models. +- **Collaborative Real-Time Editing**: Enhance the web GUI to support real-time, multi-user collaborative editing of system models. - **Risk Quantification and Financial Impact Analysis**: Integrate with risk quantification models (e.g., FAIR) to estimate the potential financial impact of identified threats. -- **Automated Attack Path Analysis**: Generate potential attack paths by chaining together identified MITRE ATT&CK techniques across all nodes in the threat model. This will involve enriching the CAPEC to ATT&CK mappings to build a comprehensive attack graph. +- **Automated Attack Path Analysis**: Generate potential attack paths by chaining together identified MITRE ATT&CK techniques across all nodes in the system model. This will involve enriching the CAPEC to ATT&CK mappings to build a full attack graph. - **Advanced Severity Multiplier Calculation**: Implement a more granular and qualitative method for calculating the severity multiplier of each asset. This will be based on a rating of four key security criteria: * **Confidentiality**: The impact of unauthorized disclosure of data. * **Integrity**: The impact of unauthorized modification of data. @@ -48,8 +48,8 @@ This section outlines the strategic vision and planned features for future relea 2. **Dataset Creation**: Build a high-quality dataset for fine-tuning, including: - The complete knowledge bases of **MITRE ATT&CK, CAPEC, and D3FEND**. - Thousands of public **threat intelligence reports** (from CISA, Mandiant, etc.) and **CVE** descriptions. - - All existing **threat models (`.md` files)** from the project, so the LLM can learn the DSL and architectural patterns. - - Pairs of "input threat model" -> "output list of threats and attack paths" generated by the current tool for supervised learning. + - All existing **system models (`.md` files)** from the project, so the LLM can learn the DSL and architectural patterns. + - Pairs of "input system model" -> "output list of threats and attack paths" generated by the current tool for supervised learning. 3. **Training**: Fine-tune the base model on this dataset to specialize it in the "thinking" of a security expert. * **Usage Phase (Inference)**: 1. **Hybrid Approach**: The current rule engine remains the baseline for fast, deterministic analysis. The LLM acts as a second pass. diff --git a/docs/claude-code-skill.md b/docs/claude-code-skill.md index 8bfb032..bccada8 100644 --- a/docs/claude-code-skill.md +++ b/docs/claude-code-skill.md @@ -1,4 +1,4 @@ -# Claude Code Skill — AI-Assisted Threat Model Generation +# Claude Code Skill — AI-Assisted System Model Generation SecOpsTM ships a **Claude Code skill** that guides an AI agent through the full threat modeling workflow: information gathering, DSL file generation, BOM/context YAML generation, and validation. @@ -159,7 +159,7 @@ after the initial description. ### Step 6 — model.md Generation -Claude writes the threat model DSL file. All output is in English regardless of the language +Claude writes the system model DSL file. All output is in English regardless of the language used during information gathering. Key rules enforced automatically: @@ -230,7 +230,7 @@ Target: **0 failures**. Claude fixes any failures and re-runs before reporting c ``` {system-name}/ ├── .tm-raw/ # scratch files (not committed) -├── model.md # threat model DSL +├── model.md # system model DSL ├── context/ │ └── {system-name}_context.yaml # GDAF objectives, actors, risk criteria ├── BOM/ diff --git a/docs/customizing_prompts.md b/docs/customizing_prompts.md index 99649ec..a3eac9f 100644 --- a/docs/customizing_prompts.md +++ b/docs/customizing_prompts.md @@ -51,7 +51,7 @@ converts it into the SecOpsTM Markdown DSL. dsl_generation: system: | You are an expert cybersecurity architect for a payment-processing organisation. - Generate threat models in the SecOpsTM Markdown DSL. + Generate system models in the SecOpsTM Markdown DSL. Naming conventions: - Boundaries: use zone names (e.g. "DMZ", "CardholderDataEnvironment", "Management") @@ -311,7 +311,7 @@ rag: ## Threat Intelligence {user_threat_intelligence} - ## Architecture (Threat Model) + ## Architecture (System Model) {threat_model_markdown} ## Retrieved Security Knowledge diff --git a/docs/data_collection_guide.md b/docs/data_collection_guide.md index fede11d..2b315c7 100644 --- a/docs/data_collection_guide.md +++ b/docs/data_collection_guide.md @@ -1,11 +1,11 @@ # Data Collection Guide for Threat Modeling -This guide tells you exactly what information to gather before building a SecOpsTM threat model. +This guide tells you exactly what information to gather before building a SecOpsTM system model. It is written for security architects and DevOps engineers who are conducting or commissioning a threat modeling session. You can hand this document to a system owner and expect to receive back everything the tool needs. -A well-filled threat model produces: a STRIDE threat report with MITRE ATT&CK mappings, a +A well-filled system model produces: a STRIDE threat report with MITRE ATT&CK mappings, a severity-colored architecture diagram, goal-driven attack scenarios (GDAF), and exports to STIX 2.1 and ATT&CK Navigator. @@ -26,7 +26,7 @@ SecOpsTM models a system by describing: 6. **What the system inventory looks like** (BOM) 7. **What attackers want** and **who they are** (GDAF context) -Each piece of information maps directly to a section in the threat model file. Missing information +Each piece of information maps directly to a section in the system model file. Missing information produces lower-quality threat coverage — for example, a server without CIA ratings gets no GDAF path scoring, and a dataflow without `is_encrypted` gets a default of `False` which adds threat noise. @@ -557,7 +557,7 @@ build the DSL file. **Minimal viable model template to fill in:** ```markdown -# Threat Model: [SYSTEM NAME] +# System Model: [SYSTEM NAME] ## Description [2–5 sentences describing the system, its purpose, user count, deployment, compliance scope] @@ -627,7 +627,7 @@ A small e-commerce application with a public API, an internal database, and an a **Resulting model.md:** ```markdown -# Threat Model: E-Commerce Platform +# System Model: E-Commerce Platform ## Description Customer-facing e-commerce API built on Node.js and PostgreSQL. Approximately 50,000 customers diff --git a/docs/defining_threat_models.md b/docs/defining_threat_models.md index 41cfeed..8d663aa 100644 --- a/docs/defining_threat_models.md +++ b/docs/defining_threat_models.md @@ -2,7 +2,7 @@ This document is the authoritative reference for the SecOpsTM Markdown Domain-Specific Language (DSL). It covers every section, every attribute, and every accepted value. After reading this, you -can write any threat model from scratch without consulting the source code. +can write any system model from scratch without consulting the source code. **Target audience:** DevOps engineers, security architects, threat modelers. @@ -10,7 +10,8 @@ can write any threat model from scratch without consulting the source code. ## Overview -A SecOpsTM threat model is a Markdown file. The file is parsed in three passes: +A SecOpsTM system model is a Markdown file describing the architecture to analyze; SecOpsTM +generates the threat model from it. The file is parsed in three passes: 1. **Pass 0** — `## Context` (standalone key-value block, model-level settings) 2. **First pass** — `## Boundaries`, `## Actors`, `## Servers`, `## Data` (element definitions) @@ -36,7 +37,7 @@ This order matters: dataflows can only reference elements that were defined in t ## File Structure ```markdown -# Threat Model: My System Name +# System Model: My System Name ## Description Free-form text describing the system under analysis. @@ -70,8 +71,9 @@ bom_directory = BOM - **Attack Name**: {"tactics": ["Lateral Movement"], "techniques": [{"id": "T1021.002", "name": "SMB/Windows Admin Shares"}]} ``` -The title (`# Threat Model: ...`) is used as the model name in reports and diagrams. Everything -else is controlled by the sections described below. +The title (`# System Model: ...`) is used as the model name in reports and diagrams. Everything +else is controlled by the sections described below. The legacy `# Threat Model: ...` form still +works (with a deprecation warning) — always write new models with `# System Model:`. --- @@ -666,7 +668,7 @@ The recommended structure for a project-mode model (with GDAF context, BOM, and ``` My_System/ - model.md # Main threat model DSL file + model.md # Main system model DSL file context/ # GDAF context YAML (auto-discovered) context.yaml BOM/ # Per-asset Bill of Materials (auto-discovered) @@ -994,7 +996,7 @@ features. It models a large enterprise with 9 security zones, Active Directory, OT/SCADA. ```markdown -# Threat Model: On-Prem Enterprise Network +# System Model: On-Prem Enterprise Network ## Description Large multinational manufacturing enterprise with on-premises Active Directory infrastructure, diff --git a/docs/dsl-syntax.md b/docs/dsl-syntax.md index da28666..e80a79f 100644 --- a/docs/dsl-syntax.md +++ b/docs/dsl-syntax.md @@ -18,7 +18,7 @@ ## Overview -SecOpsTM threat model is a Markdown file parsed in three passes: +A SecOpsTM system model is a Markdown file parsed in three passes: 1. **Pass 0** — `## Context` (model-level configuration) 2. **First pass** — `## Boundaries`, `## Actors`, `## Servers`, `## Data` (element definitions) @@ -40,14 +40,16 @@ SecOpsTM threat model is a Markdown file parsed in three passes: **⚠️ CRITICAL**: model.md must use the following title format: ```markdown -# Threat Model: {System Name} +# System Model: {System Name} ``` **Format requirements**: -- Must start with `# Threat Model: ` (English colon `:` + space) +- Must start with `# System Model: ` (English colon `:` + space) - System name follows after the space -- **Correct**: `# Threat Model: DJI Power Inspection Drone System` -- **Incorrect**: `# DJI Power Inspection Drone System - Threat Model` +- **Correct**: `# System Model: DJI Power Inspection Drone System` +- **Incorrect**: `# DJI Power Inspection Drone System - System Model` +- The legacy `# Threat Model: {System Name}` form still works (with a deprecation warning) — + always write new models with `# System Model:`. **Impact**: Incorrect title format causes SecOpsTM to fail generating report filenames correctly. @@ -163,7 +165,7 @@ Actors represent people, external systems, or roles interacting with the system, ## Section: Servers -Servers represent assets, components, or systems in the threat model. +Servers represent assets, components, or systems in the system model. ### Server Attributes @@ -482,7 +484,7 @@ Use `config/asset_types_community.yaml` and `config/protocols_community.yaml` to **⚠️ CRITICAL FORMAT REQUIREMENT**: SecOpsTM parser requires EXACT DSL syntax. DO NOT use custom Markdown formats like `### Boundary:` or `**Name**:`. Use ONLY the list format shown below. ```markdown -# Threat Model: {System English Name} +# System Model: {System English Name} ## Description {High-level system description in English} diff --git a/docs/examples.md b/docs/examples.md index 7bdf3af..a1e4b1a 100644 --- a/docs/examples.md +++ b/docs/examples.md @@ -1,9 +1,9 @@ -# Threat Model DSL & Examples +# System Model DSL & Examples -This framework leverages PyTM's modeling primitives. For a comprehensive reference of all available attributes and their usage, please refer to the [PyTM Documentation](https://owasp.github.io/PyTM/Threat-Model-as-Code/). Note that this framework may extend PyTM with additional attributes or behaviors. +This framework uses PyTM's modeling primitives. For a full reference of all available attributes and their usage, see the [PyTM Documentation](https://owasp.github.io/PyTM/Threat-Model-as-Code/). This framework may extend PyTM with additional attributes or behaviors. ```markdown -# Threat Model: Advanced DMZ Architecture +# System Model: Advanced DMZ Architecture ## Description A network with a DMZ, external/internal firewalls, and a command zone. The goal is to identify STRIDE threats and map them to MITRE ATT&CK. @@ -39,7 +39,7 @@ A network with a DMZ, external/internal firewalls, and a command zone. The goal ## Common Attributes -Several elements in the threat model DSL support common attributes to enhance their representation and behavior: +Several elements in the system model DSL support common attributes that affect their representation and behavior: - **`color`**: Specifies the color of the element in the generated diagrams. - **Applies to**: Boundaries, Actors, Dataflows, Protocol Styles. @@ -59,7 +59,7 @@ Several elements in the threat model DSL support common attributes to enhance th ## Protocol Styles and Legends -To ensure that protocols are correctly styled in diagrams and appear in the legend, you must define them in the `## Protocol Styles` section of your threat model. The system **intentionally does not** assign default colors to new protocols. This gives you full control over the final visualization. +To ensure that protocols are correctly styled in diagrams and appear in the legend, you must define them in the `## Protocol Styles` section of your system model. The system **intentionally does not** assign default colors to new protocols. This gives you full control over the final visualization. **How it works:** @@ -212,6 +212,6 @@ After running the analysis, you will find a timestamped folder in `output/` (e.g ### Using Pre-defined Templates -To accelerate the creation of new threat models, the framework includes a set of pre-defined templates for common architectures. You can load these templates directly from the web interface. +To speed up the creation of new system models, the framework includes a set of pre-defined templates for common architectures. You can load these templates directly from the web interface. ![Loading a template](../../output/example/gui_example.gif) diff --git a/docs/extensibility.md b/docs/extensibility.md index 111231d..dac075f 100644 --- a/docs/extensibility.md +++ b/docs/extensibility.md @@ -29,7 +29,7 @@ ## PyTM-Based Extensions & Evolutions -You can leverage and extend all PyTM features, including: +You can use and extend all PyTM features, including: - **Custom Threat Patterns**: Define new threat types and detection logic. - **Advanced Dataflow Modeling**: Use PyTM's dataflow and element types. diff --git a/docs/features.md b/docs/features.md index 2d964a8..de1ef7c 100644 --- a/docs/features.md +++ b/docs/features.md @@ -53,18 +53,22 @@ GDAF is a top-down attack scenario generator that works from attacker objectives - **GDAF in the HTML report**: GDAF scenarios appear in the HTML threat report as a collapsible `
` accordion (closed by default), placed after the Attack Chain Analysis section. The table shows Risk level (CRITICAL/HIGH/MEDIUM/LOW badge), Objective, Actor and sophistication, Attack Path (`A → B → C`), Score, Hop count, and Detection coverage. Each row is expandable to show per-hop details: node name, asset type, protocol, cleartext/no-auth flags, and assigned MITRE ATT&CK techniques. - **Project mode**: In multi-model projects, the attack graph spans all sub-models. Servers with `submodel=` references get bridging edges so paths can traverse into component internals. - **Fully offline**: Only reads `enterprise-attack.json` from disk — no network calls. +- **Red/Blue adversarial debate** (opt-in, `debate.enabled` in `ai_config.yaml`, on by default): the top GDAF scenarios are stress-tested by two LLM personas — Red attempts to advance the attack, Blue counters with SIEM/EDR/IDS controls and detection gaps, grounded only in facts already in the model. The outcome adjusts each scenario's score/risk_level (never invents new threats), and the `.afb` files are re-written to match. -See [docs/gdaf.md](gdaf.md) for the complete reference including the context YAML schema, scoring algorithm, and asset type table. +See [docs/gdaf.md](gdaf.md) for the complete reference including the context YAML schema, scoring algorithm, debate mechanics, and asset type table. ## Reporting & Export - **HTML report**: Integrated threat statistics, STRIDE/MITRE mapping, D3FEND mitigations, severity breakdown, source tagging (`pytm` / `AI` / `LLM`), risk signals (`CVE`, `CWE⚠`, `NET`, `D3F`), executive summary with KPIs + top-5 risks, interactive severity filter (CRITICAL/HIGH/MEDIUM/LOW), risk matrix 5×5. Screenshots: [risk matrix](../../output/example/risk_matrix.jpg) · [top-5 threats](../../output/example/top5_threat.jpg) · [threat graph](../../output/example/threat_graph.jpg) · [CISO briefing](../../output/example/ciso_risk_briefing.jpg) +- **SOC Analyst detection pass**: Per-threat Sigma / Splunk SPL / KQL rule suggestions and IOCs, generated by a dedicated LLM persona and rendered in the report's "SOC Analysis" section. Requires AI enabled; skipped silently offline. +- **CISO Triage**: Board-level risk briefing summarizing the highest-priority threats (including GDAF scenarios and debate outcomes when available), generated by a dedicated LLM persona. - **Executive View toggle**: A single click hides all technical sections (Attack Chain Analysis, GDAF Scenarios, ATT&CK ID Validation, Threat Graph, Severity Calculation, Legend) for clean management presentations. Implemented as a pure CSS `.exec-view` class toggle — no layout reflow. - **Copy-as-ticket button**: Each top-5 threat row has a "Copy ticket" button that copies a GitHub Issue–formatted markdown block to clipboard (title, severity, STRIDE category, target, description, action checklist). Supports Clipboard API with `execCommand` fallback. - **Collapsible report sections**: `📋 Model Completeness` and `📖 Severity Calculation Explained` are wrapped in `
` — collapsed by default to reduce visual noise, summary line visible when closed. - **⛓️ Attack Chain Analysis**: Dedicated section in the HTML report identifying multi-step attack paths that chain threats across dataflows. Each chain shows entry point, pivot component, attack scores, and CRITICAL/HIGH/MEDIUM/LOW severity label. - **GDAF scenarios accordion**: GDAF attack scenarios are embedded in the HTML report as an expandable section between Attack Chain Analysis and Severity Calculation Explained. Only shown when GDAF scenarios have been generated (requires a valid context YAML with `attack_objectives` and `threat_actors`). +- **🔍 Automatically Discovered Attack Paths**: Best (highest-severity) attack path per STRIDE category, found by following each threat's MITRE ATT&CK tactic progression through the architecture — no GDAF context needed, works from the threat data alone (pytm + AI + LLM sources together). With AI enabled, each path also gets a short grounded narrative and business-impact summary (opt-out via `attack_flows.include_narrative` in `ai_config.yaml`); the persona is instructed to never cite an ID (ATT&CK/CVE/CAPEC/D3FEND) — the exact IDs are already shown deterministically next to each hop, and any response that emits one anyway is discarded rather than trusted. - **Report diff page** (`/diff`): Web page served at `/diff` that accepts two JSON exports (paste or file upload) and displays a visual comparison — new threats `[+]`, resolved threats `[-]`, severity changes `[~]` — with counts by category at the top. Also available via CLI: `secopstm --diff old_report.json new_report.json`. - **Versioned JSON export** (`schema_version: "1.0"`): Stable structure for SIEM, dashboards, and ticketing tools. Schema defined at `threat_analysis/schemas/v1/threat_model_report.schema.json`. Threats carry stable IDs (`T-0001`). - **JSON export REST API** (`POST /api/export_json`): Returns the versioned JSON report directly from the API without generating a ZIP bundle. Accepts `{"markdown_content": "..."}` and returns the schema-validated report. @@ -78,7 +82,7 @@ See [docs/gdaf.md](gdaf.md) for the complete reference including the context YAM ## GitHub Action -SecOpsTM ships as an **official GitHub Action** (`action.yml`) for threat-model-as-code CI/CD: +SecOpsTM ships as an **official GitHub Action** (`action.yml`) for threat-modeling-as-code CI/CD: ```yaml - uses: your-org/secopstm@v1 @@ -120,7 +124,7 @@ See `.github/workflows/threat-model.yml` for the example workflow. - **Graphical Editor**: Visual drag-and-drop canvas for building models without writing Markdown. - **Reports are fully self-contained** and work offline. -## Threat Model Templates +## System Model Templates Ready-to-use DSL templates in `threatModel_Template/`: diff --git a/docs/gdaf.md b/docs/gdaf.md index 1c88db8..05b7fdd 100644 --- a/docs/gdaf.md +++ b/docs/gdaf.md @@ -23,7 +23,7 @@ GDAF reads a YAML file that defines `attack_objectives`, `threat_actors`, and `r ### Step 2 — Build a directed graph -`GDAFEngine._build_graph()` creates a node for every actor and server defined in the threat model. Dataflows become directed edges carrying metadata (protocol, encryption, authentication). In project mode, the graph is unified across the main model and all sub-models. +`GDAFEngine._build_graph()` creates a node for every actor and server defined in the system model. Dataflows become directed edges carrying metadata (protocol, encryption, authentication). In project mode, the graph is unified across the main model and all sub-models. ### Step 3 — Identify entry points @@ -94,6 +94,45 @@ Clicking a row expands an inline detail panel showing one entry per hop: --- +## Red/Blue Adversarial Debate + +An opt-in pass (`config/ai_config.yaml → debate.enabled: true`, on by default) that runs after +GDAF and before the HTML report is written. One configured LLM provider plays two personas — +**Red** (attack) and **Blue** (defence) — that argue over whether the top-scoring GDAF scenarios +actually hold up, using only facts already present in the model (dataflow auth/encryption, hop +protocols, BOM `known_cves`). No new lookups, no invented evidence — everything either persona +cites must come from that grounding block. + +**Scope:** only the top `debate.top_n` scenarios (default 5) with an initial viability +≥ `debate.min_viability_threshold` (default 0.5) are debated — `initial_viability = (path_score / +CRITICAL_threshold) × (1 − 0.5 × detection_coverage)`. Debating every scenario would be slow and +wasteful; the point is to stress-test the ones already flagged as most dangerous. + +**Per scenario**, up to `debate.max_rounds` rounds (default 3) run: +1. **Red** proposes techniques to advance the attack, citing grounding evidence; if Blue blocked + the previous approach, Red must try a genuinely different technique or record it under + `failed_alternatives`. +2. **Blue** decides which of Red's techniques are blocked/detected, naming the control family + (SIEM / EDR / IDS), and lists `detection_gaps` for anything it cannot cover. +3. Viability for the round is `red_viability − 0.15 × len(techniques_blocked)`. If the change + from the previous round is below `debate.viability_delta_threshold` (default 0.1), the debate + is considered converged and stops early. + +**Outcome:** the final viability is mapped to a `debate_factor` in +`[debate.debate_factor_min, debate.debate_factor_max]` (default `[0.5, 1.5]`) and applied +directly to the scenario: `scenario.path_score = round(path_score × debate_factor, 2)`, then +`risk_level` is recomputed from the same CRITICAL/HIGH/MEDIUM/LOW thresholds GDAF itself uses. +A scenario Blue fully shuts down score *lower* than GDAF's original estimate; one Red can push +through despite Blue's controls scores *higher*. The debate never creates new threats or +scenarios — it only adjusts the score/risk_level of the ones GDAF already produced, and the +`.afb` Attack Flow files are re-written afterward so they stay consistent with the debated score. + +**Fully offline-compatible:** if no AI provider is configured or reachable, the debate pass is +skipped entirely and GDAF's original scores stand unchanged — this is additive enrichment, never +a requirement for GDAF to produce output. + +--- + ## Context YAML Schema ### Linking a context file to a model diff --git a/docs/getting_started.md b/docs/getting_started.md index ba8ebfb..0fd4591 100644 --- a/docs/getting_started.md +++ b/docs/getting_started.md @@ -30,7 +30,7 @@ docker run -p 5000:5000 -v $(pwd)/output:/app/output \ # With LLM inference — default provider is NVIDIA NIM (free tier available) # See https://build.nvidia.com/meta/llama-3_3-70b-instruct for a free API key docker run -p 5000:5000 \ - -e NVIDIA_API_KEY=your_key \ + -e NVIDIA_NIM_API_KEY=your_key \ -v $(pwd)/output:/app/output \ ellipse2v/secopstm:latest @@ -48,7 +48,7 @@ docker run --rm \ --init-rag docker run -p 5000:5000 \ - -e NVIDIA_API_KEY=your_key \ + -e NVIDIA_NIM_API_KEY=your_key \ -v secopstm-rag:/app/rag \ -v $(pwd)/output:/app/output \ ellipse2v/secopstm:latest @@ -116,5 +116,5 @@ The framework includes a web-based interface for interactive threat modeling, ac 2. **Open your browser** to the address shown in the console (usually `http://127.0.0.1:5000/`). 3. **Choose a mode:** - - **Simple Mode**: Ideal for quick visualization and editing of threat models written in Markdown. It features a live preview and now supports multi-file projects through a tabbed interface, allowing you to edit a main model and its sub-models together. - - **Graphical Editor**: A visual, drag-and-drop canvas for building threat models from scratch without writing Markdown. This mode is under active development. + - **Simple Mode**: For quick visualization and editing of system models written in Markdown. It features a live preview and supports multi-file projects through a tabbed interface, so you can edit a main model and its sub-models together. + - **Graphical Editor**: A visual, drag-and-drop canvas for building system models from scratch without writing Markdown. This mode is under active development. diff --git a/docs/index.md b/docs/index.md index 54e34d5..7302443 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,12 +1,12 @@ -# Elevating Cyber Resilience with Automated Threat Modeling +# SecOpsTM Documentation -Welcome to the official documentation for **SecOpsTM** — a STRIDE threat modeling framework -with MITRE ATT&CK mapping, AI-enhanced threat generation, and interactive diagram export. +**SecOpsTM** is a STRIDE threat modeling framework with MITRE ATT&CK mapping, AI-enhanced +threat generation, and interactive diagram export. -In an era of escalating cyber threats and rapid development cycles, traditional security practices -often fall short. SecOpsTM bridges that gap by embodying **Threat Modeling as Code (TMasC)**: -threat models are human-readable, version-controlled Markdown files that integrate naturally into -CI/CD pipelines and cross-functional workflows. +It follows a **Threat Modeling as Code (TMasC)** approach: you describe your system in a +human-readable, version-controlled Markdown file (the system model), and SecOpsTM generates +the threat model from it — a workflow that fits naturally into CI/CD pipelines and +cross-functional review. ## Why Automated Threat Modeling? @@ -23,17 +23,18 @@ CI/CD pipelines and cross-functional workflows. - **AI-enhanced generation** — LLM + RAG pipeline surfaces threats beyond rule-based analysis - **Context-aware severity** — scoring adjusts for encryption, authentication, network exposure, CVE signals, and D3FEND mitigations - **Hierarchical modeling** — decompose large systems into linked sub-models with drill-down diagrams -- **Comprehensive exports** — HTML reports, STIX 2.1, ATT&CK Navigator layers, SVG diagrams, ZIP bundles -- **IaC integration** — generate threat models directly from Ansible and Terraform configurations +- **Exports** — HTML reports, STIX 2.1, ATT&CK Navigator layers, SVG diagrams, ZIP bundles +- **IaC integration** — generate system models directly from Ansible and Terraform configurations ## Documentation - [Getting Started](getting_started.md) — Installation, first run, web editor -- [**Claude Code Skill**](claude-code-skill.md) — AI-assisted threat model generation: describe your system, Claude writes the DSL -- [**Workflow: Conception to Run**](workflow.md) — Complete end-to-end guide: conception, modeling, enrichment, CI/CD +- [**Claude Code Skill**](claude-code-skill.md) — AI-assisted system model generation: describe your system, Claude writes the DSL +- [**Workflow: Conception to Run**](workflow.md) — End-to-end guide: conception, modeling, enrichment, CI/CD - [Usage](usage.md) — CLI flags, project mode, export formats -- [Features](features.md) — Full feature list: AI engines, diagrams, exports -- [Defining Threat Models](defining_threat_models.md) — Markdown DSL reference +- [Features](features.md) — Feature list: AI engines, diagrams, exports +- [Goal-Driven Attack Flows (GDAF)](gdaf.md) — context YAML schema, scoring algorithm, asset type table +- [Defining Your System Model](defining_threat_models.md) — Markdown DSL reference - [Enriching AI Threats](enriching_ai_threats.md) — DSL attributes, BOM, and context files that improve AI-generated threats - [Data Collection Guide](data_collection_guide.md) — What information to gather before threat modeling - [Examples](examples.md) — Ready-to-use model templates diff --git a/docs/project_onboarding_questionnaire.md b/docs/project_onboarding_questionnaire.md index 45a9106..e71308e 100644 --- a/docs/project_onboarding_questionnaire.md +++ b/docs/project_onboarding_questionnaire.md @@ -234,4 +234,4 @@ Before handing this to the security team, confirm: --- -*SecOpsTM — Threat Model as Code. For questions, see `docs/getting_started.md` or open an issue.* +*SecOpsTM — Threat Modeling as Code. For questions, see `docs/getting_started.md` or open an issue.* diff --git a/docs/technical_documentation/architecture_overview.md b/docs/technical_documentation/architecture_overview.md index 0501492..d7fe56f 100644 --- a/docs/technical_documentation/architecture_overview.md +++ b/docs/technical_documentation/architecture_overview.md @@ -229,7 +229,7 @@ sub-model internals transparently. ## Sub-model Drill-down -Sub-model drill-down connects a server in a parent model to a child threat model file via the +Sub-model drill-down connects a server in a parent model to a child system model file via the `submodel=./path/to/model.md` DSL keyword. The child IS the parent server at higher detail. ### Parent diagram (hyperlink nodes) @@ -278,7 +278,7 @@ The framework uses a pluggable AI provider architecture to support various Large 2. **`threat_analysis/ai_engine/providers/litellm_client.py`**: * **Role**: A unified client leveraging the `litellm` library. * **Function**: Provides a consistent interface to interact with numerous AI providers (OpenAI, Anthropic, Google Gemini, Mistral, Ollama, etc.) using a single completion function. - * **Configuration**: Loads settings from `config/ai_config.yaml` and handles API keys via environment variables (e.g., `NVIDIA_API_KEY` for NVIDIA NIM, `GEMINI_API_KEY` for Gemini). + * **Configuration**: Loads settings from `config/ai_config.yaml` and handles API keys via environment variables (e.g., `NVIDIA_NIM_API_KEY` for NVIDIA NIM, `GEMINI_API_KEY` for Gemini). 3. **`threat_analysis/ai_engine/providers/litellm_provider.py`**: * **Role**: A concrete implementation of `BaseLLMProvider` that wraps `LiteLLMClient`. @@ -288,9 +288,9 @@ The framework uses a pluggable AI provider architecture to support various Large * **Role**: A dedicated provider for local Ollama instances. * **Function**: Uses direct HTTP calls to the Ollama API for low-latency local inference. -## AI-Powered Threat Model Generation and Modification +## AI-Powered System Model Generation and Modification -This section details the architecture behind the AI-driven generation and modification of threat models, available primarily through the "simple mode" interface. +This section details the architecture behind the AI-driven generation and modification of system models, available primarily through the "simple mode" interface. ### Workflow: @@ -306,7 +306,7 @@ This section details the architecture behind the AI-driven generation and modifi 3. **Service Layer Logic (`threat_analysis/server/ai_service.py`)**: * The `AiService` (formerly integrated into the threat model service) is the central orchestrator. * It utilizes `LiteLLMClient` to communicate with the configured AI provider. - * It supports streaming responses, allowing the user to see the threat model being generated in real-time. + * It supports streaming responses, allowing the user to see the system model being generated in real-time. 4. **AI Provider Interaction**: * Based on `config/ai_config.yaml`, the `LiteLLMClient` identifies the enabled provider (e.g., `gemini`, `openai`, or `ollama`). @@ -316,7 +316,7 @@ This section details the architecture behind the AI-driven generation and modifi * The generated Markdown content is streamed back to the frontend. * The CodeMirror editor is updated, and the diagram is re-rendered. -This architecture enables an iterative threat modeling process, allowing users to start with a basic model and refine it incrementally through natural language commands, significantly enhancing the usability and flexibility of the tool compared to an overwrite-only approach. +This architecture enables an iterative threat modeling process, allowing users to start with a basic model and refine it incrementally through natural language commands, improving the usability and flexibility of the tool compared to an overwrite-only approach. ## Frontend JavaScript Architecture (Graphical Editor) @@ -355,13 +355,13 @@ The JavaScript codebase for the graphical editor (`threat_analysis/server/templa * **Details**: Utilizes `NodeManager` to create new nodes and `PropertiesPanelManager` to immediately display their properties upon creation. * **`ThreatModelGenerator.js`**: - * **Role**: Handles the process of converting the visual graph into a structured threat model representation and initiating the backend generation process. - * **Function**: Collects all nodes and connections from the canvas, constructs a JSON representation of the threat model, converts it to Markdown, and sends it to the `/api/generate_all` endpoint. + * **Role**: Handles the process of converting the visual graph into a structured system model representation and initiating the backend generation process. + * **Function**: Collects all nodes and connections from the canvas, constructs a JSON representation of the system model, converts it to Markdown, and sends it to the `/api/generate_all` endpoint. * **Details**: Also responsible for displaying generation status and results. * **`ModelManager.js`**: - * **Role**: Manages loading and saving threat models from the server or local files. + * **Role**: Manages loading and saving system models from the server or local files. * **Function**: Handles interactions with the "Open Model" modal, fetches lists of saved models, loads models via API calls, and handles local file uploads (Markdown and metadata JSON). * **Details**: Utilizes `NodeManager` and `ConnectionManager` to repopulate the graph from loaded model data and positions. -This modular design promotes reusability, testability, and a clear separation of concerns, making the graphical editor more robust and easier to extend. \ No newline at end of file +This modular design promotes reusability, testability, and a clear separation of concerns, making the graphical editor easier to maintain and extend. \ No newline at end of file diff --git a/docs/technical_documentation/index.md b/docs/technical_documentation/index.md index 2f1d2fd..6a762ad 100644 --- a/docs/technical_documentation/index.md +++ b/docs/technical_documentation/index.md @@ -4,17 +4,17 @@ ### 1.1. The Challenge: From Manual Diagrams to Automated Analysis -As software systems grow in complexity, proactively identifying security vulnerabilities during the design phase is significantly more effective than reacting to them post-deployment. Threat modeling provides a structured process for this, but traditional approaches often rely on manual diagramming and static documents that are difficult to maintain and impossible to integrate into automated development pipelines. +As software systems grow in complexity, proactively identifying security vulnerabilities during the design phase is more effective than reacting to them post-deployment. Threat modeling provides a structured process for this, but traditional approaches often rely on manual diagramming and static documents that are difficult to maintain and impossible to integrate into automated development pipelines. -This document provides a detailed technical overview of the **SecOpsTM** framework, a tool designed to address these challenges by treating the threat model as a living artifact that evolves with the system itself. +This document provides a detailed technical overview of the **SecOpsTM** framework, a tool designed to address these challenges by treating the system model as a living artifact that evolves with the system itself. -### 1.2. Core Philosophy: System-Level Threat Model as Code +### 1.2. Core Philosophy: System-Level Threat Modeling as Code -The guiding philosophy of this framework is **Threat Model as Code (TMaC)**, applied at the **system level**. Instead of focusing on abstract application components, our approach defines the entire system architecture—including infrastructure, network boundaries, and data flows—in a simple, version-controllable format. +The guiding philosophy of this framework is **Threat Modeling as Code (TMasC)**, applied at the **system level**. Instead of focusing on abstract application components, our approach defines the entire system architecture—including infrastructure, network boundaries, and data flows—in a simple, version-controllable format. -This is particularly powerful when generating models directly from **Infrastructure as Code (IaC)** sources like Ansible playbooks. By parsing the same files that define the deployed environment, the framework creates a threat model that is a true representation of the running system. This enables a seamless, automated workflow where changes in infrastructure are immediately reflected in the threat analysis. +This is especially useful when generating models directly from **Infrastructure as Code (IaC)** sources like Ansible playbooks. By parsing the same files that define the deployed environment, the framework creates a system model that is a true representation of the running system. This enables an automated workflow where changes in infrastructure are immediately reflected in the threat analysis. -By defining the system in a Markdown DSL, the threat model becomes: +By defining the system in a Markdown DSL, the system model becomes: - **Versioned**: Stored in Git to track its evolution alongside the source code and infrastructure code. - **Automated**: Integrated directly into CI/CD pipelines to run analysis on every change. - **Collaborative**: Developers and operations engineers can contribute using the same tools and workflows they use for code. @@ -25,17 +25,17 @@ To understand the unique value of SecOpsTM, it's useful to compare it to other p | Feature | Microsoft TMT | OWASP Threat Dragon | SecOpsTM (This Tool) | | :--- | :--- | :--- | :--- | -| **Primary Paradigm** | GUI-based Diagramming | Web-based Diagramming | **Threat Model as Code (TMaC)** | +| **Primary Paradigm** | GUI-based Diagramming | Web-based Diagramming | **Threat Modeling as Code (TMasC)** | | **Input Format** | Proprietary `.tm7` format | JSON, with a web UI | **Markdown (DSL) / IaC Playbooks** | | **Automation & CI/CD** | None. Fully manual process. | Limited. Has an API but is not designed for pipeline integration. | **Core Feature**. Designed to be run from the CLI in a pipeline. | | **IaC Integration** | None. | None. | **Yes (Ansible)**. Can generate a model directly from infrastructure definitions. | -| **Version Control** | Possible by archiving the `.tm7` model file. However, the format is complex (XML-based) and not well-suited for line-by-line diffing or merging. | Feasible (JSON), but the diagram is the primary source of truth, not the code. | **Seamless**. Markdown is text-based and ideal for Git. | +| **Version Control** | Possible by archiving the `.tm7` model file. However, the format is complex (XML-based) and not well-suited for line-by-line diffing or merging. | Feasible (JSON), but the diagram is the primary source of truth, not the code. | **Direct**. Markdown is text-based and ideal for Git. | | **Extensibility** | Limited to templates. | Good. Open-source and extensible. | **High**. Mappings and logic are in simple Python dictionaries and modules. | | **Visualization & Reporting** | Basic reports. | Printable report, basic threat view. | Rich HTML reports, STIX 2.1 export, **MITRE ATT&CK Navigator layers**. | ## 3. High-Level Architecture -The framework is a Python-based application that can be run as a command-line tool (for automation) or a web server (for interactive editing). It ingests a threat model source and produces a suite of artifacts. +The framework is a Python-based application that can be run as a command-line tool (for automation) or a web server (for interactive editing). It ingests a system model source and produces a suite of artifacts. ```mermaid graph TD @@ -76,7 +76,7 @@ graph TD ## 4. Technical Deep Dive: Module by Module -This section provides a comprehensive breakdown of each component of the SecOpsTM framework. +This section provides a detailed breakdown of each component of the SecOpsTM framework. ### 4.1. Entrypoint and Orchestration (`threat_analysis/__main__.py`) @@ -93,7 +93,7 @@ The execution of the framework begins in `__main__.py`. This script is responsib The `ThreatModel` class is the heart of the framework, serving as the in-memory representation of the system under analysis. It is designed to be a rich, stateful object that not only holds the architectural components but also orchestrates the analysis process. -- **A Wrapper Around PyTM**: At its core, the `ThreatModel` class wraps a `pytm.TM` object. This allows the framework to leverage the foundational threat generation logic of the PyTM library while extending it with custom features, more detailed component attributes, and advanced analysis capabilities. +- **A Wrapper Around PyTM**: At its core, the `ThreatModel` class wraps a `pytm.TM` object. This allows the framework to use the foundational threat generation logic of the PyTM library while extending it with custom features, more detailed component attributes, and advanced analysis capabilities. - **Self-Initialization of MitreMapping**: The `ThreatModel` class now self-initializes its `MitreMapping` instance. This simplifies the constructor of `ThreatModel` and centralizes the management of the `MitreMapping` dependency within the `ThreatModel` itself, improving modularity and testability. @@ -113,12 +113,12 @@ The `ThreatModel` class is the heart of the framework, serving as the in-memory - **The Analysis Engine: `process_threats()`**: This is the most important method in the class. It serves as the orchestrator for the entire threat analysis pipeline, executing the following steps in order: 1. **Validation**: It first calls the `ModelValidator` to ensure the model is consistent and complete (e.g., no dataflows pointing to non-existent elements). 2. **PyTM Threat Generation**: It calls `self.tm.process()` to trigger the standard PyTM threat generation. - 3. **Target Expansion**: It calls `_expand_class_targets()`, a crucial helper method that takes generic threats targeted at a *class* of objects (e.g., a threat against all `Server` objects) and creates a specific threat instance for every `Server` defined in the model. This ensures that all components are evaluated correctly. + 3. **Threat Collection**: It reads `tm.findings`, PyTM's condition-matched findings, which are already resolved per-instance rather than per-class. 4. **Custom Threat Generation**: It integrates custom threats defined in the rule-based engine (`custom_threats.py`). 5. **Grouping**: It groups all generated threats (both from PyTM and custom rules) by their STRIDE category. 6. **MITRE Analysis**: Finally, it triggers the MITRE ATT&CK mapping and enrichment process. -- **`get_all_threats_details()`**: This method provides a clean, comprehensive list of all identified threats, including their description, target, STRIDE category, severity, and associated MITRE techniques. This is the primary data source for the report generation modules. +- **`get_all_threats_details()`**: This method provides a clean, complete list of all identified threats, including their description, target, STRIDE category, severity, and associated MITRE techniques. This is the primary data source for the report generation modules. ### 4.3. Model Parsing and Validation @@ -141,19 +141,19 @@ The framework's custom threat generation is driven by a flexible, rule-based eng The core of the framework leverages the `pytm` library for its foundational threat generation capabilities. When `threat_model.process_threats()` is called, `pytm` automatically analyzes the defined architecture (actors, servers, dataflows, and their properties) to identify potential STRIDE threats. **How PyTM Generates Threats:** -`pytm` applies a set of predefined rules based on the relationships and properties of elements in the threat model. For example: +`pytm` applies a set of predefined rules based on the relationships and properties of elements in the system model. For example: * A dataflow between an actor and a server might trigger "Spoofing" or "Repudiation" threats. * Dataflows marked as unencrypted (`is_encrypted=False`) can lead to "Information Disclosure" threats. * Servers with specific stereotypes (e.g., "Database") might generate threats related to data tampering or unauthorized access. **Influencing PyTM's Threats:** -To "add" or "remove" threats generated directly by `pytm`, you primarily need to modify the underlying architecture of your threat model. This includes: +To "add" or "remove" threats generated directly by `pytm`, you primarily need to modify the underlying architecture of your system model. This includes: * **Adding/Removing Elements:** Introducing new actors, servers, or dataflows can trigger new `pytm` threats. Conversely, removing elements can eliminate threats associated with them. * **Modifying Element Properties:** Changing properties like `is_encrypted` for dataflows, or `stereotype` for servers, can alter the set of threats `pytm` generates. * **Structuring Boundaries:** How elements are placed within trust boundaries can also influence `pytm`'s threat identification. **Filtering PyTM Threats:** -While `pytm` generates a comprehensive set of threats, the framework allows for post-processing and filtering. The `_expand_class_targets` method in `models_module.py` and the overall threat processing pipeline can be extended to filter or modify `pytm`-generated threats before they are presented in reports. +While `pytm` generates a broad set of threats, the framework allows for post-processing and filtering. The `process_threats()` pipeline in `models_module.py` can be extended to filter or modify `pytm`-generated threats before they are presented in reports. ### 4.4.2. Customizing PyTM's Threat Database @@ -232,7 +232,7 @@ The framework does not currently expose a direct `-ignore` parameter for `pytm`- ``` This rule states: "If a dataflow is not encrypted AND uses the HTTP protocol, then generate an 'Information Disclosure' threat." -- **`custom_threats.py`: The Engine Itself**: This module contains the `get_custom_threats` function which acts as the engine that interprets the rules. This module has been refactored to use a more robust property lookup mechanism (`_get_property`) that supports nested attributes (e.g., `source.inBoundary.isTrusted`), making the rule application more powerful and flexible. +- **`custom_threats.py`: The Engine Itself**: This module contains the `get_custom_threats` function which acts as the engine that interprets the rules. This module has been refactored to use a more resilient property lookup mechanism (`_get_property`) that supports nested attributes (e.g., `source.inBoundary.isTrusted`), making the rule application more flexible. - **`get_custom_threats(threat_model)`**: This function is the main entry point for the custom threat generation process. It takes the fully parsed `ThreatModel` object as input. - **Iteration and Matching**: The function iterates through every component (server, dataflow, actor, etc.) in the `threat_model`. For each component, it retrieves the relevant rules from `THREAT_RULES` (based on the component's type). It then checks if the component's properties match all the `conditions` specified in a rule. - **Boundary-Aware Logic**: A key feature of the engine is its ability to handle complex conditions, especially for dataflows. It can check the properties of the source and sink of a dataflow, including which boundary they are in. For example, a rule can be written to only trigger a threat if a dataflow crosses from an untrusted boundary (like the "Internet") to a trusted one (like the "Internal Network"). @@ -264,7 +264,7 @@ graph TD ### 4.5. The STRIDE, CAPEC, and ATT&CK Mapping (`mitre_mapping_module.py`) -This is the most complex and critical module for enriching the raw threat data. It transforms high-level STRIDE threats into specific, actionable MITRE ATT&CK techniques through a chained mapping process that leverages established cybersecurity knowledge bases. The module was recently updated to fix a bug where the D3FEND mitigation name was not being correctly displayed in the report. The fix involved updating the regex to correctly parse the mitigation name from the `ATTACK_D3FEND_MAPPING` and updating the report template to display the name instead of the description. +This is the most complex and critical module for enriching the raw threat data. It transforms high-level STRIDE threats into specific, actionable MITRE ATT&CK techniques through a chained mapping process that uses established cybersecurity knowledge bases. The module was recently updated to fix a bug where the D3FEND mitigation name was not being correctly displayed in the report. The fix involved updating the regex to correctly parse the mitigation name from the `ATTACK_D3FEND_MAPPING` and updating the report template to display the name instead of the description. - **`MitreMapping` Class**: The central class that orchestrates the entire enrichment pipeline. - **Initialization**: When instantiated, it pre-loads and processes several external data sources to build its mapping tables: @@ -308,7 +308,7 @@ The `SeverityCalculator` provides a nuanced risk score for each threat. - **Multi-Factor Calculation**: The final score is not a static value but a composite calculated from: 1. **Base Score**: A default score for each STRIDE category. 2. **Rule-Defined Score**: The impact and likelihood values (1-5) defined in the `threat_rules.py` entry for that threat. - 3. **Target Multipliers**: The score can be increased by multipliers defined in the `## Severity Multipliers` section of the threat model, which are loaded from the markdown file. + 3. **Target Multipliers**: The score can be increased by multipliers defined in the `## Severity Multipliers` section of the system model, which are loaded from the markdown file. 4. **Protocol Adjustments**: The protocol of a dataflow can adjust the score (e.g., HTTP increases it, HTTPS decreases it). 5. **Data Classification**: The classification of the data in a flow (`PUBLIC`, `SECRET`, etc.) acts as a final multiplier. - **Normalization**: The final score is clamped between 1.0 and 10.0 and assigned a qualitative level (e.g., "HIGH", "CRITICAL"). @@ -356,7 +356,7 @@ Results appear in a **"⛓️ Attack Chain Analysis"** section in the HTML repor - It uses a Jinja2 template (`threat_model.dot.j2`) to generate Graphviz DOT language code from the `ThreatModel` object. - **Trust Boundary Colors**: Trusted boundaries use `color="#2e7d32"` (dark green, solid); untrusted use `color="#c62828"` (dark red, dashed). These values are baked into the DOT template and thus appear in both exported SVG and HTML diagrams. - **Severity Heat Map**: `_generate_html_with_legend()` accepts `severity_map` (dict of component → severity label) and `report_url`. The HTML diagram template injects these as JavaScript variables, enabling a toggle button that applies/restores per-component colour overlays and shows hover tooltips with "View threats →" deep-links into the HTML report. - - **Visual Styling**: The generator includes sophisticated logic for rich visual styling, combining native Graphviz shapes with embedded SVG icons. The layout of the icon and text is adjusted based on the element type for maximum clarity: + - **Visual Styling**: The generator includes detailed logic for rich visual styling, combining native Graphviz shapes with embedded SVG icons. The layout of the icon and text is adjusted based on the element type for maximum clarity: - **Native Shapes & Sizing**: It assigns semantic shapes to elements and sets their sizes for a clean visual hierarchy: - **Actors**: Rendered as fixed-size **circles**. - **Switches and Firewalls**: Rendered as fixed-size diamonds and hexagons, respectively, which are smaller than other nodes for visual distinction. @@ -399,7 +399,7 @@ graph TD ``` - **`stix_generator.py`**: This module provides interoperability. - It translates the framework's findings into STIX 2.1, a standardized language for cyber threat intelligence. - - It leverages the `attack-flow` STIX extension to create a structured representation of the attack chains, creating `attack-action` and `attack-asset` objects and linking them with relationships. + - It uses the `attack-flow` STIX extension to create a structured representation of the attack chains, creating `attack-action` and `attack-asset` objects and linking them with relationships. - **`attack_navigator_generator.py`**: This module creates a JSON layer file compatible with the [MITRE ATT&CK Navigator](https://mitre-attack.github.io/attack-navigator/) to visualize the results of the analysis. - The `AttackNavigatorGenerator` class takes the threat model's name and a list of all detailed threats. - It processes the threats to extract all unique ATT&CK techniques. For each technique, it aggregates the findings, using the highest severity score as the technique's score and compiling the descriptions of all threats mapped to it in the comments. @@ -409,7 +409,7 @@ graph TD - **`server.py`**: A simple Flask application that defines the API endpoints: - `/`: Serves the main `web_interface.html`. - - `/fullGUI`: Serves the `full_gui.html` with a more comprehensive interface. + - `/fullGUI`: Serves the `full_gui.html` with a fuller interface. - `/api/update`: Receives Markdown from the editor, triggers a live analysis, and returns the resulting SVG diagram and legend. - `/api/export` & `/api/export_all`: Handle requests to download the generated artifacts. - **`threat_model_service.py`**: This service layer acts as a bridge between the web server and the core analysis engine. It encapsulates the logic for handling web requests, calling the appropriate framework components, and managing temporary files, keeping the Flask app clean and focused on routing. It has been updated to align with the new `create_threat_model` signature, removing the `mitre_mapping` argument from its calls. @@ -426,8 +426,8 @@ This module provides actionable mitigation advice for the threats identified dur - **`get_framework_mitigation_suggestions()` Function**: This is the primary function exposed by the module. - It takes a list of ATT&CK technique IDs (extracted from the threats during the report generation phase). - - It now leverages the `MITRE_TO_CIS_MAP` (generated by `_create_mitre_to_cis_map`) to provide CIS control suggestions. - - The `FRAMEWORK_MITIGATION_MAP` has been updated with more relevant OWASP ASVS entries, particularly for Information Disclosure threats, ensuring more comprehensive and accurate mitigation advice. + - It now uses the `MITRE_TO_CIS_MAP` (generated by `_create_mitre_to_cis_map`) to provide CIS control suggestions. + - The `FRAMEWORK_MITIGATION_MAP` has been updated with more relevant OWASP ASVS entries, particularly for Information Disclosure threats, giving more complete and accurate mitigation advice. - **Architecture Flow**: ```mermaid @@ -456,7 +456,7 @@ Key configurations include: ### 4.11. IaC Plugin Architecture (`iac_plugins/`) -The framework is designed to be extensible through a dedicated Infrastructure as Code (IaC) plugin system, allowing it to generate threat models from various IaC sources. +The framework is designed to be extensible through a dedicated Infrastructure as Code (IaC) plugin system, allowing it to generate system models from various IaC sources. - **Abstract Base Class**: The `iac_plugins/__init__.py` file defines an abstract base class called `IaCPlugin`. To create a new plugin, a developer must create a class that inherits from `IaCPlugin`. - **Required Implementations**: Any new plugin must implement three key methods: @@ -467,14 +467,14 @@ The framework is designed to be extensible through a dedicated Infrastructure as ### 4.12. Ansible Plugin and Metadata (`iac_plugins/ansible_plugin.py`) -The Ansible plugin is a concrete implementation of the IaC plugin architecture, designed to translate an existing Ansible project into a threat model. It works by combining information from the Ansible playbook and inventory with a dedicated metadata structure that describes the security-relevant aspects of the architecture. +The Ansible plugin is a concrete implementation of the IaC plugin architecture, designed to translate an existing Ansible project into a system model. It works by combining information from the Ansible playbook and inventory with a dedicated metadata structure that describes the security-relevant aspects of the architecture. - **How it Works**: The plugin is triggered when the `--ansible-path` argument is used, pointing to a main playbook file (e.g., `playbook.yml`). 1. **Parsing**: The plugin first parses the specified playbook. It also looks for a corresponding inventory file named `hosts.ini` in the same directory. - 2. **Metadata Extraction**: The crucial step is the extraction of a special variable named `threat_model_metadata` from the `vars` section of the playbook. This variable must be a dictionary that contains the threat model definition. - 3. **Model Generation**: The plugin then uses the data from the `threat_model_metadata` dictionary to generate the components of the threat model (Boundaries, Actors, Servers, Dataflows) in the Markdown DSL format. + 2. **Metadata Extraction**: The crucial step is the extraction of a special variable named `threat_model_metadata` from the `vars` section of the playbook. This variable must be a dictionary that contains the system model definition. + 3. **Model Generation**: The plugin then uses the data from the `threat_model_metadata` dictionary to generate the components of the system model (Boundaries, Actors, Servers, Dataflows) in the Markdown DSL format. -- **The `threat_model_metadata` Structure**: This is the core concept for the Ansible integration. Instead of trying to infer the entire architecture from Ansible tasks and roles (which can be ambiguous), the framework requires the user to explicitly define the threat model's structure within the playbook itself. This approach keeps the threat model definition alongside the infrastructure code that it describes. +- **The `threat_model_metadata` Structure**: This is the core concept for the Ansible integration. Instead of trying to infer the entire architecture from Ansible tasks and roles (which can be ambiguous), the framework requires the user to explicitly define the system model's structure within the playbook itself. This approach keeps the system model definition alongside the infrastructure code that it describes. - The `threat_analysis/iac_plugins/ansible_threat_model_config.yml` file serves as a **template or example** of what this `threat_model_metadata` variable should look like. It is **not** a configuration file that is read by the plugin. - The user is expected to copy and adapt this structure into the `vars` section of their own Ansible playbook. @@ -502,7 +502,7 @@ The Ansible plugin is a concrete implementation of the IaC plugin architecture, roles: - webserver ``` - In this example, the `threat_model_metadata` variable is defined directly within the playbook. The plugin will parse this variable to create the "Public DMZ" and "Internal Network" boundaries, the "WebApp Server" component, and the "User Traffic" dataflow. The use of Ansible variables like `{{ ansible_default_ipv4.address }}` within the metadata is also supported, allowing the threat model to be dynamically updated with information from the inventory. + In this example, the `threat_model_metadata` variable is defined directly within the playbook. The plugin will parse this variable to create the "Public DMZ" and "Internal Network" boundaries, the "WebApp Server" component, and the "User Traffic" dataflow. The use of Ansible variables like `{{ ansible_default_ipv4.address }}` within the metadata is also supported, allowing the system model to be dynamically updated with information from the inventory. ## 5. Tooling and Data Maintenance @@ -534,7 +534,7 @@ The framework relies on external data from MITRE ATT&CK, CAPEC, and other source ### 4.13. Attack Flow Generation (`generation/attack_flow_generator.py`) -This module is responsible for generating visual, end-to-end attack scenarios based on the identified threats, compatible with the [Attack Flow](https://attackflow.io/) tool. It has been significantly enhanced to provide more realistic and optimized attack paths. +This module is responsible for generating visual, end-to-end attack scenarios based on the identified threats, compatible with the [Attack Flow](https://attackflow.io/) tool. It has been enhanced to provide more realistic and optimized attack paths. - **Core Philosophy**: The generator's goal is to transform a flat list of threats into meaningful narratives. It achieves this by sequencing threats according to the logical progression of adversary tactics as defined by the MITRE ATT&CK framework. - **Path Discovery Logic (`_find_attack_paths`)**: @@ -550,7 +550,7 @@ This module is responsible for generating visual, end-to-end attack scenarios ba - It filters out threats targeting generic classes (e.g., `pytm.Server` class itself) or tuples of classes, focusing only on threats against specific asset instances. - **Diagram Structure and Content**: - Each generated `.afb` file represents a single, complete attack path. - - The diagram illustrates the progression by alternating between **Actions** (MITRE ATT&CK techniques) and **Assets** (the specific components from your threat model). + - The diagram illustrates the progression by alternating between **Actions** (MITRE ATT&CK techniques) and **Assets** (the specific components from your system model). - The final node in the chain is a conceptual `asset` representing the adversary's objective, derived from the STRIDE category of the final threat (e.g., "Impact: Tampering"). - The resulting flow visually communicates a clear narrative: `(Action 1) -> (targets Asset A) -> (enabling Action 2) -> (targets Asset B) -> ... -> (achieves Impact)`. - **Asset Name Resolution**: Uses dedicated helper methods (`_get_target_name`, `_extract_name_from_object`) to accurately resolve asset names from raw `pytm` objects, ensuring human-readable labels in the generated attack flows. diff --git a/docs/threat_modeling_guide.md b/docs/threat_modeling_guide.md index c55bcfa..90c85aa 100644 --- a/docs/threat_modeling_guide.md +++ b/docs/threat_modeling_guide.md @@ -1,6 +1,6 @@ -# How to Create a Threat Model with SecOpsTM DSL +# How to Create a System Model with SecOpsTM DSL -This guide walks you through building a complete, production-quality threat model for a real +This guide walks you through building a complete, production-quality system model for a real system. By the end, you will have a working SecOpsTM model file (`.md`) that produces STRIDE threat reports, MITRE ATT&CK mappings, GDAF attack paths, and navigable HTML diagrams. @@ -68,7 +68,7 @@ Multi-line definitions are supported — indent continuation lines: Start the file with: ```markdown -# Threat Model: My Application Name +# System Model: My Application Name ## Description Brief description of the system under analysis. Include: diff --git a/docs/usage.md b/docs/usage.md index 1acfb3f..d8c6dad 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -2,13 +2,13 @@ This framework supports two modes of operation: Command Line Interface (CLI) for automated analysis and a Web-based Graphical User Interface (GUI) for interactive editing and visualization. -## Threat Model as Code Philosophy +## Threat Modeling as Code Philosophy -This framework is designed to be used in a "Threat Model as Code" workflow. This means that the threat model is defined in a simple, version-controllable format (Markdown), and the analysis is performed by running a script. This approach has several advantages: +This framework is designed to be used in a "Threat Modeling as Code" workflow. This means that the system model is defined in a simple, version-controllable format (Markdown), and the threat analysis is performed by running a script. This approach has several advantages: -- **Version Control**: Threat models can be stored in a Git repository, allowing you to track changes over time. +- **Version Control**: System models can be stored in a Git repository, allowing you to track changes over time. - **Automation**: The threat modeling process can be integrated into your CI/CD pipeline, allowing you to automatically update your threat model whenever your architecture changes. -- **Collaboration**: Developers can collaborate on the threat model using the same tools they use for code. +- **Collaboration**: Developers can collaborate on the system model using the same tools they use for code. ## 0. Data layers and installation @@ -36,7 +36,7 @@ docker run -p 5000:5000 -v $(pwd)/output:/app/output \ # With LLM inference — default provider is NVIDIA NIM (free tier available) # See https://build.nvidia.com/meta/llama-3_3-70b-instruct for a free API key docker run -p 5000:5000 \ - -e NVIDIA_API_KEY=your_key \ + -e NVIDIA_NIM_API_KEY=your_key \ -v $(pwd)/output:/app/output \ ellipse2v/secopstm:latest @@ -49,7 +49,7 @@ docker run -p 5000:5000 \ # With LLM + RAG docker run -p 5000:5000 \ - -e NVIDIA_API_KEY=your_key \ + -e NVIDIA_NIM_API_KEY=your_key \ -v secopstm-rag:/app/rag \ -v $(pwd)/output:/app/output \ ellipse2v/secopstm:latest @@ -74,10 +74,12 @@ The named volume `secopstm-rag` persists across container restarts and image reb | Provider | Environment variable | Free tier | |---|---|---| -| NVIDIA NIM *(default)* | `NVIDIA_API_KEY` | ✅ [build.nvidia.com](https://build.nvidia.com/meta/llama-3_3-70b-instruct) | +| NVIDIA NIM *(default)* | `NVIDIA_NIM_API_KEY` | ✅ [build.nvidia.com](https://build.nvidia.com/meta/llama-3_3-70b-instruct) | | Google Gemini | `GEMINI_API_KEY` | ✅ | | OpenAI | `OPENAI_API_KEY` | ❌ | | Mistral | `MISTRAL_API_KEY` | ✅ | +| Groq | `GROQ_API_KEY` | ✅ [console.groq.com](https://console.groq.com) | +| xAI | `XAI_API_KEY` | ❌ [console.x.ai](https://console.x.ai) | | Ollama (local) | — no key needed — | ✅ | **Volumes and mounts** @@ -87,7 +89,7 @@ The named volume `secopstm-rag` persists across container restarts and image reb | Output reports | `-v $(pwd)/output:/app/output` | Files land in `output//` on the host | | AI config | `-v $(pwd)/ai_config.yaml:/app/config/ai_config.yaml:ro` | Switch provider/model without rebuilding | | Prompts | `-v $(pwd)/prompts.yaml:/app/config/prompts.yaml:ro` | Override LLM prompts | -| Threat model files | `-v $(pwd)/models:/models` | Pass `--model-file /models/model.md` | +| System model files | `-v $(pwd)/models:/models` | Pass `--model-file /models/model.md` | | CVE definitions | `-v $(pwd)/cve_definitions.yml:/app/cve_definitions.yml:ro` | Per-asset CVE list | | RAG vector store | `-v secopstm-rag:/app/rag` | Named volume, required for RAG | @@ -174,7 +176,7 @@ secopstm --server You can also still use `python -m threat_analysis` with all the same flags — they are 100% equivalent. -1. **Learn how to define your threat model in Markdown** by reading the [Defining Your Threat Model](defining_threat_models.md) guide. +1. **Learn how to define your system model in Markdown** by reading the [Defining Your System Model](defining_threat_models.md) guide. 2. **Generate Attack Flow diagrams:** Add the `--attack-flow` flag to generate `.afb` files for key STRIDE objectives (Tampering, Spoofing, Information Disclosure, Repudiation). ```bash secopstm --model-file path/to/your_model.md --attack-flow @@ -200,7 +202,7 @@ Example: python -m threat_analysis --model-file path/to/your_model.md --implemented-mitigations-file path/to/your/mitigations.txt --cve-definitions-file path/to/your/cves.yml ``` -Here's a more comprehensive example using test files: +Here's another example using test files: ```bash python -m threat_analysis --model-file threatModel_Template/threat_model.md \ --navigator \ @@ -209,9 +211,9 @@ python -m threat_analysis --model-file threatModel_Template/threat_model.md \ --cve-definitions-file tests/cve_definitions.yml ``` -### 2. Project Mode: Hierarchical Threat Models +### 2. Project Mode: Hierarchical System Models -The framework excels at handling complex projects with multiple, nested threat models. While you can run project-based analysis from the CLI, the recommended workflow is to use the **Web-based User Interface (Server Mode)**, which provides a more interactive and intuitive experience. +The framework handles projects with multiple, nested system models. You can run project-based analysis from the CLI, but the recommended workflow is the **Web-based User Interface (Server Mode)**, which is more interactive. 1. **Organize your project** in a directory, with a `main.md` at the root and sub-models in sub-directories (e.g., `my_project/main.md`, `my_project/backend/model.md`). 2. **Launch the server with your project path:** @@ -222,7 +224,7 @@ The framework excels at handling complex projects with multiple, nested threat m ### 3. Infrastructure as Code (IaC) Integration (Ansible Example) -This framework can automatically generate a complete threat model directly from IaC configurations. It automatically includes a set of default protocol styles from `threatModel_Template/base_protocol_styles.md` to ensure consistent visualization. +This framework can generate a complete system model directly from IaC configurations. It automatically includes a set of default protocol styles from `threatModel_Template/base_protocol_styles.md` to ensure consistent visualization. Here's how to use the Ansible plugin with a sample playbook: @@ -231,7 +233,7 @@ Here's how to use the Ansible plugin with a sample playbook: ```bash python -m threat_analysis --ansible-path tests/ansible_playbooks/simple_web_server/simple_web_server.yml ``` - This command will generate a complete threat model based on the Ansible playbook. The generated Markdown model will be saved in the `output/` directory with a filename derived from your Ansible playbook (e.g., `simple_web_server.md`). + This command will generate a complete system model based on the Ansible playbook. The generated Markdown model will be saved in the `output/` directory with a filename derived from your Ansible playbook (e.g., `simple_web_server.md`). If you wish to specify a different output file for the generated model, you can use the `--model-file` option: ```bash @@ -241,7 +243,7 @@ Here's how to use the Ansible plugin with a sample playbook: ### 4. CVE-Based Threat Generation (Optional) -This framework can generate threats based on a list of Common Vulnerabilities and Exposures (CVEs) that you provide for specific components in your threat model. +This framework can generate threats based on a list of Common Vulnerabilities and Exposures (CVEs) that you provide for specific components in your system model. #### CVE data source @@ -260,7 +262,7 @@ The tool maps your specified CVEs to CAPEC attack patterns, which are then used 1. **Create `cve_definitions.yml`**: By default, the tool looks for `cve_definitions.yml` in the directory of the model or project. You can override this path using the `--cve-definitions-file` command-line argument. -2. **Define CVEs for your equipment**: In this file, list the equipment (servers or actors from your threat model) and the CVEs associated with them. +2. **Define CVEs for your equipment**: In this file, list the equipment (servers or actors from your system model) and the CVEs associated with them. **Example `cve_definitions.yml`:** ```yaml @@ -301,8 +303,8 @@ For a more interactive experience, the framework provides a web-based UI that ru The console will display the address (e.g., `http://127.0.0.1:5000`) to open in your web browser. 2. **Choose a Mode from the Menu:** - - **Simple Mode**: An interface designed for editing and visualizing threat models described in Markdown. It features a tabbed editor, a live interactive diagram, and full reporting capabilities. When a project is loaded, all model files are automatically opened in separate tabs. - - **Graphical Editor**: A full-featured, interactive canvas to build, modify, and analyze threat models from scratch directly in the browser. It includes a toolbar for adding elements, a properties panel for editing, and the ability to generate all artifacts without touching Markdown directly. + - **Simple Mode**: An interface for editing and visualizing system models described in Markdown. It features a tabbed editor, a live interactive diagram, and full reporting. When a project is loaded, all model files open automatically in separate tabs. + - **Graphical Editor**: An interactive canvas to build, modify, and analyze system models from scratch directly in the browser. It includes a toolbar for adding elements, a properties panel for editing, and the ability to generate all artifacts without touching Markdown directly. ### Working with Projects and Sub-models (Simple Mode) @@ -316,11 +318,11 @@ The Simple Mode is optimized for working with complex, multi-file projects. - **Backend Services**: submodel=backend/model.md, boundary="Internal" ``` -3. **Generate the Full Project**: Click the **"Generate All"** button. The system is designed to be robust: +3. **Generate the Full Project**: Click the **"Generate All"** button: - It gathers the content from all open tabs. - - It intelligently detects if any model references a sub-model that is not currently open. - - If a missing sub-model is found, it will prompt you to select your project's root directory. It then scans this directory to find the missing files and includes them in the generation process. - - This ensures that a complete, unified, and navigable set of reports and diagrams is always generated. + - It detects if any model references a sub-model that is not currently open. + - If a missing sub-model is found, it prompts you to select your project's root directory, then scans this directory to find the missing files and includes them in the generation process. + - The result is a complete, unified, navigable set of reports and diagrams. #### Loading a Project with the Directory Picker @@ -372,7 +374,7 @@ Both the web and CLI interfaces expect files generated by `--output-format json` ## 4. GitHub Action (CI/CD) -SecOpsTM ships as an official GitHub Action for threat-model-as-code workflows: +SecOpsTM ships as an official GitHub Action for threat-modeling-as-code workflows: ```yaml # .github/workflows/threat-model.yml diff --git a/docs/workflow.md b/docs/workflow.md index 82f1eb0..0712a93 100644 --- a/docs/workflow.md +++ b/docs/workflow.md @@ -18,7 +18,7 @@ whiteboard conversation with the system owner to automated reports in CI/CD. gathering) ``` -Phases 1 and 2 are mandatory. Phase 3 is optional but significantly improves AI-generated threats. +Phases 1 and 2 are mandatory. Phase 3 is optional but improves AI-generated threats. Phase 4 can be automated. --- @@ -442,7 +442,7 @@ secopstm --server ## Related Documentation -- [Defining Threat Models](defining_threat_models.md) — complete DSL reference (all attributes, all sections) +- [Defining Your System Model](defining_threat_models.md) — complete DSL reference (all attributes, all sections) - [Enriching AI Threats](enriching_ai_threats.md) — which attributes improve AI output and how - [Data Collection Guide](data_collection_guide.md) — questionnaire for Phase 1 information gathering - [Examples](examples.md) — ready-to-use model templates diff --git a/pyproject.toml b/pyproject.toml index c7c6a7c..0d650c8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,7 @@ classifiers = [ "Topic :: Software Development :: Quality Assurance", ] dependencies = [ - "pytm>=1.3.1", + "pytm==1.3.1", # pinned exact — see requirements.txt comment / decisions.md "Flask[async]", "Jinja2", "graphviz", diff --git a/requirements.txt b/requirements.txt index 62c4a28..6dc0b01 100644 --- a/requirements.txt +++ b/requirements.txt @@ -21,7 +21,10 @@ # - Includes jq and unstructured[all] (tooling / data pipeline only) # - Versions are pinned with >= lower-bounds (not exact) to stay compatible with # fresh environments while avoiding known-bad releases. -pytm>=1.3.1 +pytm==1.3.1 # pinned exact: newer releases turned pytm.Threat into a strict-typed + # pydantic model (likelihood/impact must be str), breaking ExtendedThreat + # construction in ai_service.py and causing a pydantic __eq__ RecursionError + # in pytm's _match_responses(). Revisit when upgrading — see decisions.md. Flask[async]>=3.1.1 openpyxl>=3.1.0 Jinja2>=3.1.6 diff --git a/tests/test_ai_service.py b/tests/test_ai_service.py index f81c311..8e34c45 100644 --- a/tests/test_ai_service.py +++ b/tests/test_ai_service.py @@ -47,6 +47,30 @@ async def _run(): mock_rag.assert_called_once() asyncio.run(_run()) + +def test_init_ai_force_disable_rag_skips_rag_even_when_enabled_in_config(): + """--no-rag must skip RAGThreatGenerator entirely, even with rag.enabled: true — + used to test AI enrichment with/without RAG, and to speed up iteration (RAG + cold-start — chromadb + embeddings — can take minutes). + """ + async def _run(): + with patch("os.path.exists", return_value=True): + with patch("builtins.open", mock_open(read_data=AI_CONFIG_WITH_RAG)): + ai_service = AIService(config_path="dummy_config.yaml", force_disable_rag=True) + + with patch("threat_analysis.server.ai_service.LiteLLMProvider") as mock_provider_class, \ + patch("threat_analysis.server.ai_service.RAGThreatGenerator") as mock_rag: + + mock_provider = MagicMock() + mock_provider.check_connection = AsyncMock(return_value=True) + mock_provider_class.return_value = mock_provider + + await ai_service.init_ai() + + mock_rag.assert_not_called() + assert ai_service.rag_generator is None + asyncio.run(_run()) + def test_generate_markdown_from_prompt(ai_service): async def _run(): ai_service.ai_online = True @@ -121,6 +145,156 @@ def __init__(self, name, description, stereotype): assert "SQLi" in actor.threats[0].description asyncio.run(_run()) + +def test_enrich_with_ai_threats_missing_title_falls_back_consistently(ai_service, caplog): + """A threat dict missing 'title' (the LLM occasionally drops it under batch token + pressure) must still get a sensible SID/description, and the 'Added AI threat ...' + log line must show that same fallback — not the raw None from threat_json.get('title'). + """ + async def _run(): + ai_service.ai_online = True + ai_service.provider = MagicMock() + ai_service.provider.check_connection = AsyncMock(return_value=True) + + class MockElement: + def __init__(self, name): + self.name = name + self.description = "" + self.stereotype = "Actor" + self.threats = [] + + actor = MockElement("Actor 1") + threat_model = MagicMock() + threat_model.actors = [{'object': actor}] + threat_model.servers = [] + threat_model.dataflows = [] + threat_model.tm.description = "System desc" + + no_title = {"description": "Some threat with no title field", "category": "Spoofing"} + ai_service.provider.generate_threats_batch = AsyncMock( + return_value={"Actor 1": [no_title]} + ) + + with caplog.at_level("INFO"): + await ai_service._enrich_with_ai_threats(threat_model) + + assert len(actor.threats) == 1 + assert actor.threats[0].description.startswith("(AI) Unknown AI Threat:") + assert "None" not in caplog.text + assert "Added AI threat 'Unknown AI Threat' to Actor 1" in caplog.text + asyncio.run(_run()) + + +def test_enrich_with_ai_threats_null_title_falls_back_consistently(ai_service, caplog): + """The actually-observed real-world case: 'title' present but null (not absent) — + dict.get('title', default) alone does NOT catch this, only 'or default' does. + """ + async def _run(): + ai_service.ai_online = True + ai_service.provider = MagicMock() + ai_service.provider.check_connection = AsyncMock(return_value=True) + + class MockElement: + def __init__(self, name): + self.name = name + self.description = "" + self.stereotype = "Actor" + self.threats = [] + + actor = MockElement("Actor 1") + threat_model = MagicMock() + threat_model.actors = [{'object': actor}] + threat_model.servers = [] + threat_model.dataflows = [] + threat_model.tm.description = "System desc" + + null_title = {"title": None, "description": "Some threat", "category": "Spoofing"} + ai_service.provider.generate_threats_batch = AsyncMock( + return_value={"Actor 1": [null_title]} + ) + + with caplog.at_level("INFO"): + await ai_service._enrich_with_ai_threats(threat_model) + + assert len(actor.threats) == 1 + assert actor.threats[0].description.startswith("(AI) Unknown AI Threat:") + assert "None" not in caplog.text + asyncio.run(_run()) + + +def test_enrich_with_ai_threats_confidence_threshold_filters_low_confidence(ai_service): + """threat_generation.confidence_threshold was previously dead config — threats + below the threshold must be dropped before being attached to the element. + """ + async def _run(): + ai_service.ai_online = True + ai_service.provider = MagicMock() + ai_service.provider.check_connection = AsyncMock(return_value=True) + ai_service.confidence_threshold = 0.7 + + class MockElement: + def __init__(self, name): + self.name = name + self.description = "" + self.stereotype = "Actor" + self.threats = [] + + actor = MockElement("Actor 1") + threat_model = MagicMock() + threat_model.actors = [{'object': actor}] + threat_model.servers = [] + threat_model.dataflows = [] + threat_model.tm.description = "System desc" + + low_conf = {"title": "Low", "description": "d", "category": "Spoofing", "confidence": 0.3} + high_conf = {"title": "High", "description": "d", "category": "Spoofing", "confidence": 0.9} + ai_service.provider.generate_threats_batch = AsyncMock( + return_value={"Actor 1": [low_conf, high_conf]} + ) + + await ai_service._enrich_with_ai_threats(threat_model) + + assert len(actor.threats) == 1 + assert "High" in actor.threats[0].description + asyncio.run(_run()) + + +def test_enrich_with_ai_threats_max_threats_per_component_caps_and_keeps_highest_confidence(ai_service): + """threat_generation.max_threats_per_component was previously dead config — when + set, only the top-N threats by confidence are kept per component. + """ + async def _run(): + ai_service.ai_online = True + ai_service.provider = MagicMock() + ai_service.provider.check_connection = AsyncMock(return_value=True) + ai_service.max_threats_per_component = 1 + + class MockElement: + def __init__(self, name): + self.name = name + self.description = "" + self.stereotype = "Actor" + self.threats = [] + + actor = MockElement("Actor 1") + threat_model = MagicMock() + threat_model.actors = [{'object': actor}] + threat_model.servers = [] + threat_model.dataflows = [] + threat_model.tm.description = "System desc" + + low = {"title": "Low", "description": "d", "category": "Spoofing", "confidence": 0.4} + high = {"title": "High", "description": "d", "category": "Spoofing", "confidence": 0.95} + ai_service.provider.generate_threats_batch = AsyncMock( + return_value={"Actor 1": [low, high]} + ) + + await ai_service._enrich_with_ai_threats(threat_model) + + assert len(actor.threats) == 1 + assert "High" in actor.threats[0].description + asyncio.run(_run()) + def test_load_ai_config_not_found(ai_service): with patch("os.path.exists", return_value=False): config = ai_service._load_ai_config("nonexistent.yaml") @@ -153,6 +327,128 @@ async def _run(): assert threats[0].source == "LLM" asyncio.run(_run()) +def test_generate_rag_threats_uses_top_k_examples_default(ai_service): + """rag.top_k_examples was previously dead config — generate_threats() must be + called with the configured k, defaulting to 5 when the key is absent. + """ + async def _run(): + ai_service.rag_generator = MagicMock() + ai_service.rag_generator.generate_threats.return_value = [] + + threat_model = MagicMock() + threat_model.tm.name = "Test TM" + threat_model.tm.description = "Test desc" + threat_model.actors = [] + threat_model.servers = [] + threat_model.dataflows = [] + + await ai_service._generate_rag_threats(threat_model) + + call_kwargs = ai_service.rag_generator.generate_threats.call_args.kwargs + assert call_kwargs["k"] == 5 + asyncio.run(_run()) + + +def test_generate_rag_threats_uses_configured_top_k_examples(ai_service): + async def _run(): + ai_service.ai_config["rag"]["top_k_examples"] = 8 + ai_service.rag_generator = MagicMock() + ai_service.rag_generator.generate_threats.return_value = [] + + threat_model = MagicMock() + threat_model.tm.name = "Test TM" + threat_model.tm.description = "Test desc" + threat_model.actors = [] + threat_model.servers = [] + threat_model.dataflows = [] + + await ai_service._generate_rag_threats(threat_model) + + call_kwargs = ai_service.rag_generator.generate_threats.call_args.kwargs + assert call_kwargs["k"] == 8 + asyncio.run(_run()) + + +def test_generate_rag_threats_skipped_below_min_components_threshold(ai_service): + """rag.min_components — tiny models mostly restate single-component threats + already covered by the per-component AI pass, so RAG's system-level pass is + skipped when actors+servers (main model + sub-models) is below the threshold. + """ + async def _run(): + ai_service.rag_min_components = 3 + ai_service.rag_generator = MagicMock() + ai_service.rag_generator.generate_threats.return_value = [ + {"name": "RAG Threat", "description": "rag desc", "category": "Tampering"} + ] + + threat_model = MagicMock() + threat_model.tm.name = "Test TM" + threat_model.tm.description = "Test desc" + threat_model.actors = [MagicMock()] + threat_model.servers = [MagicMock()] + threat_model.dataflows = [] + threat_model.sub_models = [] + + threats = await ai_service._generate_rag_threats(threat_model) + + assert threats == [] + ai_service.rag_generator.generate_threats.assert_not_called() + asyncio.run(_run()) + + +def test_generate_rag_threats_runs_at_or_above_min_components_threshold(ai_service): + async def _run(): + ai_service.rag_min_components = 3 + ai_service.rag_generator = MagicMock() + ai_service.rag_generator.generate_threats.return_value = [ + {"name": "RAG Threat", "description": "rag desc", "category": "Tampering"} + ] + + threat_model = MagicMock() + threat_model.tm.name = "Test TM" + threat_model.tm.description = "Test desc" + threat_model.actors = [MagicMock()] + threat_model.servers = [MagicMock(), MagicMock()] + threat_model.dataflows = [] + threat_model.sub_models = [] + + threats = await ai_service._generate_rag_threats(threat_model) + + assert len(threats) == 1 + ai_service.rag_generator.generate_threats.assert_called_once() + asyncio.run(_run()) + + +def test_generate_rag_threats_counts_sub_model_components_toward_threshold(ai_service): + """Cross-model RAG (project mode) sends main + sub-model markdown, so the + threshold must count components across all of them, not just the main model. + """ + async def _run(): + ai_service.rag_min_components = 3 + ai_service.rag_generator = MagicMock() + ai_service.rag_generator.generate_threats.return_value = [] + + sub_model = MagicMock() + sub_model.tm.name = "Sub" + sub_model.tm.description = "Sub desc" + sub_model.actors = [MagicMock()] + sub_model.servers = [MagicMock()] + sub_model.dataflows = [] + + threat_model = MagicMock() + threat_model.tm.name = "Test TM" + threat_model.tm.description = "Test desc" + threat_model.actors = [MagicMock()] + threat_model.servers = [] + threat_model.dataflows = [] + threat_model.sub_models = [sub_model] + + await ai_service._generate_rag_threats(threat_model) + + ai_service.rag_generator.generate_threats.assert_called_once() + asyncio.run(_run()) + + def test_enrich_with_ai_threats_rag_enabled(ai_service): async def _run(): ai_service.ai_online = True @@ -174,6 +470,45 @@ async def _run(): ai_service.rag_generator.generate_threats.assert_called_once() asyncio.run(_run()) + +def test_enrich_with_ai_threats_rag_failure_does_not_abort_component_enrichment(ai_service): + """A RAG-side exception (e.g. the prompts.yaml template bug this regression-tests) + must not abort the whole enrichment — per-component threats must still be attached. + Previously an unhandled exception from step 1 (RAG) propagated out of + _enrich_with_ai_threats() entirely, so steps 2-5 (including per-component + enrichment, the main value driver) never ran. + """ + async def _run(): + ai_service.ai_online = True + ai_service.provider = MagicMock() + ai_service.provider.check_connection = AsyncMock(return_value=True) + ai_service.rag_generator = MagicMock() + ai_service.rag_generator.generate_threats.side_effect = KeyError('\n "name"') + + class MockElement: + def __init__(self, name): + self.name = name + self.description = "" + self.stereotype = "Actor" + self.threats = [] + + actor = MockElement("Actor 1") + threat_model = MagicMock() + threat_model.actors = [{'object': actor}] + threat_model.servers = [] + threat_model.dataflows = [] + threat_model.tm.description = "System desc" + threat_model.tm.global_threats_llm = [] + + threat_payload = {"title": "SQLi", "description": "SQL injection", "category": "Information Disclosure"} + ai_service.provider.generate_threats_batch = AsyncMock(return_value={"Actor 1": [threat_payload]}) + + await ai_service._enrich_with_ai_threats(threat_model) + + assert len(actor.threats) == 1 + assert "SQLi" in actor.threats[0].description + asyncio.run(_run()) + def test_enrich_with_ai_threats_json_fallback(ai_service): async def _run(): ai_service.ai_online = True diff --git a/tests/test_attack_flow_generator.py b/tests/test_attack_flow_generator.py index 7ce66d6..be29df4 100644 --- a/tests/test_attack_flow_generator.py +++ b/tests/test_attack_flow_generator.py @@ -19,9 +19,9 @@ from threat_analysis.generation.attack_flow_generator import AttackFlowGenerator # Helper to create mock threats for testing -def create_mock_threat(tech_id, tech_name, tactics, stride, target, severity_score): +def create_mock_threat(tech_id, tech_name, tactics, stride, target, severity_score, threat_id=None): """Creates a simplified threat dictionary for testing purposes.""" - return { + threat = { "threat_name": f"{stride} against {target}", "description": f"A threat involving {tech_name}", "stride_category": stride, @@ -33,6 +33,9 @@ def create_mock_threat(tech_id, tech_name, tactics, stride, target, severity_sco "tactics": tactics }] } + if threat_id: + threat["id"] = threat_id + return threat @pytest.fixture def output_dir(tmp_path): @@ -129,7 +132,31 @@ class MockServerClass: # Mock a class object # Assert assert len(generator.threats) == 1 - assert generator.threats[0]["target"] == "User" + + +def test_allowed_categories_defaults_to_all_six_stride_categories(): + threats = [ + create_mock_threat("T1595", "Active Scanning", ["Reconnaissance"], "Spoofing", "Network", 3.0), + create_mock_threat("T1499", "Endpoint DoS", ["Impact"], "Denial of Service", "Server", 3.0), + ] + generator = AttackFlowGenerator(threats, model_name="DefaultCategoriesTest") + assert len(generator.threats) == 2 + + +def test_allowed_categories_narrows_to_configured_set(): + """attack_flows.generate_for_categories was previously dead config — when passed, + only threats in those STRIDE categories are considered for path-finding. + """ + threats = [ + create_mock_threat("T1595", "Active Scanning", ["Reconnaissance"], "Spoofing", "Network", 3.0), + create_mock_threat("T1499", "Endpoint DoS", ["Impact"], "Denial of Service", "Server", 3.0), + ] + generator = AttackFlowGenerator( + threats, model_name="NarrowedCategoriesTest", allowed_categories=["Spoofing"] + ) + assert len(generator.threats) == 1 + assert generator.threats[0]["stride_category"] == "Spoofing" + assert generator.threats[0]["target"] == "Network" def test_asset_hopping_logic(output_dir): """ @@ -161,6 +188,66 @@ def test_asset_hopping_logic(output_dir): if prop[0] == 'name' ), "The attack path should have moved to the new asset 'Workstation'." +# --------------------------------------------------------------------------- +# get_paths_summary — HTML rendering data, no .afb files written +# --------------------------------------------------------------------------- + +def test_get_paths_summary_no_threats_returns_empty_list(): + generator = AttackFlowGenerator([], model_name="EmptyTest") + assert generator.get_paths_summary() == [] + + +def test_get_paths_summary_matches_best_path_selection(output_dir): + """get_paths_summary() must pick the same best path per category as + generate_and_save_flows() — same underlying _compute_best_paths(). + """ + threats = [ + create_mock_threat("T1595", "Active Scanning", ["Reconnaissance"], "Spoofing", "Network", 3.0, "T-0001"), + create_mock_threat("T1485", "Data Destruction", ["Impact"], "Tampering", "File System", 9.0, "T-0002"), + create_mock_threat("T1055", "Process Injection", ["Impact"], "Tampering", "Database", 1.0, "T-0003"), + ] + generator = AttackFlowGenerator(threats, model_name="SummaryTest") + + summary = generator.get_paths_summary() + + assert len(summary) >= 1 + tampering_entry = next(p for p in summary if p["stride_category"] == "Tampering") + # The higher-severity Data Destruction (9.0) hop must win over Process Injection (1.0). + assert any(h["technique_id"] == "T1485" for h in tampering_entry["hops"]) + assert not any(h["technique_id"] == "T1055" for h in tampering_entry["hops"]) + assert "T-0002" in tampering_entry["threat_ids"] + assert tampering_entry["hop_count"] == len(tampering_entry["hops"]) + assert "File System" in tampering_entry["path_label"] + + +def test_get_paths_summary_hops_carry_their_own_threat_id(): + """Each hop must self-report its threat_id — narrative grounding (report_generator) + looks this up per hop and must not rely on threat_ids[i] positional alignment, + since that list silently skips threats with no id. + """ + threats = [ + create_mock_threat("T1485", "Data Destruction", ["Impact"], "Tampering", "File System", 9.0, "T-0001"), + ] + generator = AttackFlowGenerator(threats, model_name="HopIdTest") + + summary = generator.get_paths_summary() + + tampering_entry = next(p for p in summary if p["stride_category"] == "Tampering") + assert tampering_entry["hops"][0]["threat_id"] == "T-0001" + assert "threat_description" in tampering_entry["hops"][0] + + +def test_get_paths_summary_does_not_write_files(output_dir, monkeypatch): + """get_paths_summary() must be read-only — never create the afb/ directory.""" + monkeypatch.chdir(output_dir) + threats = [ + create_mock_threat("T1485", "Data Destruction", ["Impact"], "Tampering", "File System", 9.0, "T-0001"), + ] + generator = AttackFlowGenerator(threats, model_name="NoWriteTest") + generator.get_paths_summary() + assert not (output_dir / "afb").exists() + + def test_no_paths_found(output_dir, capsys): """ Tests the behavior when no logical paths can be constructed. diff --git a/tests/test_attack_id_validator.py b/tests/test_attack_id_validator.py index 1372889..38ecebe 100644 --- a/tests/test_attack_id_validator.py +++ b/tests/test_attack_id_validator.py @@ -293,6 +293,81 @@ def test_attack_url_subtechnique(self): # ValidationReport properties # --------------------------------------------------------------------------- +class TestCapecValidation: + """CAPEC IDs are validated against the union of every committed CAPEC source + (see data_loader.load_capec_catalog_ids) — there's no per-test override like + the STIX bundle, so these use known-real and known-fake IDs from that fixed + corpus. + """ + + def test_real_capec_id_no_issue(self, tmp_path): + AttackIdValidator._reset_cache() + v = AttackIdValidator() + threat = { + "id": "T-0001", "name": "x", "target": "y", + "capecs": [{"capec_id": "CAPEC-1"}], # real: "Accessing Functionality Not Properly Constrained by ACLs" + } + r = v.validate_all([threat], stix_path=tmp_path / "missing.json") + assert not any(i.technique_id == "CAPEC-1" for i in r.invalid) + + def test_capec_id_only_in_stride_to_capec_json_not_flagged(self, tmp_path): + """Regression test: CAPEC-166 is real (used by pytm-derived threats via + stride_to_capec.json) but absent from CAPEC_VIEW_ATT&CK_Related_Patterns.csv + (no ATT&CK mapping). Validating against the CSV alone previously flagged + it — and 41 other legitimate IDs — as hallucinated, which they were not. + """ + AttackIdValidator._reset_cache() + v = AttackIdValidator() + threat = { + "id": "T-0001", "name": "x", "target": "y", + "capecs": [{"capec_id": "CAPEC-166"}], + } + r = v.validate_all([threat], stix_path=tmp_path / "missing.json") + assert not any(i.technique_id == "CAPEC-166" for i in r.invalid) + + def test_hallucinated_capec_id_flagged_invalid(self, tmp_path): + AttackIdValidator._reset_cache() + v = AttackIdValidator() + threat = { + "id": "T-0001", "name": "x", "target": "y", + "capecs": [{"capec_id": "CAPEC-999999"}], # does not exist in the catalog + } + r = v.validate_all([threat], stix_path=tmp_path / "missing.json") + assert any(i.technique_id == "CAPEC-999999" and i.issue_type == INVALID for i in r.invalid) + + def test_capec_url_points_to_capec_mitre_org(self): + issue = IdIssue("CAPEC-66", INVALID, "T-0001", "x", "y") + assert issue.attack_url == "https://capec.mitre.org/data/definitions/66.html" + + +class TestD3fendValidation: + """D3FEND IDs are validated against the real, committed d3fend.csv.""" + + def test_real_d3fend_id_no_issue(self, tmp_path): + AttackIdValidator._reset_cache() + v = AttackIdValidator() + threat = { + "id": "T-0001", "name": "x", "target": "y", + "mitre_techniques": [{"id": "T1000", "defend_mitigations": [{"id": "D3-OAM"}]}], + } + r = v.validate_all([threat], stix_path=tmp_path / "missing.json") + assert not any(i.technique_id == "D3-OAM" for i in r.invalid) + + def test_fake_d3fend_id_flagged_invalid(self, tmp_path): + AttackIdValidator._reset_cache() + v = AttackIdValidator() + threat = { + "id": "T-0001", "name": "x", "target": "y", + "mitre_techniques": [{"id": "T1000", "defend_mitigations": [{"id": "D3-NOTREAL"}]}], + } + r = v.validate_all([threat], stix_path=tmp_path / "missing.json") + assert any(i.technique_id == "D3-NOTREAL" and i.issue_type == INVALID for i in r.invalid) + + def test_d3fend_url_points_to_d3fend_mitre_org(self): + issue = IdIssue("D3-OAM", INVALID, "T-0001", "x", "y") + assert issue.attack_url == "https://d3fend.mitre.org/technique/d3f:D3-OAM/" + + class TestValidationReport: def test_has_issues_false_when_empty(self): r = ValidationReport(total_techniques_checked=10) diff --git a/tests/test_attack_path_narrative.py b/tests/test_attack_path_narrative.py new file mode 100644 index 0000000..ad98ffe --- /dev/null +++ b/tests/test_attack_path_narrative.py @@ -0,0 +1,344 @@ +# Copyright 2025 ellipse2v +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the discovered-attack-path narrative AI pass +(ReportGenerator._generate_path_narratives) and its ID-leakage guard. +""" + +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from threat_analysis.ai_engine.providers.base_provider import BaseLLMProvider +from threat_analysis.ai_engine.providers.litellm_provider import LiteLLMProvider +from threat_analysis.generation.report_generator import _narrative_has_id_leakage + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_path(stride_category="Tampering", threat_id="T-0001"): + return { + "stride_category": stride_category, + "path_score": 6.5, + "hop_count": 2, + "hops": [ + { + "technique_id": "T1190", + "technique_name": "Exploit Public-Facing Application", + "tactic": "Initial Access", + "target": "WebApp", + "threat_id": threat_id, + "threat_description": "Unpatched web app accepts malicious input.", + }, + { + "technique_id": "T1485", + "technique_name": "Data Destruction", + "tactic": "Impact", + "target": "Database", + "threat_id": "T-0002", + "threat_description": "DB has no backup and is reachable from WebApp.", + }, + ], + "threat_ids": [threat_id, "T-0002"], + "path_label": "WebApp → Database", + } + + +def _make_report_generator(provider=None, include_narrative=True): + from threat_analysis.generation.report_generator import ReportGenerator + + rg = object.__new__(ReportGenerator) + rg.ai_provider = provider + rg._attack_flows_config = {"include_narrative": include_narrative} + return rg + + +def _make_provider(return_value: dict) -> MagicMock: + provider = MagicMock(spec=BaseLLMProvider) + provider.check_connection = AsyncMock(return_value=True) + provider.generate_attack_path_narrative = AsyncMock(return_value=return_value) + return provider + + +VALID_NARRATIVE = { + "narrative": "An attacker exploits the public-facing web app to pivot into the database.", + "business_impact": "Loss of customer data integrity and availability.", +} + + +# --------------------------------------------------------------------------- +# _narrative_has_id_leakage — pure function +# --------------------------------------------------------------------------- + +class TestNarrativeHasIdLeakage: + def test_clean_text_has_no_leakage(self): + assert not _narrative_has_id_leakage(VALID_NARRATIVE["narrative"]) + + def test_empty_text_has_no_leakage(self): + assert not _narrative_has_id_leakage("") + assert not _narrative_has_id_leakage(None) + + def test_attack_technique_id_detected(self): + assert _narrative_has_id_leakage("The attacker uses T1190 to gain access.") + + def test_attack_subtechnique_id_detected(self): + assert _narrative_has_id_leakage("Leverages T1078.001 for valid accounts.") + + def test_cve_id_detected(self): + assert _narrative_has_id_leakage("Exploiting CVE-2023-12345 on the server.") + + def test_capec_id_detected(self): + assert _narrative_has_id_leakage("This matches CAPEC-66 injection pattern.") + + def test_d3fend_id_detected(self): + assert _narrative_has_id_leakage("Mitigated by D3-NTA network analysis.") + + def test_plain_word_starting_with_t_is_not_a_false_positive(self): + assert not _narrative_has_id_leakage("The team escalates to management.") + + +# --------------------------------------------------------------------------- +# _generate_path_narratives — gating +# --------------------------------------------------------------------------- + +class TestGeneratePathNarrativesGating: + def test_no_ai_provider_skips(self): + rg = _make_report_generator(provider=None) + paths = [_make_path()] + asyncio.run(rg._generate_path_narratives(paths)) + assert "narrative" not in paths[0] + + def test_empty_paths_skips(self): + provider = _make_provider(VALID_NARRATIVE) + rg = _make_report_generator(provider=provider) + asyncio.run(rg._generate_path_narratives([])) + provider.generate_attack_path_narrative.assert_not_called() + + def test_include_narrative_disabled_skips(self): + provider = _make_provider(VALID_NARRATIVE) + rg = _make_report_generator(provider=provider, include_narrative=False) + paths = [_make_path()] + asyncio.run(rg._generate_path_narratives(paths)) + assert "narrative" not in paths[0] + provider.generate_attack_path_narrative.assert_not_called() + + def test_provider_offline_skips(self): + provider = MagicMock(spec=BaseLLMProvider) + provider.check_connection = AsyncMock(return_value=False) + rg = _make_report_generator(provider=provider) + paths = [_make_path()] + asyncio.run(rg._generate_path_narratives(paths)) + assert "narrative" not in paths[0] + provider.generate_attack_path_narrative.assert_not_called() + + +# --------------------------------------------------------------------------- +# _generate_path_narratives — success + grounding +# --------------------------------------------------------------------------- + +class TestGeneratePathNarrativesSuccess: + def test_valid_result_applied_to_path(self): + provider = _make_provider(VALID_NARRATIVE) + rg = _make_report_generator(provider=provider) + paths = [_make_path()] + + with patch("threat_analysis.ai_engine.prompt_loader.get", return_value="stub <>"): + asyncio.run(rg._generate_path_narratives(paths)) + + assert paths[0]["narrative"] == VALID_NARRATIVE["narrative"] + assert paths[0]["business_impact"] == VALID_NARRATIVE["business_impact"] + + def test_prompt_includes_hop_details_and_no_technique_ids(self): + """The grounding block gives technique names, not IDs — reduces the model's + raw material to echo back a malformed or hallucinated ID. + """ + captured = {} + + async def _capture(prompt, system_prompt): + captured["prompt"] = prompt + return VALID_NARRATIVE + + provider = MagicMock(spec=BaseLLMProvider) + provider.check_connection = AsyncMock(return_value=True) + provider.generate_attack_path_narrative = _capture + rg = _make_report_generator(provider=provider) + paths = [_make_path()] + + def _fake_prompt(section, key): + if key == "template": + return "<> <> <> <> <>" + return "system" + + with patch("threat_analysis.ai_engine.prompt_loader.get", side_effect=_fake_prompt): + asyncio.run(rg._generate_path_narratives(paths)) + + prompt = captured["prompt"] + assert "Exploit Public-Facing Application" in prompt + assert "WebApp" in prompt + assert "T1190" not in prompt # technique IDs are not fed into the grounding text + + def test_multiple_paths_each_get_a_call(self): + provider = _make_provider(VALID_NARRATIVE) + rg = _make_report_generator(provider=provider) + paths = [_make_path("Tampering"), _make_path("Spoofing")] + + with patch("threat_analysis.ai_engine.prompt_loader.get", return_value="stub"): + asyncio.run(rg._generate_path_narratives(paths)) + + assert provider.generate_attack_path_narrative.await_count == 2 + assert all("narrative" in p for p in paths) + + +# --------------------------------------------------------------------------- +# _generate_path_narratives — hallucination guard +# --------------------------------------------------------------------------- + +class TestGeneratePathNarrativesGroundingViolation: + def test_narrative_with_technique_id_is_discarded(self): + provider = _make_provider({ + "narrative": "Attacker uses T1190 to breach the app.", + "business_impact": "Data loss.", + }) + rg = _make_report_generator(provider=provider) + paths = [_make_path()] + + with patch("threat_analysis.ai_engine.prompt_loader.get", return_value="stub"): + asyncio.run(rg._generate_path_narratives(paths)) + + assert "narrative" not in paths[0] + assert "business_impact" not in paths[0] + + def test_business_impact_with_cve_is_discarded(self): + provider = _make_provider({ + "narrative": "A realistic pivot from the app to the database.", + "business_impact": "Exploits CVE-2024-00001 for full compromise.", + }) + rg = _make_report_generator(provider=provider) + paths = [_make_path()] + + with patch("threat_analysis.ai_engine.prompt_loader.get", return_value="stub"): + asyncio.run(rg._generate_path_narratives(paths)) + + assert "narrative" not in paths[0] + + def test_non_dict_result_ignored(self): + provider = MagicMock(spec=BaseLLMProvider) + provider.check_connection = AsyncMock(return_value=True) + provider.generate_attack_path_narrative = AsyncMock(return_value=[]) + rg = _make_report_generator(provider=provider) + paths = [_make_path()] + + with patch("threat_analysis.ai_engine.prompt_loader.get", return_value="stub"): + asyncio.run(rg._generate_path_narratives(paths)) + + assert "narrative" not in paths[0] + + def test_empty_narrative_and_impact_ignored(self): + provider = _make_provider({"narrative": "", "business_impact": ""}) + rg = _make_report_generator(provider=provider) + paths = [_make_path()] + + with patch("threat_analysis.ai_engine.prompt_loader.get", return_value="stub"): + asyncio.run(rg._generate_path_narratives(paths)) + + assert "narrative" not in paths[0] + + def test_provider_exception_is_non_fatal_and_other_paths_still_processed(self): + provider = MagicMock(spec=BaseLLMProvider) + provider.check_connection = AsyncMock(return_value=True) + provider.generate_attack_path_narrative = AsyncMock( + side_effect=[RuntimeError("LLM error"), VALID_NARRATIVE] + ) + rg = _make_report_generator(provider=provider) + paths = [_make_path("Tampering"), _make_path("Spoofing")] + + with patch("threat_analysis.ai_engine.prompt_loader.get", return_value="stub"): + asyncio.run(rg._generate_path_narratives(paths)) + + assert "narrative" not in paths[0] + assert paths[1]["narrative"] == VALID_NARRATIVE["narrative"] + + def test_missing_prompt_key_skips_all_paths(self): + provider = _make_provider(VALID_NARRATIVE) + rg = _make_report_generator(provider=provider) + paths = [_make_path()] + + with patch("threat_analysis.ai_engine.prompt_loader.get", side_effect=KeyError("attack_path_narrative")): + asyncio.run(rg._generate_path_narratives(paths)) + + assert "narrative" not in paths[0] + provider.generate_attack_path_narrative.assert_not_called() + + +# --------------------------------------------------------------------------- +# BaseLLMProvider default +# --------------------------------------------------------------------------- + +class TestBaseLLMProviderDefault: + def test_default_returns_empty_dict(self): + class _ConcreteProvider(BaseLLMProvider): + async def check_connection(self): return True + async def generate_threats(self, c, ctx): return [] + async def generate_markdown(self, p, m=None): + yield "" + + provider = _ConcreteProvider() + result = asyncio.run(provider.generate_attack_path_narrative("prompt", "system")) + assert result == {} + + +# --------------------------------------------------------------------------- +# LiteLLMProvider.generate_attack_path_narrative +# --------------------------------------------------------------------------- + +class TestLiteLLMProviderAttackPathNarrative: + def test_returns_dict_from_client(self): + provider = object.__new__(LiteLLMProvider) + provider._client = None + provider._config = {} + + async def _fake_generate_content(prompt, system_prompt, output_format, **kw): + yield VALID_NARRATIVE + + mock_client = MagicMock() + mock_client.generate_content = _fake_generate_content + + async def _run(): + provider._client = mock_client + return await provider.generate_attack_path_narrative("prompt", "system") + + result = asyncio.run(_run()) + assert result["narrative"] == VALID_NARRATIVE["narrative"] + + def test_client_exception_returns_empty(self): + provider = object.__new__(LiteLLMProvider) + provider._client = None + provider._config = {} + + async def _fake_generate_content(*a, **kw): + raise RuntimeError("timeout") + yield # make it an async generator + + mock_client = MagicMock() + mock_client.generate_content = _fake_generate_content + + async def _run(): + provider._client = mock_client + return await provider.generate_attack_path_narrative("prompt", "system") + + result = asyncio.run(_run()) + assert result == {} diff --git a/tests/test_ciso_triage.py b/tests/test_ciso_triage.py index da2abd8..d9a0513 100644 --- a/tests/test_ciso_triage.py +++ b/tests/test_ciso_triage.py @@ -94,7 +94,6 @@ def test_default_returns_empty_dict(self): class _ConcreteProvider(BaseLLMProvider): async def check_connection(self): return True async def generate_threats(self, c, ctx): return [] - async def generate_attack_flow(self, t, c, ctx): return {} async def generate_markdown(self, p, m=None): yield "" @@ -173,6 +172,93 @@ async def _capture_triage(prompt, system_prompt): lines = [l for l in captured.get("prompt", "").splitlines() if l.startswith("- [")] assert len(lines) <= 20 + def test_prompt_has_no_scenarios_message_when_gdaf_empty(self): + """No GDAF scenarios provided -> the prompt says so explicitly rather than + silently omitting the section (the LLM must not assume scenarios exist).""" + captured = {} + + async def _capture_triage(prompt, system_prompt): + captured["prompt"] = prompt + return VALID_TRIAGE + + provider = MagicMock(spec=BaseLLMProvider) + provider.check_connection = AsyncMock(return_value=True) + provider.generate_ciso_triage = _capture_triage + + with patch("threat_analysis.ai_engine.prompt_loader.get", return_value="<>"): + rg = _make_report_generator(provider=provider) + asyncio.run(rg._run_ciso_triage([_make_threat()], gdaf_scenarios=None, debate_results=None)) + + assert "No GDAF attack scenarios were generated" in captured["prompt"] + + def test_prompt_contains_gdaf_scenario_summary(self): + """GDAF scenarios provided -> the prompt lists them (id, objective, actor, risk, score).""" + scenario = MagicMock() + scenario.scenario_id = "GDAF-ABC1" + scenario.objective_name = "Exfiltrate customer data" + scenario.actor_name = "External Attacker" + scenario.risk_level = "CRITICAL" + scenario.path_score = 4.2 + scenario.path_score_pre_debate = None + captured = {} + + async def _capture_triage(prompt, system_prompt): + captured["prompt"] = prompt + return VALID_TRIAGE + + provider = MagicMock(spec=BaseLLMProvider) + provider.check_connection = AsyncMock(return_value=True) + provider.generate_ciso_triage = _capture_triage + + with patch("threat_analysis.ai_engine.prompt_loader.get", return_value="<>"): + rg = _make_report_generator(provider=provider) + asyncio.run(rg._run_ciso_triage([_make_threat()], gdaf_scenarios=[scenario], debate_results=[])) + + prompt = captured["prompt"] + assert "GDAF-ABC1" in prompt + assert "Exfiltrate customer data" in prompt + assert "CRITICAL" in prompt + assert "4.20" in prompt + + def test_prompt_includes_debate_outcome_for_debated_scenario(self): + """A scenario with a matching debate result must show the pre-debate score and + whether the residual path stayed viable — the whole point of enriching CISO + triage with debate context. + """ + scenario = MagicMock() + scenario.scenario_id = "GDAF-ABC1" + scenario.objective_name = "Exfiltrate customer data" + scenario.actor_name = "External Attacker" + scenario.risk_level = "MEDIUM" + scenario.path_score = 2.1 + scenario.path_score_pre_debate = 4.2 + + debate_result = MagicMock() + debate_result.scenario_id = "GDAF-ABC1" + debate_result.residual_path_viable = False + debate_result.round_count = 2 + + captured = {} + + async def _capture_triage(prompt, system_prompt): + captured["prompt"] = prompt + return VALID_TRIAGE + + provider = MagicMock(spec=BaseLLMProvider) + provider.check_connection = AsyncMock(return_value=True) + provider.generate_ciso_triage = _capture_triage + + with patch("threat_analysis.ai_engine.prompt_loader.get", return_value="<>"): + rg = _make_report_generator(provider=provider) + asyncio.run(rg._run_ciso_triage( + [_make_threat()], gdaf_scenarios=[scenario], debate_results=[debate_result] + )) + + prompt = captured["prompt"] + assert "was 4.20 pre-debate" in prompt + assert "largely blocked by Blue" in prompt + assert "2 round(s)" in prompt + # --------------------------------------------------------------------------- # _run_ciso_triage — result handling diff --git a/tests/test_data_loader.py b/tests/test_data_loader.py index ddacd20..72434cd 100644 --- a/tests/test_data_loader.py +++ b/tests/test_data_loader.py @@ -21,6 +21,7 @@ from threat_analysis.core.data_loader import ( load_attack_techniques, load_capec_to_mitre_mapping, + load_capec_catalog_ids, load_stride_to_capec_map, load_d3fend_mapping, load_nist_mappings, @@ -63,6 +64,52 @@ def test_load_capec_to_mitre_mapping_invalid_entry(): mapping = load_capec_to_mitre_mapping() assert mapping == {} +def test_load_capec_catalog_ids_file_not_found(): + with patch("builtins.open", side_effect=FileNotFoundError): + ids = load_capec_catalog_ids() + assert ids == frozenset() + +def test_load_capec_catalog_ids_parses_id_column(): + csv_data = "'ID,Name,Abstraction\n1,Some Pattern,Standard\n66,SQL Injection,Standard\n" + with patch("builtins.open", mock_open(read_data=csv_data)): + ids = load_capec_catalog_ids() + assert ids == frozenset({"CAPEC-1", "CAPEC-66"}) + +def test_load_capec_catalog_ids_skips_blank_id_rows(): + csv_data = "'ID,Name,Abstraction\n1,Some Pattern,Standard\n,Empty ID Row,Standard\n" + with patch("builtins.open", mock_open(read_data=csv_data)): + ids = load_capec_catalog_ids() + assert ids == frozenset({"CAPEC-1"}) + +def test_load_capec_catalog_ids_unions_stride_and_mitre_mapping_sources(): + """The CSV (CAPEC_VIEW_ATT&CK_Related_Patterns.csv) is only a narrow + ATT&CK-linked view of CAPEC, not the master catalog — stride_to_capec.json + and capec_to_mitre_structured_mapping.json must also contribute IDs, or a + real, legitimate CAPEC ID that pytm-derived threats reference (but which + happens to have no ATT&CK mapping) gets wrongly flagged as invalid. + """ + csv_data = "'ID,Name,Abstraction\n1,Some Pattern,Standard\n" + stride_json = json.dumps({ + "Tampering": [{"capec_id": "CAPEC-166", "name": "Force the System to Reset Values"}] + }) + mitre_json = json.dumps([{ + "capec_id": "CAPEC-519", + "techniques": [{"taxonomy": "ATT&CK", "id": "T1565"}], + }]) + + def _fake_open(path, *args, **kwargs): + path_str = str(path) + if "stride_to_capec" in path_str: + return mock_open(read_data=stride_json)() + if "capec_to_mitre_structured_mapping" in path_str: + return mock_open(read_data=mitre_json)() + return mock_open(read_data=csv_data)() + + with patch("builtins.open", side_effect=_fake_open): + ids = load_capec_catalog_ids() + + assert ids == frozenset({"CAPEC-1", "CAPEC-166", "CAPEC-519"}) + def test_load_stride_to_capec_map_file_not_found(): with patch("builtins.open", side_effect=FileNotFoundError): mapping = load_stride_to_capec_map() diff --git a/tests/test_debate_engine.py b/tests/test_debate_engine.py new file mode 100644 index 0000000..743193f --- /dev/null +++ b/tests/test_debate_engine.py @@ -0,0 +1,413 @@ +# Copyright 2025 ellipse2v +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for threat_analysis/core/debate_engine.py""" + +import asyncio +from unittest.mock import MagicMock, AsyncMock + +from threat_analysis.core.debate_engine import RedBlueDebateEngine +from threat_analysis.core.gdaf_engine import AttackScenario, AttackHop +from threat_analysis.core.asset_technique_mapper import ScoredTechnique + + +def make_scenario(path_score=4.5, detection_coverage=0.2, risk_level="CRITICAL"): + hop = AttackHop( + asset_name="WebServer", + asset_type="web_server", + techniques=[ScoredTechnique( + id="T1190", name="Exploit Public-Facing Application", + tactics=["initial-access"], score=1.5, rationale="", + )], + dataflow_name="UserToWeb", + protocol="https", + is_encrypted=True, + is_authenticated=False, + hop_score=1.8, + hop_position="entry", + ) + return AttackScenario( + scenario_id="GDAF-TEST01", + objective_id="obj1", + objective_name="Exfiltrate customer data", + objective_description="", + objective_business_impact="High", + objective_mitre_final_tactic="exfiltration", + actor_id="actor1", + actor_name="Organized Crime", + actor_sophistication="high", + entry_point="ExternalUser", + target_asset="WebServer", + hops=[hop], + path_score=path_score, + risk_level=risk_level, + detection_coverage=detection_coverage, + unacceptable_risk=True, + ) + + +class FakeProvider: + """Returns queued canned responses in call order.""" + def __init__(self, responses): + self._responses = list(responses) + self.calls = [] + + async def generate_debate_turn(self, prompt, system_prompt): + self.calls.append((prompt, system_prompt)) + if not self._responses: + return {} + return self._responses.pop(0) + + +RED_TURN_HIGH = { + "viability_score": 0.8, + "techniques_attempted": ["T1190"], + "failed_alternatives": [], + "rationale": "Exploiting unauthenticated HTTPS endpoint", + "evidence": [ + {"claim": "endpoint has no auth", "evidence_type": "misconfig", + "evidence_ref": "hop 1 is_authenticated=false", "confidence": "high"}, + ], +} +BLUE_TURN_BLOCK_ALL = { + "viability_score": 0.1, + "techniques_blocked": ["T1190"], + "detection_gaps": [ + {"step": "WebServer entry", "control_family": "EDR", "covered": True, + "detail": "EDR rule X blocks T1190 payloads", "confidence": "high"}, + ], + "rationale": "Blocked by EDR rule", + "evidence": [], +} +BLUE_TURN_GAP = { + "viability_score": 0.6, + "techniques_blocked": [], + "detection_gaps": [ + {"step": "WebServer entry", "control_family": "SIEM", "covered": False, + "detail": "No rule for this payload", "confidence": "medium"}, + ], + "rationale": "No coverage", + "evidence": [], +} +RED_TURN_ALT = { + "viability_score": 0.3, + "techniques_attempted": ["T1210"], + "failed_alternatives": [], + "rationale": "Switching to a different lateral movement technique", + "evidence": [], +} + + +def test_selects_only_scenarios_above_viability_threshold(): + low = make_scenario(path_score=0.5, detection_coverage=0.9, risk_level="LOW") + high = make_scenario(path_score=4.5, detection_coverage=0.0, risk_level="CRITICAL") + engine = RedBlueDebateEngine(FakeProvider([]), config={"min_viability_threshold": 0.5}) + selected = engine._select_scenarios([low, high]) + assert selected == [high] + + +def test_top_n_caps_selection(): + scenarios = [make_scenario(path_score=4.0 + i, detection_coverage=0.0) for i in range(10)] + engine = RedBlueDebateEngine(FakeProvider([]), config={"top_n": 3, "min_viability_threshold": 0.0}) + selected = engine._select_scenarios(scenarios) + assert len(selected) == 3 + assert selected[0].path_score == max(s.path_score for s in scenarios) + + +def test_convergence_stops_before_max_rounds(): + scenario = make_scenario() + provider = FakeProvider([RED_TURN_HIGH, BLUE_TURN_BLOCK_ALL]) + engine = RedBlueDebateEngine(provider, config={"max_rounds": 3, "viability_delta_threshold": 0.5}) + result = asyncio.run(engine._debate_scenario(scenario)) + assert result is not None + assert result.converged is True + assert len(result.rounds) == 2 + + +def test_max_rounds_respected_when_not_converging(): + scenario = make_scenario() + # Viability must genuinely keep moving round to round — identical canned turns would + # converge on repetition (delta hits an exact 0.0 fixed point), which is not what this + # test is exercising. Alternating RED_TURN_HIGH/RED_TURN_ALT keeps the delta above + # threshold every round. + provider = FakeProvider([ + RED_TURN_HIGH, BLUE_TURN_GAP, + RED_TURN_ALT, BLUE_TURN_GAP, + RED_TURN_HIGH, BLUE_TURN_GAP, + ]) + engine = RedBlueDebateEngine(provider, config={"max_rounds": 3, "viability_delta_threshold": 0.001}) + result = asyncio.run(engine._debate_scenario(scenario)) + assert result is not None + assert len(result.rounds) == 6 + assert result.converged is False + + +def test_malformed_turn_json_skipped_gracefully(): + scenario = make_scenario() + provider = FakeProvider([{"unexpected": "shape"}]) + engine = RedBlueDebateEngine(provider, config={"max_rounds": 3}) + result = asyncio.run(engine._debate_scenario(scenario)) + assert result is None + + +def test_debate_factor_bounded_and_risk_level_recomputed(): + scenario = make_scenario(path_score=4.5, risk_level="CRITICAL") + provider = FakeProvider([RED_TURN_HIGH, BLUE_TURN_BLOCK_ALL]) + # BLUE_TURN_BLOCK_ALL only blocks 1 technique, so the round viability (0.8 - 0.15*1 = 0.65) + # stays above the engine's default min_viability_threshold (0.5) — residual_path_viable would + # be True and the >1.0 cap would never engage. Raising the threshold to 0.7 for this test + # makes 0.65 register as "not viable," which is what "Blue blocked everything -> capped at + # 1.0" is meant to exercise. + engine = RedBlueDebateEngine( + provider, + config={"max_rounds": 1, "debate_factor_min": 0.5, "debate_factor_max": 1.5, + "viability_delta_threshold": 0.5, "min_viability_threshold": 0.7}, + ) + results = asyncio.run(engine.run([scenario])) + assert len(results) == 1 + assert scenario.path_score_pre_debate == 4.5 + assert 0.5 <= scenario.debate_factor <= 1.5 + assert scenario.path_score <= 4.5 # Blue blocked everything -> factor capped at 1.0 + assert scenario.risk_level in {"CRITICAL", "HIGH", "MEDIUM", "LOW"} + + +def test_evidence_without_ref_marked_unverified(): + scenario = make_scenario() + turn_with_bad_evidence = dict(RED_TURN_HIGH) + turn_with_bad_evidence["evidence"] = [ + {"claim": "unsupported claim", "evidence_type": "none", "evidence_ref": "", "confidence": "low"}, + ] + provider = FakeProvider([turn_with_bad_evidence, BLUE_TURN_BLOCK_ALL]) + engine = RedBlueDebateEngine(provider, config={"max_rounds": 1}) + result = asyncio.run(engine._debate_scenario(scenario)) + assert result is not None + red_turn = result.rounds[0] + assert red_turn.evidence[0].verified is False + + +def test_grounding_includes_hop_facts_and_bom_cves(tmp_path): + bom_dir = tmp_path / "BOM" + bom_dir.mkdir() + (bom_dir / "webserver.yaml").write_text("known_cves:\n - CVE-2024-1234\n", encoding="utf-8") + scenario = make_scenario() + engine = RedBlueDebateEngine(FakeProvider([]), bom_directory=str(bom_dir)) + grounding = engine._build_grounding(scenario) + assert "unauthenticated" in grounding + assert "https" in grounding + assert "CVE-2024-1234" in grounding + + +from threat_analysis.generation.report_generator import ReportGenerator + + +def test_run_debate_disabled_by_default(): + rg = ReportGenerator(MagicMock(), MagicMock()) + rg.ai_provider = MagicMock() + rg._debate_config = {} + threat_model = MagicMock() + threat_model.gdaf_scenarios = [make_scenario()] + result = asyncio.run(rg._run_debate(threat_model)) + assert result == [] + + +def test_run_debate_no_provider(): + rg = ReportGenerator(MagicMock(), MagicMock()) + rg.ai_provider = None + rg._debate_config = {"enabled": True} + threat_model = MagicMock() + threat_model.gdaf_scenarios = [make_scenario()] + result = asyncio.run(rg._run_debate(threat_model)) + assert result == [] + + +def test_run_debate_no_scenarios(): + rg = ReportGenerator(MagicMock(), MagicMock()) + rg.ai_provider = MagicMock() + rg._debate_config = {"enabled": True} + threat_model = MagicMock() + threat_model.gdaf_scenarios = [] + result = asyncio.run(rg._run_debate(threat_model)) + assert result == [] + + +def test_run_debate_runs_engine_when_enabled_and_online(): + rg = ReportGenerator(MagicMock(), MagicMock()) + mock_client = MagicMock() + mock_client.ai_online = True + rg.ai_provider = MagicMock() + rg.ai_provider._get_client = AsyncMock(return_value=mock_client) + rg.ai_provider.generate_debate_turn = AsyncMock(side_effect=[RED_TURN_HIGH, BLUE_TURN_BLOCK_ALL]) + rg._debate_config = {"enabled": True, "max_rounds": 1, "viability_delta_threshold": 0.5} + threat_model = MagicMock() + threat_model.gdaf_scenarios = [make_scenario()] + threat_model.context_config = {} + threat_model._model_file_path = None + result = asyncio.run(rg._run_debate(threat_model)) + assert len(result) == 1 + assert result[0].scenario_id == "GDAF-TEST01" + + +import re +import jinja2 +from pathlib import Path +from unittest.mock import mock_open, patch + +_TEMPLATE_PATH = Path(__file__).resolve().parents[1] / "threat_analysis" / "templates" / "report_template.html" + + +def _extract_debate_section() -> str: + text = _TEMPLATE_PATH.read_text(encoding="utf-8") + match = re.search( + r"(.*?)", + text, re.DOTALL, + ) + assert match, "DEBATE_SECTION markers not found in report_template.html" + return match.group(1) + + +def test_debate_section_hidden_when_no_results(): + tmpl = jinja2.Template(_extract_debate_section()) + html = tmpl.render(debate_results=[]) + assert "Red/Blue Adversarial Debate" not in html + + +def test_debate_section_shows_persuasion_order(): + tmpl = jinja2.Template(_extract_debate_section()) + debate_results = [{ + "objective_name": "Exfiltrate data", + "entry_point": "ExternalUser", + "target_asset": "WebServer", + "residual_path_viable": True, + "final_viability": 0.6, + "blocked_paths": ["T1078"], + "residual_detection_gaps": [ + {"step": "hop1", "control_family": "SIEM", "covered": False, + "detail": "no rule", "confidence": "medium"}, + ], + "red_failed_attempts": ["Tried T1110 — blocked"], + "round_count": 2, + "converged": True, + "convergence_delta": 0.05, + "debate_factor": 1.2, + "rounds": [], + }] + html = tmpl.render(debate_results=debate_results) + assert "Red/Blue Adversarial Debate" in html + idx_blocked = html.index("Paths Blue successfully blocked") + idx_gaps = html.index("Detection gaps on the residual path") + idx_failed = html.index("Red's failed alternative attempts") + idx_footnote = html.index("round(s)") + assert idx_blocked < idx_gaps < idx_failed < idx_footnote + + +def _debate_ready_report_generator(): + """A ReportGenerator wired to actually run the debate (1 round, deterministic turns).""" + rg = ReportGenerator(MagicMock(), MagicMock()) + rg.severity_calculator.get_severity_info.return_value = {"level": "High", "score": 8.0} + rg.mitre_mapping.map_threat_to_mitre.return_value = {"techniques": [], "capecs": []} + + mock_client = MagicMock() + mock_client.ai_online = True + rg.ai_provider = MagicMock() + rg.ai_provider._get_client = AsyncMock(return_value=mock_client) + rg.ai_provider.generate_debate_turn = AsyncMock(side_effect=[RED_TURN_HIGH, BLUE_TURN_BLOCK_ALL]) + rg._debate_config = {"enabled": True, "max_rounds": 1, "viability_delta_threshold": 0.5} + return rg + + +def _make_gdaf_threat_model(gdaf_scenarios): + threat_model = MagicMock() + threat_model.mitre_analysis_results = { + "total_threats": 0, "mitre_techniques_count": 0, "stride_distribution": {} + } + threat_model.tm.name = "Test Architecture" + threat_model.gdaf_scenarios = gdaf_scenarios + threat_model.context_config = {} + threat_model._model_file_path = None + return threat_model + + +def test_afb_files_rewritten_after_debate_changes_scores(): + """generate_html_report() must re-write the .afb Attack Flow files after the debate + runs, so they reflect debate-adjusted scores instead of the pre-debate ones written + earlier by run_gdaf_engine() (called before generate_html_report in every caller). + """ + rg = _debate_ready_report_generator() + scenario = make_scenario() + threat_model = _make_gdaf_threat_model([scenario]) + + with patch.object(rg.env, "get_template") as mock_get_template, \ + patch("threat_analysis.generation.report_generator.get_framework_mitigation_suggestions", return_value=[]), \ + patch("threat_analysis.generation.attack_flow_builder.AttackFlowBuilder") as mock_builder_cls: + mock_get_template.return_value = MagicMock() + with patch("builtins.open", mock_open()): + rg.generate_html_report(threat_model, {}, "output/run1/test_report.html") + + mock_builder_cls.assert_called_once_with(threat_model.gdaf_scenarios, model_name="Test Architecture") + mock_builder_cls.return_value.generate_and_save.assert_called_once_with("output/run1") + + +def test_afb_files_not_rewritten_when_debate_disabled(): + """No debate results (disabled, offline, or no scenarios) → no redundant .afb rewrite.""" + rg = ReportGenerator(MagicMock(), MagicMock()) + rg.severity_calculator.get_severity_info.return_value = {"level": "High", "score": 8.0} + rg.mitre_mapping.map_threat_to_mitre.return_value = {"techniques": [], "capecs": []} + rg.ai_provider = None + rg._debate_config = {"enabled": False} + + threat_model = _make_gdaf_threat_model([make_scenario()]) + + with patch.object(rg.env, "get_template") as mock_get_template, \ + patch("threat_analysis.generation.report_generator.get_framework_mitigation_suggestions", return_value=[]), \ + patch("threat_analysis.generation.attack_flow_builder.AttackFlowBuilder") as mock_builder_cls: + mock_get_template.return_value = MagicMock() + with patch("builtins.open", mock_open()): + rg.generate_html_report(threat_model, {}, "output/run1/test_report.html") + + mock_builder_cls.assert_not_called() + + +def test_threat_graph_gdaf_paths_reflect_debate_outcome(): + """The threat graph passed to the template must include a gdaf_paths entry for the + debated scenario, built AFTER the debate ran (so residual_path_viable reflects it — + not the pre-debate default). + """ + from types import SimpleNamespace + + rg = _debate_ready_report_generator() + scenario = make_scenario() + threat_model = _make_gdaf_threat_model([scenario]) + # Needed so _build_threat_graph_data() produces non-empty nodes (matching the + # scenario's entry_point/hop asset names) instead of short-circuiting to {}. + threat_model.actors = [{"name": "ExternalUser", "object": SimpleNamespace(name="ExternalUser")}] + threat_model.servers = [{"name": "WebServer", "object": SimpleNamespace(name="WebServer")}] + + with patch.object(rg.env, "get_template") as mock_get_template, \ + patch("threat_analysis.generation.report_generator.get_framework_mitigation_suggestions", return_value=[]), \ + patch("threat_analysis.generation.attack_flow_builder.AttackFlowBuilder"): + mock_template = MagicMock() + mock_get_template.return_value = mock_template + with patch("builtins.open", mock_open()): + rg.generate_html_report(threat_model, {}, "output/run1/test_report.html") + + render_kwargs = mock_template.render.call_args.kwargs + threat_graph = render_kwargs["threat_graph"] + assert len(threat_graph["gdaf_paths"]) == 1 + path = threat_graph["gdaf_paths"][0] + assert path["scenario_id"] == "GDAF-TEST01" + assert path["nodes"] == ["ExternalUser", "WebServer"] + assert path["debated"] is True + # Must reflect the actual DebateResult computed by _run_debate(), not a pre-debate default. + debate_result = threat_model.debate_results[0] + assert path["residual_path_viable"] == debate_result.residual_path_viable diff --git a/tests/test_diagram_generator.py b/tests/test_diagram_generator.py index bb59c68..fae9680 100644 --- a/tests/test_diagram_generator.py +++ b/tests/test_diagram_generator.py @@ -66,6 +66,22 @@ def test_get_element_name(diagram_generator): assert diagram_generator._get_element_name("String Name") == "String Name" assert diagram_generator._get_element_name(None) is None +def test_get_element_name_unstringable_element_logs_warning(diagram_generator, caplog): + """The last-resort str(element) fallback must log why it failed, not silently + return None — a bare `except:` there previously masked real extraction bugs + (and could swallow KeyboardInterrupt/SystemExit). + """ + class Unstringable: + def __str__(self): + raise RuntimeError("boom") + + with caplog.at_level("WARNING"): + result = diagram_generator._get_element_name(Unstringable()) + + assert result is None + assert "Could not convert element to string" in caplog.text + assert "boom" in caplog.text + def test_extract_data_info(diagram_generator): # Test with single Data object mock_data_single = MagicMock(spec=[]) @@ -408,6 +424,30 @@ def test_generate_diagram_from_dot_subprocess_error(diagram_generator): result = diagram_generator.generate_diagram_from_dot("digraph G {}", "output", "svg") assert result is None +def test_generate_diagram_from_dot_subprocess_error_uses_unique_tempfile(diagram_generator, tmp_path, caplog): + """The Graphviz stderr dump must go to a unique tempfile.mkstemp() path, not a + fixed /tmp/graphviz_error.log — a fixed world-writable path lets concurrent + requests clobber each other's dump and is vulnerable to a symlink attack on a + shared host. + """ + import tempfile as tempfile_module + with patch('subprocess.run') as mock_run: + mock_run.side_effect = subprocess.CalledProcessError(1, 'dot', stderr="Graphviz error detail") + with patch.object(diagram_generator, 'check_graphviz_installation', return_value=True): + with patch.object(tempfile_module, 'mkstemp', wraps=tempfile_module.mkstemp) as mock_mkstemp: + with caplog.at_level("ERROR"): + # "png" (not "svg") to hit the standard `dot` subprocess path + # directly — "svg" first tries CustomSVGGenerator, which has its + # own separate subprocess/error handling. + result = diagram_generator.generate_diagram_from_dot("digraph G {}", str(tmp_path / "output"), "png") + + assert result is None + mock_mkstemp.assert_called_once() + call_kwargs = mock_mkstemp.call_args.kwargs + assert call_kwargs.get("prefix") == "graphviz_error_" + assert "/tmp/graphviz_error.log" not in caplog.text + assert "Graphviz error detail" in caplog.text + def test_generate_diagram_from_dot_output_file_not_created(diagram_generator): with patch('subprocess.run') as mock_run: mock_run.return_value = MagicMock(returncode=0) diff --git a/tests/test_export_service.py b/tests/test_export_service.py index 6463d44..e073160 100644 --- a/tests/test_export_service.py +++ b/tests/test_export_service.py @@ -68,7 +68,7 @@ def test_export_files_logic_markdown(mock_create, mock_validator, export_service with patch('os.makedirs'): with patch('builtins.open', MagicMock()): path, filename = export_service.export_files_logic("md", "markdown") - assert filename == "threat_model.md" + assert filename == "system_model.md" @patch('threat_analysis.server.export_service.ModelValidator') @patch('threat_analysis.server.export_service.create_threat_model') @@ -88,6 +88,70 @@ def test_generate_full_project_export_single(mock_create, mock_validator, export assert "reports" in res assert res["reports"]["html"] == "stride_mitre_report.html" +@patch('threat_analysis.server.export_service.ModelValidator') +@patch('threat_analysis.server.export_service.create_threat_model') +def test_generate_full_project_export_html_runs_before_json_and_exports(mock_create, mock_validator, export_service): + """generate_html_report() must run before generate_json_export()/Navigator/STIX/ + Attack Flow — it's what populates threat_model._report_all_detailed_threats with the + AI-enriched, debate-aware threat list those exports should reuse (get_enriched_threats). + Previously JSON/Navigator/STIX/Attack Flow ran BEFORE the HTML report in this single- + file path, so the cache never existed yet and they silently fell back to pytm-only. + """ + tm = MagicMock() + tm.tm.name = "TestModel" + tm._report_all_detailed_threats = [{"id": "T-0001", "source": "AI"}] + mock_create.return_value = tm + mock_validator.return_value.validate.return_value = [] + + with patch('pathlib.Path.write_text'): + with patch('threat_analysis.server.export_service.AttackNavigatorGenerator') as mock_nav: + mock_nav.return_value.save_layer_to_file.return_value = None + with patch('threat_analysis.server.export_service.StixGenerator') as mock_stix: + mock_stix.return_value.generate_stix_bundle.return_value = {"type": "bundle"} + with patch('threat_analysis.server.export_service.AttackFlowGenerator') as mock_afg: + export_service.generate_full_project_export("md", Path("/tmp")) + + call_order = [c[0] for c in export_service.report_generator.method_calls] + assert call_order.index("generate_html_report") < call_order.index("generate_json_export") + + # Navigator/STIX/AttackFlowGenerator must receive the cached AI-enriched list. + mock_nav.assert_called_once_with(threat_model_name="TestModel", all_detailed_threats=tm._report_all_detailed_threats) + mock_stix.assert_called_once_with(threat_model=tm, all_detailed_threats=tm._report_all_detailed_threats) + mock_afg.assert_called_once_with(threats=tm._report_all_detailed_threats, model_name="TestModel") + + +@patch('threat_analysis.server.export_service.ModelValidator') +@patch('threat_analysis.server.export_service.create_threat_model') +def test_generate_full_project_export_diagram_html_built_after_html_report(mock_create, mock_validator, export_service): + """The diagram's severity heat map (tm_diagram.html, via _generate_html_with_legend) + must be built AFTER generate_html_report() — that's what attaches AI-enriched + threats to element.threats, which _compute_severity_map() reads. Building the + diagram HTML earlier would permanently miss AI-derived severities. + """ + tm = MagicMock() + tm.tm.name = "TestModel" + tm._report_all_detailed_threats = [{"id": "T-0001", "source": "AI"}] + mock_create.return_value = tm + mock_validator.return_value.validate.return_value = [] + + call_order = [] + export_service.report_generator.generate_html_report.side_effect = ( + lambda *a, **kw: call_order.append("generate_html_report") + ) + export_service.diagram_generator._generate_html_with_legend.side_effect = ( + lambda *a, **kw: call_order.append("_generate_html_with_legend") + ) + + with patch('pathlib.Path.write_text'): + with patch('threat_analysis.server.export_service.AttackNavigatorGenerator'): + with patch('threat_analysis.server.export_service.StixGenerator') as mock_stix: + mock_stix.return_value.generate_stix_bundle.return_value = {"type": "bundle"} + with patch('threat_analysis.server.export_service.AttackFlowGenerator'): + export_service.generate_full_project_export("md", Path("/tmp")) + + assert call_order == ["generate_html_report", "_generate_html_with_legend"] + + def test_export_all_files_logic(export_service): with patch.object(export_service, 'generate_full_project_export'): with patch('threat_analysis.server.export_service.create_threat_model') as mock_create: diff --git a/tests/test_gdaf_engine.py b/tests/test_gdaf_engine.py index 0df063a..45c3462 100644 --- a/tests/test_gdaf_engine.py +++ b/tests/test_gdaf_engine.py @@ -23,6 +23,7 @@ GDAFEngine, AttackScenario, AttackHop, + compute_risk_level, _CLASSIFICATION_SCORE, _TRAVERSAL_BONUS, _DETECTION_COVERAGE, @@ -1445,3 +1446,38 @@ def test_bfs_handles_disconnected_graph(self): graph = engine._build_graph() paths = engine._bfs_paths(graph, "AttackerA", "ServerB", max_hops=10) assert paths == [] + + +# --------------------------------------------------------------------------- +# compute_risk_level function and debate fields +# --------------------------------------------------------------------------- + +class TestComputeRiskLevel: + def test_compute_risk_level_thresholds(self): + thresholds = {"CRITICAL": 4.0, "HIGH": 2.8, "MEDIUM": 1.8} + assert compute_risk_level(4.5, thresholds) == "CRITICAL" + assert compute_risk_level(4.0, thresholds) == "CRITICAL" + assert compute_risk_level(3.0, thresholds) == "HIGH" + assert compute_risk_level(2.0, thresholds) == "MEDIUM" + assert compute_risk_level(1.0, thresholds) == "LOW" + + +class TestGDAFEngineGetRiskThresholds: + def test_get_risk_thresholds_defaults(self): + GDAFEngine._scoring_config = None # reset class-level cache + with patch.object(GDAFEngine, "_load_scoring_config", return_value={}): + thresholds = GDAFEngine.get_risk_thresholds() + assert thresholds == {"CRITICAL": 4.0, "HIGH": 2.8, "MEDIUM": 1.8} + + +class TestAttackScenarioDebateFields: + def test_attack_scenario_debate_fields_default(self): + scenario = AttackScenario( + scenario_id="s1", objective_id="o1", objective_name="Obj", + objective_description="", objective_business_impact="", objective_mitre_final_tactic="", + actor_id="a1", actor_name="Actor", actor_sophistication="medium", + entry_point="Entry", target_asset="Target", hops=[], path_score=3.0, + risk_level="HIGH", detection_coverage=0.5, unacceptable_risk=False, + ) + assert scenario.path_score_pre_debate is None + assert scenario.debate_factor == 1.0 diff --git a/tests/test_generation_utils.py b/tests/test_generation_utils.py index 5406d6c..08294e5 100644 --- a/tests/test_generation_utils.py +++ b/tests/test_generation_utils.py @@ -15,7 +15,8 @@ """Tests for threat_analysis/generation/utils.py""" import pytest -from threat_analysis.generation.utils import extract_name_from_object, get_target_name +from types import SimpleNamespace +from threat_analysis.generation.utils import extract_name_from_object, get_target_name, get_enriched_threats # --------------------------------------------------------------------------- @@ -138,3 +139,38 @@ def test_empty_tuple(self): # () is not None, not string, has no .name → "Unspecified" result = get_target_name(()) assert result == "Unspecified" + + +# --------------------------------------------------------------------------- +# get_enriched_threats +# --------------------------------------------------------------------------- + +class TestGetEnrichedThreats: + def test_prefers_cached_list_when_present(self): + cached = [{"id": "T-0001", "source": "AI"}] + threat_model = SimpleNamespace( + _report_all_detailed_threats=cached, + get_all_threats_details=lambda: (_ for _ in ()).throw(AssertionError("should not be called")), + ) + assert get_enriched_threats(threat_model) is cached + + def test_falls_back_when_cache_missing(self): + pytm_only = [{"id": "T-0001", "source": "pytm"}] + threat_model = SimpleNamespace(get_all_threats_details=lambda: pytm_only) + assert get_enriched_threats(threat_model) is pytm_only + + def test_falls_back_when_cache_is_none(self): + pytm_only = [{"id": "T-0001", "source": "pytm"}] + threat_model = SimpleNamespace( + _report_all_detailed_threats=None, + get_all_threats_details=lambda: pytm_only, + ) + assert get_enriched_threats(threat_model) is pytm_only + + def test_falls_back_when_cache_is_empty_list(self): + pytm_only = [{"id": "T-0001", "source": "pytm"}] + threat_model = SimpleNamespace( + _report_all_detailed_threats=[], + get_all_threats_details=lambda: pytm_only, + ) + assert get_enriched_threats(threat_model) is pytm_only diff --git a/tests/test_litellm_provider.py b/tests/test_litellm_provider.py index a70e786..0873f8c 100644 --- a/tests/test_litellm_provider.py +++ b/tests/test_litellm_provider.py @@ -182,8 +182,12 @@ async def _run(): # --- enable_thinking config tests --- -def test_enable_thinking_defaults_to_false(): - """enable_thinking must be False when not set in provider config.""" +def test_enable_thinking_omitted_when_not_configured(): + """enable_thinking must be absent from the request when not set in provider config. + + Some providers (e.g. Anthropic) reject unknown parameters outright, so the key + must only be sent when a provider block explicitly opts in. + """ async def _run(): client = LiteLLMClient() client.ai_online = True @@ -201,7 +205,7 @@ async def _run(): pass call_kwargs = client._litellm_module.acompletion.call_args[1] - assert call_kwargs["enable_thinking"] is False + assert "enable_thinking" not in call_kwargs asyncio.run(_run()) @@ -272,16 +276,50 @@ async def mock_gen(**kwargs): assert result == [{"id": "T1"}] asyncio.run(_run()) -def test_litellm_provider_generate_attack_flow(): +def test_litellm_provider_generate_debate_turn(): + async def _run(): + with patch("threat_analysis.ai_engine.providers.litellm_client.LiteLLMClient.create", new_callable=AsyncMock) as mock_create: + mock_client = MagicMock() + async def mock_gen(**kwargs): + yield {"viability_score": 0.6, "techniques_attempted": ["T1190"]} + mock_client.generate_content = mock_gen + mock_create.return_value = mock_client + + provider = LiteLLMProvider({}) + result = await provider.generate_debate_turn("prompt", "system") + assert result == {"viability_score": 0.6, "techniques_attempted": ["T1190"]} + asyncio.run(_run()) + +def test_litellm_provider_generate_soc_analysis_success(): + async def _run(): + with patch("threat_analysis.ai_engine.providers.litellm_client.LiteLLMClient.create", new_callable=AsyncMock) as mock_create: + mock_client = MagicMock() + async def mock_gen(**kwargs): + yield [{"threat_id": "t-0", "detectability": "medium"}] + mock_client.generate_content = mock_gen + mock_create.return_value = mock_client + + provider = LiteLLMProvider({}) + result = await provider.generate_soc_analysis("prompt", "system") + assert result == [{"threat_id": "t-0", "detectability": "medium"}] + asyncio.run(_run()) + +def test_litellm_provider_generate_soc_analysis_logs_error_string(caplog): + """A JSON-parsing failure (truncated response, etc.) must be logged, not silently + swallowed into an empty list indistinguishable from 'nothing to analyze'. + """ async def _run(): with patch("threat_analysis.ai_engine.providers.litellm_client.LiteLLMClient.create", new_callable=AsyncMock) as mock_create: mock_client = MagicMock() async def mock_gen(**kwargs): - yield {"flow": "steps"} + yield "Error: No valid JSON found in AI response (finish_reason=length)." mock_client.generate_content = mock_gen mock_create.return_value = mock_client provider = LiteLLMProvider({}) - result = await provider.generate_attack_flow({}, {}, {}) - assert result == {"flow": "steps"} + with caplog.at_level("WARNING"): + result = await provider.generate_soc_analysis("prompt", "system") + assert result == [] + assert "generate_soc_analysis" in caplog.text + assert "No valid JSON found" in caplog.text asyncio.run(_run()) diff --git a/tests/test_main.py b/tests/test_main.py index e69de29..a197cb6 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -0,0 +1,71 @@ +# Copyright 2025 ellipse2v +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for threat_analysis/__main__.py""" + +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch + +from threat_analysis.__main__ import _build_project_ai_service + + +class TestBuildProjectAiService: + """secopstm --project mode must build an AIService so cross-model RAG threat + generation (Decision A2) runs — previously only the web UI (export_service.py) + passed ai_service to generate_project_reports(), so the CLI silently skipped it. + """ + + def test_builds_and_initializes_ai_service(self): + mock_ai_service_instance = MagicMock() + mock_ai_service_instance.init_ai = AsyncMock() + mock_ai_service_cls = MagicMock(return_value=mock_ai_service_instance) + mock_module = MagicMock(AIService=mock_ai_service_cls) + + with patch("importlib.import_module", return_value=mock_module): + result = _build_project_ai_service(Path("/some/ai_config.yaml")) + + mock_ai_service_cls.assert_called_once_with( + config_path=str(Path("/some/ai_config.yaml")), force_disable_rag=False + ) + mock_ai_service_instance.init_ai.assert_awaited_once() + assert result is mock_ai_service_instance + + def test_passes_force_disable_rag_through(self): + mock_ai_service_instance = MagicMock() + mock_ai_service_instance.init_ai = AsyncMock() + mock_ai_service_cls = MagicMock(return_value=mock_ai_service_instance) + mock_module = MagicMock(AIService=mock_ai_service_cls) + + with patch("importlib.import_module", return_value=mock_module): + _build_project_ai_service(Path("/some/ai_config.yaml"), force_disable_rag=True) + + mock_ai_service_cls.assert_called_once_with( + config_path=str(Path("/some/ai_config.yaml")), force_disable_rag=True + ) + + def test_returns_none_on_import_failure(self): + with patch("importlib.import_module", side_effect=ImportError("no module")): + result = _build_project_ai_service(Path("/some/ai_config.yaml")) + assert result is None + + def test_returns_none_on_init_ai_failure(self): + mock_ai_service_instance = MagicMock() + mock_ai_service_instance.init_ai = AsyncMock(side_effect=RuntimeError("boom")) + mock_ai_service_cls = MagicMock(return_value=mock_ai_service_instance) + mock_module = MagicMock(AIService=mock_ai_service_cls) + + with patch("importlib.import_module", return_value=mock_module): + result = _build_project_ai_service(Path("/some/ai_config.yaml")) + + assert result is None diff --git a/tests/test_mitre_mapping_module.py b/tests/test_mitre_mapping_module.py index bd60b7a..093cb7a 100644 --- a/tests/test_mitre_mapping_module.py +++ b/tests/test_mitre_mapping_module.py @@ -58,6 +58,40 @@ def test_map_threat_to_mitre_spoofing(mitre_mapping_instance): technique_ids = [t['id'] for t in techniques] assert "T1566" in technique_ids +def test_map_threat_to_mitre_caps_capecs_without_explicit_ids(mitre_mapping_instance): + """A threat with no explicit capec_ids must not get the entire STRIDE-category CAPEC + bucket (Tampering alone has ~149 entries) — only the ones relevant to its description, + capped at _CAPEC_RELEVANCE_MAX. + """ + full_bucket_size = len(mitre_mapping_instance.stride_to_capec.get("Tampering", [])) + assert full_bucket_size > MitreMapping._CAPEC_RELEVANCE_MAX, "fixture data too small to exercise the cap" + + threat = { + "description": "Untrusted actor may attempt to inject malicious data.", + "stride_category": "Tampering", + } + mapping_results = mitre_mapping_instance.map_threat_to_mitre(threat) + capecs = mapping_results["capecs"] + + assert 0 < len(capecs) <= MitreMapping._CAPEC_RELEVANCE_MAX + assert len(capecs) < full_bucket_size + + +def test_rank_capecs_by_relevance_prefers_keyword_overlap(): + """Ranking must prefer CAPECs whose own description shares words with the threat + description, not just return the list truncated in its original order. + """ + capec_list = [ + {"capec_id": "CAPEC-1", "description": "Buffer Overflow via Environment Variables"}, + {"capec_id": "CAPEC-2", "description": "Injection of malicious data into network traffic"}, + {"capec_id": "CAPEC-3", "description": "Unrelated pattern about physical access"}, + ] + ranked = MitreMapping._rank_capecs_by_relevance( + "An attacker may inject malicious data over the network.", capec_list, max_results=1 + ) + assert ranked == [capec_list[1]] + + def test_get_stride_categories(mitre_mapping_instance): """Test the get_stride_categories method.""" categories = mitre_mapping_instance.get_stride_categories() diff --git a/tests/test_model_factory.py b/tests/test_model_factory.py new file mode 100644 index 0000000..fb238fb --- /dev/null +++ b/tests/test_model_factory.py @@ -0,0 +1,119 @@ +# Copyright 2025 ellipse2v +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for threat_analysis/core/model_factory.py""" + +from unittest.mock import MagicMock, patch + +import pytest + +from threat_analysis.core.model_factory import create_threat_model + + +@pytest.fixture +def cve_service(): + return MagicMock() + + +@patch('threat_analysis.core.model_factory.ModelValidator') +@patch('threat_analysis.core.model_factory.ModelParser') +@patch('threat_analysis.core.model_factory.ThreatModel') +def test_create_threat_model_success_without_validation(mock_tm_cls, mock_parser_cls, mock_validator_cls, cve_service): + mock_tm = MagicMock() + mock_tm_cls.return_value = mock_tm + + result = create_threat_model( + markdown_content="# System Model: X\n", model_name="X", + model_description="desc", cve_service=cve_service, validate=False, + ) + + assert result is mock_tm + mock_parser_cls.return_value.parse_markdown.assert_called_once_with("# System Model: X\n") + mock_validator_cls.assert_not_called() + + +@patch('threat_analysis.core.model_factory.ModelValidator') +@patch('threat_analysis.core.model_factory.ModelParser') +@patch('threat_analysis.core.model_factory.ThreatModel') +def test_create_threat_model_success_with_validation_passing(mock_tm_cls, mock_parser_cls, mock_validator_cls, cve_service): + mock_tm = MagicMock() + mock_tm_cls.return_value = mock_tm + mock_validator_cls.return_value.validate.return_value = [] + + result = create_threat_model( + markdown_content="# System Model: X\n", model_name="X", + model_description="desc", cve_service=cve_service, validate=True, + ) + + assert result is mock_tm + mock_validator_cls.return_value.validate.assert_called_once() + + +@patch('threat_analysis.core.model_factory.ModelValidator') +@patch('threat_analysis.core.model_factory.ModelParser') +@patch('threat_analysis.core.model_factory.ThreatModel') +def test_create_threat_model_validation_failure_returns_none(mock_tm_cls, mock_parser_cls, mock_validator_cls, cve_service): + mock_tm_cls.return_value = MagicMock() + mock_validator_cls.return_value.validate.return_value = ["Dataflow points to unknown element"] + + result = create_threat_model( + markdown_content="# System Model: X\n", model_name="X", + model_description="desc", cve_service=cve_service, validate=True, + ) + + assert result is None + + +@patch('threat_analysis.core.model_factory.ModelParser') +@patch('threat_analysis.core.model_factory.ThreatModel') +def test_create_threat_model_parse_exception_returns_none(mock_tm_cls, mock_parser_cls, cve_service): + mock_tm_cls.return_value = MagicMock() + mock_parser_cls.return_value.parse_markdown.side_effect = RuntimeError("malformed DSL") + + result = create_threat_model( + markdown_content="not valid", model_name="X", + model_description="desc", cve_service=cve_service, validate=False, + ) + + assert result is None + + +@patch('threat_analysis.core.model_factory.ModelValidator') +@patch('threat_analysis.core.model_factory.ModelParser') +@patch('threat_analysis.core.model_factory.ThreatModel') +def test_create_threat_model_sets_model_file_path_when_given(mock_tm_cls, mock_parser_cls, mock_validator_cls, cve_service): + mock_tm = MagicMock() + mock_tm_cls.return_value = mock_tm + + create_threat_model( + markdown_content="# System Model: X\n", model_name="X", + model_description="desc", cve_service=cve_service, validate=False, + model_file_path="/path/to/model.md", + ) + + assert mock_tm._model_file_path == "/path/to/model.md" + + +@patch('threat_analysis.core.model_factory.ModelValidator') +@patch('threat_analysis.core.model_factory.ModelParser') +@patch('threat_analysis.core.model_factory.ThreatModel') +def test_create_threat_model_constructs_with_given_name_description_and_cve_service( + mock_tm_cls, mock_parser_cls, mock_validator_cls, cve_service +): + create_threat_model( + markdown_content="# System Model: X\n", model_name="MyModel", + model_description="My description", cve_service=cve_service, validate=False, + ) + + mock_tm_cls.assert_called_once_with("MyModel", "My description", cve_service=cve_service) diff --git a/tests/test_models_module.py b/tests/test_models_module.py index f5a5f4e..c751811 100644 --- a/tests/test_models_module.py +++ b/tests/test_models_module.py @@ -442,6 +442,19 @@ def test_add_severity_multiplier(threat_model_instance): tm.add_severity_multiplier("WebServer", 1.5) assert tm.severity_multipliers["WebServer"] == 1.5 + +def test_add_severity_multiplier_syncs_severity_calculator(threat_model_instance): + """add_severity_multiplier() must keep tm.severity_calculator.target_multipliers in + sync — this calculator is used for custom-rule threat scoring (_apply_custom_threats) + and starts empty (constructed before parsing populates severity_multipliers), so + without this sync it would never see the model's own multipliers. + """ + tm = threat_model_instance[0] + tm.add_severity_multiplier("WebServer", 1.5) + # severity_calculator is mocked by the fixture (patches models_module.SeverityCalculator), + # so target_multipliers is itself a MagicMock — assert the item-assignment happened. + tm.severity_calculator.target_multipliers.__setitem__.assert_called_once_with("WebServer", 1.5) + # --- ThreatModel add_custom_mitre_mapping Tests --- def test_add_custom_mitre_mapping(threat_model_instance): diff --git a/tests/test_ollama_provider.py b/tests/test_ollama_provider.py index 780c96c..8da2af6 100644 --- a/tests/test_ollama_provider.py +++ b/tests/test_ollama_provider.py @@ -94,33 +94,3 @@ async def _run(): assert threats == [] asyncio.run(_run()) -def test_ollama_provider_generate_attack_flow(provider): - async def _run(): - with patch("aiohttp.ClientSession.post") as mock_post: - # Success case - mock_response = AsyncMock() - mock_response.status = 200 - mock_response.raise_for_status = MagicMock() - mock_response.text.return_value = json.dumps({ - "response": json.dumps({ - "type": "attack-flow", - "name": "Test Flow" - }) - }) - mock_post.return_value.__aenter__.return_value = mock_response - - flow = await provider.generate_attack_flow({}, {}, {}) - assert flow["name"] == "Test Flow" - - # Error case - ClientError - from aiohttp import ClientError - mock_post.side_effect = ClientError("Post failed") - flow = await provider.generate_attack_flow({}, {}, {}) - assert flow == {} - - # Error case - JSONDecodeError - mock_post.side_effect = None - mock_response.text.return_value = "invalid json" - flow = await provider.generate_attack_flow({}, {}, {}) - assert flow == {} - asyncio.run(_run()) diff --git a/tests/test_prompts.py b/tests/test_prompts.py index 7fe4890..2b222d7 100644 --- a/tests/test_prompts.py +++ b/tests/test_prompts.py @@ -14,7 +14,6 @@ import pytest from threat_analysis.ai_engine.prompts.stride_prompts import build_component_prompt, STRIDE_SYSTEM_PROMPT -from threat_analysis.ai_engine.prompts.attack_flow_prompts import build_attack_flow_prompt, ATTACK_FLOW_SYSTEM_PROMPT def test_stride_prompts(): component = { @@ -65,46 +64,3 @@ def test_stride_prompts_defaults(): assert "No" in prompt # Internet facing default False assert "None specified" in prompt assert "None" in prompt # Integrations - -def test_attack_flow_prompts(): - threat = { - 'title': 'SQL Injection', - 'category': 'Information Disclosure', - 'description': 'Attacker injects SQL code', - 'attack_scenario': """1. Find input field -2. Inject code -3. Dump DB""", - 'mitre_techniques': ['T1190'] - } - component = { - 'type': 'Database', - 'name': 'UserDB', - 'description': 'Stores user data' - } - context = { - 'system_description': 'Backend system' - } - - prompt = build_attack_flow_prompt(threat, component, context) - - assert "SQL Injection" in prompt - assert "Information Disclosure" in prompt - assert "Attacker injects SQL code" in prompt - assert "Find input field" in prompt - assert "T1190" in prompt - assert "Database" in prompt - assert "UserDB" in prompt - assert "Stores user data" in prompt - assert "Backend system" in prompt - assert ATTACK_FLOW_SYSTEM_PROMPT is not None - -def test_attack_flow_prompts_defaults(): - threat = {} - component = {} - context = {} - - prompt = build_attack_flow_prompt(threat, component, context) - - assert "Unknown" in prompt - assert "No additional context" in prompt - assert "None" in prompt # mitre_techniques.join diff --git a/tests/test_rag_functionality.py b/tests/test_rag_functionality.py index cac21db..e9ec40b 100644 --- a/tests/test_rag_functionality.py +++ b/tests/test_rag_functionality.py @@ -184,6 +184,43 @@ def test_rag_threat_generator_initialization(mock_chat_litellm): # LLM model was configured (uses litellm.completion directly) assert rag_generator._llm_model is not None assert "ollama" in rag_generator._llm_model + # No max_tokens in the test config -> must not be injected (backward compat). + assert "max_tokens" not in rag_generator._llm_params + + +def test_rag_threat_generator_passes_max_tokens_when_configured(mock_chat_litellm): + """RAG calls litellm.completion() directly, bypassing LiteLLMClient entirely — it + must still honour the provider's max_tokens (already carefully tuned to avoid JSON + truncation for component-level generation, see ai_config.yaml comments) instead of + silently falling back to litellm's own per-model default. + """ + max_tokens_config_path = TEST_DIR / "ai_config_max_tokens.yaml" + max_tokens_config_path.write_text( + "ai_providers:\n" + " ollama:\n" + " enabled: true\n" + " model: \"test-llama3\"\n" + " host: \"http://localhost:11434\"\n" + " temperature: 0.5\n" + " max_tokens: 16000\n" + "embedding:\n" + " provider: huggingface\n" + " model: \"test-embedding-model\"\n" + " device: cpu\n" + "rag:\n" + " enabled: true\n", + encoding="utf-8", + ) + try: + with patch('os.path.exists', side_effect=lambda x: x == str(TEST_VECTOR_STORE_DIR) or x == str(max_tokens_config_path) or x == str(TEST_USER_CONTEXT_PATH)): + rag_generator = RAGThreatGenerator( + vector_store_dir=str(TEST_VECTOR_STORE_DIR), + user_context_path=str(TEST_USER_CONTEXT_PATH), + ai_config_path=str(max_tokens_config_path) + ) + assert rag_generator._llm_params.get("max_tokens") == 16000 + finally: + max_tokens_config_path.unlink(missing_ok=True) def test_rag_threat_generator_generates_threats(mock_chat_litellm): @@ -232,6 +269,45 @@ def test_rag_threat_generator_generates_threats(mock_chat_litellm): mock_litellm.completion.assert_called_once() +def test_rag_threat_generator_malformed_template_returns_empty_list(mock_chat_litellm): + """A human_template with an unescaped {{ / }} JSON example (like the real + prompts.yaml bug this regression-tests) must make generate_threats() return [] + instead of raising KeyError out of str.format() — that KeyError previously + propagated all the way up and aborted the entire AI enrichment pipeline (RAG is + only step 1 of 5), not just the RAG threats. + """ + with patch('os.path.exists', side_effect=lambda x: x == str(TEST_VECTOR_STORE_DIR) or x == str(TEST_AI_CONFIG_PATH) or x == str(TEST_USER_CONTEXT_PATH)): + with patch.object(RAGThreatGenerator, '_initialize_components'): + rag_generator = RAGThreatGenerator( + vector_store_dir=str(TEST_VECTOR_STORE_DIR), + user_context_path=str(TEST_USER_CONTEXT_PATH), + ai_config_path=str(TEST_AI_CONFIG_PATH) + ) + mock_collection = MagicMock() + mock_collection.query.return_value = {"documents": [["context chunk 1"]]} + rag_generator.collection = mock_collection + mock_embeddings = MagicMock() + mock_embeddings.embed_query.return_value = [0.1, 0.2, 0.3] + rag_generator.embeddings = mock_embeddings + rag_generator._llm_model = "ollama/test-model" + rag_generator._llm_params = {"temperature": 0.5} + rag_generator._rag_system_prompt = "You are a security expert." + # Unescaped literal JSON example — the exact bug pattern from prompts.yaml. + rag_generator._rag_human_template = ( + "{optional_context}\n" + "Model: {threat_model_markdown}\n" + "Context: {context}\n" + "Format:\n[\n {\n \"name\": \"...\"\n }\n]" + ) + + call_count_before = mock_litellm.completion.call_count + threats = rag_generator.generate_threats("This is a test threat model.") + + assert threats == [] + # Must fail before the LLM call — no point calling it with a broken prompt. + assert mock_litellm.completion.call_count == call_count_before + + def test_aiservice_integrates_rag_threats(mock_provider, mock_chat_litellm): """Test AIService successfully integrates RAG-generated threats.""" async def _run(): diff --git a/tests/test_report_generator.py b/tests/test_report_generator.py index 282e4da..fc90dc9 100644 --- a/tests/test_report_generator.py +++ b/tests/test_report_generator.py @@ -88,6 +88,129 @@ def test_generate_html_report(mock_get_framework_mitigations, report_generator): assert result == output_file mock_template.render.assert_called_once() + +@patch('threat_analysis.generation.report_generator.get_framework_mitigation_suggestions') +def test_generate_html_report_assigns_stable_threat_ids(mock_get_framework_mitigations, report_generator): + """generate_html_report() must assign T-0001-style ids onto all_detailed_threats + before CISO triage / threat graph build, so those consumers (and the later JSON/STIX + export, which now reuses this same cached list) all reference the same ids. + """ + threat_model = MagicMock() + threat_model.mitre_analysis_results = { + 'total_threats': 2, 'mitre_techniques_count': 0, 'stride_distribution': {} + } + threat_model.tm.name = "Test Architecture" + + def _mock_threat(desc): + m = MagicMock(description=desc, stride_category='S', + target=MagicMock(data=MagicMock(classification=MagicMock(name='Public')))) + if hasattr(m, 'mitigations'): + del m.mitigations + return m + + grouped_threats = { + 'Spoofing': [ + (_mock_threat("First threat"), MagicMock(name="Target A")), + (_mock_threat("Second threat"), MagicMock(name="Target B")), + ] + } + + report_generator.severity_calculator.get_severity_info.return_value = {'level': 'High', 'score': 8.0} + report_generator.mitre_mapping.map_threat_to_mitre.return_value = {'techniques': [], 'capecs': []} + mock_get_framework_mitigations.return_value = [] + + with patch.object(report_generator.env, 'get_template') as mock_get_template: + mock_get_template.return_value = MagicMock() + with patch("builtins.open", mock_open()): + report_generator.generate_html_report(threat_model, grouped_threats, "test_report.html") + + cached = threat_model._report_all_detailed_threats + assert [t["id"] for t in cached] == ["T-0001", "T-0002"] + + +@patch('threat_analysis.generation.report_generator.get_framework_mitigation_suggestions') +def test_generate_html_report_computes_discovered_attack_paths(mock_get_framework_mitigations, report_generator): + """generate_html_report() must compute discovered_attack_paths (AttackFlowGenerator's + best-path-per-STRIDE-category, independent of GDAF) and pass it to the template — + this is the new 'Automatically Discovered Attack Paths' HTML section. + """ + threat_model = MagicMock() + threat_model.mitre_analysis_results = { + 'total_threats': 1, 'mitre_techniques_count': 1, 'stride_distribution': {} + } + threat_model.tm.name = "Test Architecture" + + threat_mock = MagicMock(description="Data destruction", stride_category='Tampering', + target=MagicMock(data=MagicMock(classification=MagicMock(name='Public')))) + if hasattr(threat_mock, 'mitigations'): + del threat_mock.mitigations + + grouped_threats = { + 'Tampering': [ + (threat_mock, MagicMock(name="FileServer")), + ] + } + + report_generator.severity_calculator.get_severity_info.return_value = {'level': 'HIGH', 'score': 9.0} + report_generator.mitre_mapping.map_threat_to_mitre.return_value = { + 'techniques': [{'id': 'T1485', 'name': 'Data Destruction', 'tactics': ['Impact'], + 'defend_mitigations': [], 'mitre_mitigations': [], + 'owasp_mitigations': [], 'nist_mitigations': [], 'cis_mitigations': []}], + 'capecs': [], + } + mock_get_framework_mitigations.return_value = [] + + with patch.object(report_generator.env, 'get_template') as mock_get_template: + mock_template = MagicMock() + mock_get_template.return_value = mock_template + with patch("builtins.open", mock_open()): + report_generator.generate_html_report(threat_model, grouped_threats, "test_report.html") + + render_kwargs = mock_template.render.call_args.kwargs + paths = render_kwargs["discovered_attack_paths"] + assert len(paths) == 1 + assert paths[0]["stride_category"] == "Tampering" + assert paths[0]["hops"][0]["technique_id"] == "T1485" + assert "FileServer" in paths[0]["path_label"] + + +@patch('threat_analysis.generation.report_generator.get_framework_mitigation_suggestions') +def test_generate_html_report_skips_discovered_attack_paths_when_disabled(mock_get_framework_mitigations, report_generator): + """attack_flows.enabled: false was previously dead config — must actually skip + discovered-path computation. + """ + threat_model = MagicMock() + threat_model.mitre_analysis_results = { + 'total_threats': 1, 'mitre_techniques_count': 1, 'stride_distribution': {} + } + threat_model.tm.name = "Test Architecture" + + threat_mock = MagicMock(description="Data destruction", stride_category='Tampering', + target=MagicMock(data=MagicMock(classification=MagicMock(name='Public')))) + if hasattr(threat_mock, 'mitigations'): + del threat_mock.mitigations + + grouped_threats = {'Tampering': [(threat_mock, MagicMock(name="FileServer"))]} + + report_generator.severity_calculator.get_severity_info.return_value = {'level': 'HIGH', 'score': 9.0} + report_generator.mitre_mapping.map_threat_to_mitre.return_value = { + 'techniques': [{'id': 'T1485', 'name': 'Data Destruction', 'tactics': ['Impact'], + 'defend_mitigations': [], 'mitre_mitigations': [], + 'owasp_mitigations': [], 'nist_mitigations': [], 'cis_mitigations': []}], + 'capecs': [], + } + mock_get_framework_mitigations.return_value = [] + report_generator._attack_flows_config = {"enabled": False} + + with patch.object(report_generator.env, 'get_template') as mock_get_template: + mock_template = MagicMock() + mock_get_template.return_value = mock_template + with patch("builtins.open", mock_open()): + report_generator.generate_html_report(threat_model, grouped_threats, "test_report.html") + + assert mock_template.render.call_args.kwargs["discovered_attack_paths"] == [] + + def test_generate_json_export(report_generator): threat_model = MagicMock() threat_model.tm.name = "Test Architecture" @@ -135,6 +258,38 @@ def __init__(self, description, stride_category): assert result == output_file + +def test_generate_json_export_reuses_ai_enriched_cache(report_generator): + """generate_json_export must reuse threat_model._report_all_detailed_threats (set by + generate_html_report after AI enrichment) instead of recomputing pytm-only threats — + otherwise AI/LLM-sourced threats present in the HTML report silently vanish from the + JSON/STIX exports. + """ + threat_model = MagicMock() + threat_model.tm.name = "Test Architecture" + + cached_ai_threat = { + "id": "T-0001", + "description": "AI-enriched threat", + "stride_category": "Spoofing", + "target": "TestTarget", + "source": "AI", + "severity": {"level": "High", "score": 8.0}, + "confidence": 0.9, + "capec_ids": [], + "mitigations": [], + } + threat_model._report_all_detailed_threats = [cached_ai_threat] + + with patch.object(report_generator, "_get_all_threats_with_mitre_info") as mock_recompute, \ + patch("threat_analysis.generation.report_generator.ReportSerializer") as mock_serializer, \ + patch("builtins.open", mock_open()): + mock_serializer.serialize.return_value = {"threats": [cached_ai_threat]} + report_generator.generate_json_export(threat_model, {}, "test_export.json") + + mock_recompute.assert_not_called() + mock_serializer.serialize.assert_called_once_with(threat_model, [cached_ai_threat]) + def test_get_all_threats_with_mitre_info_handles_missing_url_friendly_name_source(report_generator): threat_model = MagicMock() threat_model.mitre_analysis_results = { diff --git a/tests/test_report_generator_extended.py b/tests/test_report_generator_extended.py index 290f3c9..9a8e916 100644 --- a/tests/test_report_generator_extended.py +++ b/tests/test_report_generator_extended.py @@ -18,7 +18,7 @@ import json import os from pathlib import Path -from unittest.mock import MagicMock, patch, mock_open +from unittest.mock import MagicMock, patch, mock_open, AsyncMock as _StdAsyncMock from threat_analysis.generation.report_generator import ( ReportGenerator, _is_network_exposed, @@ -79,47 +79,93 @@ def report_generator(): mitre_mapping = MagicMock() return ReportGenerator(severity_calc, mitre_mapping) -def test_enrich_threats_with_ai_no_provider(report_generator): +def test_run_ai_enrichment_no_provider(report_generator): + """No ai_provider -> _get_ai_service() must not even be attempted.""" report_generator.ai_provider = None - threats = [{"id": 1}] - result = asyncio.run(report_generator._enrich_threats_with_ai(None, threats)) - assert result == threats + with patch.object(report_generator, "_get_ai_service") as mock_get_service: + asyncio.run(report_generator._run_ai_enrichment(MagicMock())) + mock_get_service.assert_not_called() -def test_enrich_threats_with_ai_not_reachable(report_generator): - mock_client = MagicMock() - mock_client.ai_online = False - report_generator.ai_provider = MagicMock() - report_generator.ai_provider._get_client = AsyncMock(return_value=mock_client) - report_generator.ai_context = {"ctx": 1} - threats = [{"id": 1}] - result = asyncio.run(report_generator._enrich_threats_with_ai(None, threats)) - assert result == threats -def test_enrich_threats_with_ai_success(report_generator): +def test_run_ai_enrichment_service_unavailable(report_generator): + """AIService init failed (_get_ai_service returns None) -> enrichment is a no-op, + not an error — matches the AI degrade-silently contract used everywhere else. + """ report_generator.ai_provider = MagicMock() - report_generator.ai_provider.check_connection = AsyncMock(return_value=True) - report_generator.ai_provider.generate_threats = AsyncMock(return_value=[ - { - "category": "Spoofing", - "description": "AI Spoof", - "business_impact": {"impact_score": 5, "likelihood_score": 5}, - "mitre_techniques": ["T1000"], - "real_world_precedents": ["CVE-2023-1234"], - "confidence": 0.9 - } - ]) - report_generator.ai_context = {"ctx": 1} - report_generator.severity_calculator.get_severity_info.return_value = {"level": "High", "score": 8.0} + with patch.object(report_generator, "_get_ai_service", new=_StdAsyncMock(return_value=None)): + # Must not raise. + asyncio.run(report_generator._run_ai_enrichment(MagicMock())) + +def test_run_ai_enrichment_delegates_to_ai_service(report_generator): + """_run_ai_enrichment() must call AIService._enrich_with_ai_threats(threat_model) — + the real, cache- and SOC-analysis-aware pipeline — passing the threat_model through + so it can mutate element.threats / tm.global_threats_llm in place. + """ + report_generator.ai_provider = MagicMock() threat_model = MagicMock() - threat_model.servers = [{"name": "S1", "description": "desc"}] - threat_model.actors = [] - threat_model.boundaries = {} + mock_service = MagicMock() + mock_service._enrich_with_ai_threats = _StdAsyncMock() + + with patch.object(report_generator, "_get_ai_service", new=_StdAsyncMock(return_value=mock_service)): + asyncio.run(report_generator._run_ai_enrichment(threat_model)) + + mock_service._enrich_with_ai_threats.assert_awaited_once_with(threat_model) + + +def test_get_ai_service_caches_across_calls(report_generator): + """The AIService instance is built once (connection check + RAG pre-warm are not + free) and reused for the ReportGenerator's lifetime — e.g. across the per-submodel + loop in project mode. + """ + report_generator._ai_config_path = Path("/some/ai_config.yaml") + mock_service_instance = MagicMock() + mock_service_instance.init_ai = _StdAsyncMock() + mock_service_cls = MagicMock(return_value=mock_service_instance) + mock_module = MagicMock(AIService=mock_service_cls) + + with patch("importlib.import_module", return_value=mock_module): + first = asyncio.run(report_generator._get_ai_service()) + second = asyncio.run(report_generator._get_ai_service()) + + assert first is second is mock_service_instance + mock_service_cls.assert_called_once() + mock_service_instance.init_ai.assert_awaited_once() - result = asyncio.run(report_generator._enrich_threats_with_ai(threat_model, [])) - assert len(result) == 1 - assert result[0]["description"] == "AI Spoof" - assert result[0]["source"] == "AI" + +def test_get_ai_service_passes_disable_rag_through(report_generator): + """report_generator._disable_rag (set from --no-rag) must reach AIService's + force_disable_rag param. + """ + report_generator._ai_config_path = Path("/some/ai_config.yaml") + report_generator._disable_rag = True + mock_service_instance = MagicMock() + mock_service_instance.init_ai = _StdAsyncMock() + mock_service_cls = MagicMock(return_value=mock_service_instance) + mock_module = MagicMock(AIService=mock_service_cls) + + with patch("importlib.import_module", return_value=mock_module): + asyncio.run(report_generator._get_ai_service()) + + mock_service_cls.assert_called_once_with( + config_path="/some/ai_config.yaml", force_disable_rag=True + ) + + +def test_get_ai_service_returns_none_without_config_path(report_generator): + report_generator._ai_config_path = None + result = asyncio.run(report_generator._get_ai_service()) + assert result is None + + +def test_get_ai_service_returns_none_on_init_failure(report_generator): + report_generator._ai_config_path = Path("/some/ai_config.yaml") + mock_module = MagicMock() + mock_module.AIService.side_effect = RuntimeError("boom") + + with patch("importlib.import_module", return_value=mock_module): + result = asyncio.run(report_generator._get_ai_service()) + assert result is None def test_open_report_in_browser(report_generator): with patch("webbrowser.open", return_value=True): @@ -131,18 +177,42 @@ def test_open_report_in_browser(report_generator): def test_generate_stix_export(report_generator, tmp_path): threat_model = MagicMock() threat_model.tm.name = "TestModel" + threat_model._report_all_detailed_threats = None # no HTML-report cache — force recompute grouped_threats = {} - - with patch.object(report_generator, "_get_all_threats_with_mitre_info", return_value=[]), \ + + with patch.object(report_generator, "_get_all_threats_with_mitre_info", return_value=[]) as mock_recompute, \ patch("threat_analysis.generation.report_generator.StixGenerator") as mock_stix_gen: - + mock_stix_gen.return_value.generate_stix_bundle.return_value = {"type": "bundle"} - + output_dir = tmp_path / "stix" result = report_generator.generate_stix_export(threat_model, grouped_threats, output_dir) - + assert result == output_dir / "TestModel_stix_attack_flow.json" assert result.exists() + mock_recompute.assert_called_once() + + +def test_generate_stix_export_reuses_ai_enriched_cache(report_generator, tmp_path): + """generate_stix_export must reuse threat_model._report_all_detailed_threats (set by + generate_html_report after AI enrichment) instead of recomputing pytm-only threats. + """ + threat_model = MagicMock() + threat_model.tm.name = "TestModel" + cached_ai_threat = {"id": "T-0001", "description": "AI-enriched threat", "source": "AI"} + threat_model._report_all_detailed_threats = [cached_ai_threat] + grouped_threats = {} + + with patch.object(report_generator, "_get_all_threats_with_mitre_info") as mock_recompute, \ + patch("threat_analysis.generation.report_generator.StixGenerator") as mock_stix_gen: + + mock_stix_gen.return_value.generate_stix_bundle.return_value = {"type": "bundle"} + + output_dir = tmp_path / "stix" + report_generator.generate_stix_export(threat_model, grouped_threats, output_dir) + + mock_recompute.assert_not_called() + mock_stix_gen.assert_called_once_with(threat_model, [cached_ai_threat]) def test_generate_summary_stats(report_generator): # Empty @@ -202,7 +272,7 @@ def test_get_all_threats_with_mitre_info(report_generator): } result = report_generator._get_all_threats_with_mitre_info(grouped_threats, threat_model) - + # We expect: 1 from grouped_threats, 1 from AI Actor threat, 1 from Global RAG # Note: deduplication might happen. assert len(result) >= 3 @@ -211,6 +281,112 @@ def test_get_all_threats_with_mitre_info(report_generator): assert "AI" in sources assert "LLM" in sources + +def test_get_all_threats_with_mitre_info_collects_boundary_ai_threats(report_generator): + """AI threats attached to a boundary (a valid enrichment target, see decisions.md + A3) must be collected into the final report — previously the collector only + iterated actors/servers/dataflows, silently dropping every AI threat whose target + was a boundary even though AIService attaches threats to boundaries too. + """ + threat_model = MagicMock() + threat_model.actors = [] + threat_model.servers = [] + threat_model.dataflows = [] + threat_model.tm.global_threats_llm = [] + + boundary_obj = MagicMock() + boundary_obj.threats = [ + MagicMock(description="Boundary AI Threat", source="AI", category="Spoofing", + confidence=0.9, capec_ids=[]) + ] + threat_model.boundaries = { + "Internet": {"boundary": boundary_obj, "business_value": "High"} + } + + report_generator.mitre_mapping.map_threat_to_mitre.return_value = {"techniques": [], "capecs": []} + report_generator.severity_calculator.get_severity_info.return_value = {"level": "Medium", "score": 5.0} + + result = report_generator._get_all_threats_with_mitre_info({}, threat_model) + + ai_threats = [t for t in result if t["source"] == "AI"] + assert len(ai_threats) == 1 + assert ai_threats[0]["target"] == "Internet" + + +def test_get_all_threats_with_mitre_info_rag_threat_uses_affected_components_as_target(report_generator): + """RAG (system-level) threats must use the LLM-provided affected_components as + target when available, instead of the generic 'Threat Model (Global)' label the + LLM's affected_components field was previously requested but never read. + """ + threat_model = MagicMock() + threat_model.actors = [] + threat_model.servers = [] + threat_model.dataflows = [] + threat_model.boundaries = {} + + rag_threat = MagicMock( + description="Cross-component pivot", category="Elevation of Privilege", + source="LLM", capec_ids=[], impact=None, likelihood=None, + ) + rag_threat.ai_details = {"affected_components": ["External Attacker", "Target Server"]} + threat_model.tm.global_threats_llm = [rag_threat] + + report_generator.mitre_mapping.map_threat_to_mitre.return_value = {"techniques": [], "capecs": []} + report_generator.severity_calculator.get_severity_info.return_value = {"level": "High", "score": 7.0} + + result = report_generator._get_all_threats_with_mitre_info({}, threat_model) + + llm_threats = [t for t in result if t["source"] == "LLM"] + assert len(llm_threats) == 1 + assert llm_threats[0]["target"] == "External Attacker → Target Server" + + +def test_get_all_threats_with_mitre_info_rag_threat_without_affected_components_falls_back(report_generator): + threat_model = MagicMock() + threat_model.actors = [] + threat_model.servers = [] + threat_model.dataflows = [] + threat_model.boundaries = {} + + rag_threat = MagicMock( + description="Generic system-level threat", category="Tampering", + source="LLM", capec_ids=[], impact=None, likelihood=None, + ) + rag_threat.ai_details = {} + threat_model.tm.global_threats_llm = [rag_threat] + + report_generator.mitre_mapping.map_threat_to_mitre.return_value = {"techniques": [], "capecs": []} + report_generator.severity_calculator.get_severity_info.return_value = {"level": "High", "score": 7.0} + + result = report_generator._get_all_threats_with_mitre_info({}, threat_model) + + llm_threats = [t for t in result if t["source"] == "LLM"] + assert len(llm_threats) == 1 + assert llm_threats[0]["target"] == "Threat Model (Global)" + + +def test_get_all_threats_with_mitre_info_syncs_severity_multipliers(report_generator): + """_get_all_threats_with_mitre_info() must sync self.severity_calculator's target + multipliers from threat_model.severity_multipliers (parsed by ModelParser from the + model's own '## Severity Multipliers' section) before any scoring happens — + self.severity_calculator may be a long-lived singleton reused across different + models (server mode) and must not score this model using a previous one's values. + """ + threat_model = MagicMock() + threat_model.actors = [] + threat_model.servers = [] + threat_model.boundaries = {} + threat_model.tm.global_threats_llm = [] + threat_model.severity_multipliers = {"WebServer": 1.5} + + report_generator.mitre_mapping.map_threat_to_mitre.return_value = {"techniques": [], "capecs": []} + report_generator.severity_calculator.get_severity_info.return_value = {"level": "Medium", "score": 5.0} + + report_generator._get_all_threats_with_mitre_info({}, threat_model) + + report_generator.severity_calculator.update_target_multipliers.assert_called_once_with({"WebServer": 1.5}) + + def test_get_all_business_values(report_generator): threat_model = MagicMock() threat_model.boundaries = {"B1": {"business_value": "BV1"}} diff --git a/tests/test_report_serializer.py b/tests/test_report_serializer.py index bea39d4..d197ae8 100644 --- a/tests/test_report_serializer.py +++ b/tests/test_report_serializer.py @@ -127,6 +127,35 @@ def test_threat_id_single_threat(): assert report["threats"][0]["id"] == "T-0001" +def test_serialize_reuses_pre_assigned_id(): + """A threat dict carrying a pre-assigned T-NNNN id (set upstream by + generate_html_report() so CISO triage can cite ids matching the final report) + must keep that id rather than being renumbered by position. + """ + threat_model = _make_threat_model() + t0, t1 = _make_threat(), _make_threat(description="Other threat") + t0["id"] = "T-0007" # pre-assigned upstream, out of positional order + t1["id"] = "T-0003" + + report = ReportSerializer.serialize(threat_model, [t0, t1]) + + ids = [t["id"] for t in report["threats"]] + assert ids == ["T-0007", "T-0003"] + + +def test_serialize_falls_back_to_positional_id_when_missing_or_invalid(): + """Threats with no id, or a non-'T-' id, still get numbered by position.""" + threat_model = _make_threat_model() + t0 = _make_threat() # no id at all + t1 = _make_threat(description="Other threat") + t1["id"] = "not-a-valid-id" + + report = ReportSerializer.serialize(threat_model, [t0, t1]) + + ids = [t["id"] for t in report["threats"]] + assert ids == ["T-0001", "T-0002"] + + # --------------------------------------------------------------------------- # Tests — schema_version # --------------------------------------------------------------------------- diff --git a/tests/test_server.py b/tests/test_server.py index 9d790eb..d8dc186 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -15,7 +15,7 @@ import pytest import json import os -from unittest.mock import patch, MagicMock, mock_open, ANY +from unittest.mock import patch, MagicMock, mock_open, ANY, AsyncMock from io import BytesIO import base64 import sys @@ -29,6 +29,7 @@ # Now we can import the app import threat_analysis.server.server as _server_module from threat_analysis.server.server import app, run_server, DEFAULT_EMPTY_MARKDOWN, get_threat_model_service +from flask import session @pytest.fixture def client(): @@ -36,6 +37,7 @@ def client(): app.config['TESTING'] = True # Reset module-level globals so tests don't bleed state into each other _server_module.initial_model_file_path = None + _server_module.graphical_editor_enabled = False with app.test_client() as client: yield client @@ -158,6 +160,30 @@ def test_run_server_with_no_model_file(client): assert response.status_code == 200 assert b'Threat Model Editor' in response.data +def test_run_server_warns_when_bound_to_non_loopback_host(client, caplog, monkeypatch): + """No route has authentication (single-user tool, by design) — binding to + anything other than loopback (e.g. FLASK_HOST=0.0.0.0, what the Docker image + sets by default) must log a clear warning, since the README/DOCKER_HUB.md's + first onboarding command exposes exactly this configuration. + """ + monkeypatch.setenv("FLASK_HOST", "0.0.0.0") + with patch('os.path.exists', return_value=False): + with patch('threat_analysis.server.server.app.run'): + with caplog.at_level("WARNING"): + run_server(model_filepath=None) + + assert "NO authentication" in caplog.text + assert "0.0.0.0" in caplog.text + +def test_run_server_no_warning_when_bound_to_loopback(client, caplog, monkeypatch): + monkeypatch.setenv("FLASK_HOST", "127.0.0.1") + with patch('os.path.exists', return_value=False): + with patch('threat_analysis.server.server.app.run'): + with caplog.at_level("WARNING"): + run_server(model_filepath=None) + + assert "NO authentication" not in caplog.text + def test_run_server_with_non_existent_model_file(client): """Test that run_server starts with DEFAULT_EMPTY_MARKDOWN if a non-existent model file is provided.""" with patch('os.path.exists', return_value=False): # Simulate file not found @@ -195,6 +221,71 @@ def mocked_open(path, *args, **kwargs): # Verify that our specific file was indeed opened assert any(call.args[0] == model_path for call in mock_file.call_args_list) +def test_initialize_ai_in_background_logs_when_event_queue_push_fails(caplog): + """A broken ai_status_event_queue must be logged, not silently swallowed — + previously `except Exception: pass` meant a broken queue silently stopped all + further AI status updates from reaching the client with zero trace in the logs. + """ + mock_service = MagicMock() + mock_service.init_ai = AsyncMock() + mock_service.ai_online = True + + with patch('threat_analysis.server.server.get_threat_model_service', return_value=mock_service), \ + patch('threat_analysis.server.server.ai_status_broadcaster') as mock_broadcaster, \ + patch('threat_analysis.server.events.ai_status_event_queue') as mock_queue: + mock_queue.put.side_effect = RuntimeError("queue is broken") + with caplog.at_level("WARNING"): + _server_module.initialize_ai_in_background() + + mock_broadcaster.broadcast.assert_called_once_with("ai_status", {"ai_online": True}) + assert "Could not push ai_status to event queue" in caplog.text + assert "queue is broken" in caplog.text + +def test_get_active_model_file_path_prefers_session_over_global(): + """_get_active_model_file_path() (used everywhere set_project_path's effect is + read) must prefer the current session's override over the process-wide + startup global — this is the core of the fix: two browser sessions must not + share one active project via mutable global state (server runs threaded=True). + """ + _server_module.initial_model_file_path = "/startup/default.md" + with app.test_request_context(): + # No session override yet -> falls back to the startup value. + assert _server_module._get_active_model_file_path() == "/startup/default.md" + + session['model_file_path'] = "/session/override.md" + assert _server_module._get_active_model_file_path() == "/session/override.md" + +def test_get_active_project_path_prefers_session_over_global(): + _server_module.initial_project_path = "/startup/project" + with app.test_request_context(): + assert _server_module._get_active_project_path() == "/startup/project" + + session['project_path'] = "/session/project" + assert _server_module._get_active_project_path() == "/session/project" + +def test_set_project_path_isolated_between_sessions(tmp_path): + """Two independent browser sessions (different cookie jars) must each keep + their own session['project_path'] / session['model_file_path'] — simulated + directly via two request contexts since set_project_path itself is one of the + `await request.get_json()` routes this test suite documents as untestable via + the synchronous WSGI client (see note above). + """ + project_a = tmp_path / "project_a" + project_a.mkdir() + project_b = tmp_path / "project_b" + project_b.mkdir() + + with app.test_client() as client_a, app.test_client() as client_b: + with client_a.session_transaction() as sess_a: + sess_a['project_path'] = str(project_a) + with client_b.session_transaction() as sess_b: + sess_b['project_path'] = str(project_b) + + with client_a.session_transaction() as sess_a: + assert sess_a['project_path'] == str(project_a) + with client_b.session_transaction() as sess_b: + assert sess_b['project_path'] == str(project_b) + def test_export_all_api_success(client): """Test the /api/export_all endpoint for successful ZIP file generation.""" mock_markdown = "# Test Model" @@ -294,6 +385,7 @@ def test_save_model_missing_markdown(client): def test_graphical_update_success(client): """Test the /api/graphical_update endpoint with valid JSON data.""" + _server_module.graphical_editor_enabled = True with patch('threat_analysis.server.server.convert_json_to_markdown') as mock_convert, \ patch('threat_analysis.server.server.get_threat_model_service') as mock_get_service: @@ -312,6 +404,7 @@ def test_graphical_update_success(client): def test_graphical_update_empty_json(client): """Test the /api/graphical_update endpoint with empty JSON data.""" + _server_module.graphical_editor_enabled = True response = client.post('/api/graphical_update', data=json.dumps({}), content_type='application/json') assert response.status_code == 400 assert 'JSON data is empty' in response.get_json()['error'] @@ -479,6 +572,11 @@ def test_get_model_name_function(): markdown_whitespace = "# Threat Model: Test Model \n## Description\nTest" assert get_model_name(markdown_whitespace) == "Test Model" + # "System Model:" prefix is also accepted (terminology: the DSL file describes + # the system model; "Threat Model" is reserved for the analysis output). + markdown_system_model = "# System Model: My Test Model\n## Description\nTest" + assert get_model_name(markdown_system_model) == "My Test Model" + def test_convert_json_to_markdown_function(): """Test the convert_json_to_markdown function.""" @@ -505,7 +603,7 @@ def test_convert_json_to_markdown_function(): markdown = convert_json_to_markdown(json_data) # Check that all sections are present - assert '# Threat Model: Graphical Editor' in markdown + assert '# System Model: Graphical Editor' in markdown assert '## Boundaries' in markdown assert '## Actors' in markdown assert '## Servers' in markdown @@ -532,13 +630,30 @@ def test_simple_mode_route(client): # Check that the initial markdown is base64 encoded in the response assert b'Threat Model Editor' in response.data -def test_graphical_editor_route(client): - """Test the /graphical route.""" +def test_graphical_editor_route_disabled_by_default(client): + """The /graphical route is disabled unless explicitly enabled.""" + response = client.get('/graphical') + assert response.status_code == 404 + + +def test_graphical_editor_route_enabled(client): + """The /graphical route serves the editor once explicitly enabled.""" + _server_module.graphical_editor_enabled = True response = client.get('/graphical') assert response.status_code == 200 assert b'Graphical Editor' in response.data or b'threat-model' in response.data +def test_index_route_hides_graphical_editor_link_by_default(client): + """The main menu only links to the Graphical Editor when it is enabled.""" + response = client.get('/') + assert b'/graphical' not in response.data + + _server_module.graphical_editor_enabled = True + response = client.get('/') + assert b'/graphical' in response.data + + def test_static_files_route(client): """Test the /static/ route.""" # This is a basic test - in a real scenario, you'd need to ensure the static files exist @@ -1143,6 +1258,7 @@ def test_update_api_value_error_from_service(client): def test_graphical_update_exception(client): """Unhandled exception in /api/graphical_update → 500 (lines 598-600).""" + _server_module.graphical_editor_enabled = True with patch('threat_analysis.server.server.convert_json_to_markdown') as mock_convert, \ patch('threat_analysis.server.server.get_threat_model_service') as mock_get: mock_svc = MagicMock() diff --git a/tests/test_severity_calculator.py b/tests/test_severity_calculator.py index e354e25..22288a0 100644 --- a/tests/test_severity_calculator.py +++ b/tests/test_severity_calculator.py @@ -220,3 +220,54 @@ def test_explanation_reflects_custom_voc_delta(self): explanation = calc.get_calculation_explanation() # Custom cve_match delta is 1.0 assert "1.0" in explanation + + +# --------------------------------------------------------------------------- +# update_target_multipliers — must REPLACE, not merge +# --------------------------------------------------------------------------- + +class TestUpdateTargetMultipliers: + def test_replaces_previous_multipliers(self): + """A long-lived SeverityCalculator (server singleton) reused across different + models must not leak a previous model's multipliers into a later one — a merge + (.update()) would; a replace won't. + """ + calc = SeverityCalculator() + calc.target_multipliers = {"OldServer": 2.0} + calc.update_target_multipliers({"NewServer": 1.5}) + assert calc.target_multipliers == {"NewServer": 1.5} + assert "OldServer" not in calc.target_multipliers + + def test_empty_dict_clears_previous_multipliers(self): + """A model with no '## Severity Multipliers' section of its own must clear + whatever a previous model (or the constructor default) left behind. + """ + calc = SeverityCalculator() + calc.target_multipliers = {"StaleServer": 3.0} + calc.update_target_multipliers({}) + assert calc.target_multipliers == {} + + +# --------------------------------------------------------------------------- +# SeverityCalculator(markdown_content=...) — in-memory content takes precedence +# --------------------------------------------------------------------------- + +class TestSeverityCalculatorMarkdownContent: + def test_markdown_content_used_over_file(self, tmp_path): + content = "## Severity Multipliers\n- **RealServer**: 2.5\n" + # Point markdown_file_path at a file that would give DIFFERENT multipliers, + # to prove markdown_content wins. + bogus_file = tmp_path / "bogus.md" + bogus_file.write_text("## Severity Multipliers\n- **WrongServer**: 9.0\n") + calc = SeverityCalculator(markdown_file_path=str(bogus_file), markdown_content=content) + assert calc.target_multipliers == {"RealServer": 2.5} + + def test_markdown_content_none_falls_back_to_file(self, tmp_path): + f = tmp_path / "model.md" + f.write_text("## Severity Multipliers\n- **FileServer**: 1.8\n") + calc = SeverityCalculator(markdown_file_path=str(f)) + assert calc.target_multipliers == {"FileServer": 1.8} + + def test_markdown_content_with_no_multipliers_section_is_empty(self): + calc = SeverityCalculator(markdown_content="# System Model: X\n\n## Boundaries\n") + assert calc.target_multipliers == {} diff --git a/tests/test_stride_constants.py b/tests/test_stride_constants.py new file mode 100644 index 0000000..9f8128b --- /dev/null +++ b/tests/test_stride_constants.py @@ -0,0 +1,50 @@ +# Copyright 2025 ellipse2v +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for threat_analysis/core/stride_constants.py + +Used throughout the pipeline (threat grouping, summary stats, report filtering) as +the single source of truth for the 6 canonical STRIDE categories — a regression here +(e.g. a typo introduced while editing) would silently drop threats from every report. +""" + +from threat_analysis.core.stride_constants import STRIDE_CATEGORIES + + +def test_stride_categories_is_a_frozenset(): + assert isinstance(STRIDE_CATEGORIES, frozenset) + + +def test_stride_categories_has_exactly_six_entries(): + assert len(STRIDE_CATEGORIES) == 6 + + +def test_stride_categories_content(): + assert STRIDE_CATEGORIES == frozenset({ + "Spoofing", + "Tampering", + "Repudiation", + "Information Disclosure", + "Denial of Service", + "Elevation of Privilege", + }) + + +def test_stride_categories_immutable(): + """frozenset has no in-place mutation methods — any accidental `.add(...)` call + elsewhere in the codebase must fail loudly instead of silently corrupting this + shared constant. + """ + assert not hasattr(STRIDE_CATEGORIES, "add") + assert not hasattr(STRIDE_CATEGORIES, "remove") diff --git a/tests/test_svg_generator.py b/tests/test_svg_generator.py index fce76b0..4fbd537 100644 --- a/tests/test_svg_generator.py +++ b/tests/test_svg_generator.py @@ -143,6 +143,19 @@ def test_shapes_and_styles(generator): assert ' hop assets) +# --------------------------------------------------------------------------- + +def _make_hop(asset_name: str) -> SimpleNamespace: + return SimpleNamespace(asset_name=asset_name) + + +def _make_gdaf_scenario(scenario_id="GDAF-1", entry_point="Attacker", hop_names=("WebApp",), + objective_name="Exfiltrate data", actor_name="Attacker", + risk_level="CRITICAL", path_score=4.2) -> SimpleNamespace: + return SimpleNamespace( + scenario_id=scenario_id, + entry_point=entry_point, + hops=[_make_hop(h) for h in hop_names], + objective_name=objective_name, + actor_name=actor_name, + risk_level=risk_level, + path_score=path_score, + ) + + +def _make_debate_result(scenario_id="GDAF-1", residual_path_viable=True) -> SimpleNamespace: + return SimpleNamespace(scenario_id=scenario_id, residual_path_viable=residual_path_viable) + + +class TestGdafPaths: + def test_no_scenarios_returns_empty_list(self): + rg = _make_rg() + m = _make_model(servers=[_make_server("WebApp")]) + result = rg._build_threat_graph_data(m, []) + assert result["gdaf_paths"] == [] + + def test_scenario_produces_node_sequence(self): + rg = _make_rg() + m = _make_model(servers=[_make_server("WebApp")]) + scenario = _make_gdaf_scenario(entry_point="Attacker", hop_names=("WebApp", "DB")) + result = rg._build_threat_graph_data(m, [], gdaf_scenarios=[scenario]) + assert len(result["gdaf_paths"]) == 1 + path = result["gdaf_paths"][0] + assert path["nodes"] == ["Attacker", "WebApp", "DB"] + assert path["scenario_id"] == "GDAF-1" + assert path["debated"] is False + assert path["residual_path_viable"] is None + + def test_consecutive_duplicate_nodes_deduplicated(self): + """entry_point sometimes equals the first hop's asset_name — must not repeat.""" + rg = _make_rg() + m = _make_model(servers=[_make_server("WebApp")]) + scenario = _make_gdaf_scenario(entry_point="WebApp", hop_names=("WebApp", "DB")) + result = rg._build_threat_graph_data(m, [], gdaf_scenarios=[scenario]) + assert result["gdaf_paths"][0]["nodes"] == ["WebApp", "DB"] + + def test_scenario_with_single_node_excluded(self): + rg = _make_rg() + m = _make_model(servers=[_make_server("WebApp")]) + scenario = _make_gdaf_scenario(entry_point="WebApp", hop_names=()) + result = rg._build_threat_graph_data(m, [], gdaf_scenarios=[scenario]) + assert result["gdaf_paths"] == [] + + def test_debated_scenario_marks_viable(self): + rg = _make_rg() + m = _make_model(servers=[_make_server("WebApp")]) + scenario = _make_gdaf_scenario(scenario_id="GDAF-1", hop_names=("WebApp",)) + debate_result = _make_debate_result(scenario_id="GDAF-1", residual_path_viable=True) + result = rg._build_threat_graph_data( + m, [], gdaf_scenarios=[scenario], debate_results=[debate_result] + ) + path = result["gdaf_paths"][0] + assert path["debated"] is True + assert path["residual_path_viable"] is True + + def test_debated_scenario_marks_blocked(self): + rg = _make_rg() + m = _make_model(servers=[_make_server("WebApp")]) + scenario = _make_gdaf_scenario(scenario_id="GDAF-1", hop_names=("WebApp",)) + debate_result = _make_debate_result(scenario_id="GDAF-1", residual_path_viable=False) + result = rg._build_threat_graph_data( + m, [], gdaf_scenarios=[scenario], debate_results=[debate_result] + ) + assert result["gdaf_paths"][0]["residual_path_viable"] is False + + def test_debate_result_for_different_scenario_not_matched(self): + rg = _make_rg() + m = _make_model(servers=[_make_server("WebApp")]) + scenario = _make_gdaf_scenario(scenario_id="GDAF-1", hop_names=("WebApp",)) + debate_result = _make_debate_result(scenario_id="GDAF-OTHER", residual_path_viable=True) + result = rg._build_threat_graph_data( + m, [], gdaf_scenarios=[scenario], debate_results=[debate_result] + ) + assert result["gdaf_paths"][0]["debated"] is False + + def test_scenarios_capped_at_10(self): + rg = _make_rg() + m = _make_model(servers=[_make_server("WebApp")]) + scenarios = [ + _make_gdaf_scenario(scenario_id=f"GDAF-{i}", hop_names=("WebApp",)) + for i in range(15) + ] + result = rg._build_threat_graph_data(m, [], gdaf_scenarios=scenarios) + assert len(result["gdaf_paths"]) == 10 diff --git a/threatModel_Template/Advanced_Cloud_WebApp/model.md b/threatModel_Template/Advanced_Cloud_WebApp/model.md index 04ebbcf..b0ec8ff 100644 --- a/threatModel_Template/Advanced_Cloud_WebApp/model.md +++ b/threatModel_Template/Advanced_Cloud_WebApp/model.md @@ -3,7 +3,7 @@ Copyright 2024 ellipse2v — Apache License, Version 2.0 https://www.apache.org/licenses/LICENSE-2.0 --> -# Threat Model: Advanced Cloud Web Application +# System Model: Advanced Cloud Web Application ## Description Advanced cloud web app with CDN, API gateway, auth service, product microservice, and managed database. diff --git a/threatModel_Template/CI_CD_Pipeline/model.md b/threatModel_Template/CI_CD_Pipeline/model.md index d31882b..14a9bc8 100644 --- a/threatModel_Template/CI_CD_Pipeline/model.md +++ b/threatModel_Template/CI_CD_Pipeline/model.md @@ -3,7 +3,7 @@ Copyright 2024 ellipse2v — Apache License, Version 2.0 https://www.apache.org/licenses/LICENSE-2.0 --> -# Threat Model: CI/CD Pipeline (Continuous Integration and Deployment) +# System Model: CI/CD Pipeline (Continuous Integration and Deployment) ## Description This threat model focuses on a CI/CD pipeline, covering the development, integration, testing, and deployment stages. diff --git a/threatModel_Template/Cloud_Native_Architecture/model.md b/threatModel_Template/Cloud_Native_Architecture/model.md index 762d36b..b91abf6 100644 --- a/threatModel_Template/Cloud_Native_Architecture/model.md +++ b/threatModel_Template/Cloud_Native_Architecture/model.md @@ -3,7 +3,7 @@ Copyright 2024 ellipse2v — Apache License, Version 2.0 https://www.apache.org/licenses/LICENSE-2.0 --> -# Threat Model: Cloud-Native Architecture +# System Model: Cloud-Native Architecture ## Description This threat model addresses cloud-native architectures, focusing on serverless functions, API gateways, managed databases, and object storage. diff --git a/threatModel_Template/Data_Warehouse_Data_Lake/model.md b/threatModel_Template/Data_Warehouse_Data_Lake/model.md index 5443440..46bc3ca 100644 --- a/threatModel_Template/Data_Warehouse_Data_Lake/model.md +++ b/threatModel_Template/Data_Warehouse_Data_Lake/model.md @@ -1,4 +1,4 @@ -# Threat Model: Data Warehouse / Data Lake Architecture +# System Model: Data Warehouse / Data Lake Architecture ## Description Data warehouse and data lake architecture with ETL pipelines, object storage, data lakehouse, and BI tools. diff --git a/threatModel_Template/GDAF_Debate_Smoke_Test/context/gdaf_test_context.yaml b/threatModel_Template/GDAF_Debate_Smoke_Test/context/gdaf_test_context.yaml new file mode 100644 index 0000000..0c5183b --- /dev/null +++ b/threatModel_Template/GDAF_Debate_Smoke_Test/context/gdaf_test_context.yaml @@ -0,0 +1,35 @@ +system_description: "Minimal test system for GDAF and Red/Blue debate smoke testing." +sector: "technology" +internet_facing: true +compliance_requirements: [] + +attack_objectives: + - id: "OBJ-TEST" + name: "Compromise Target Server" + description: "Reach and compromise the single target server from the internet." + target_asset_names: + - "Target Server" + target_types: ["web-server"] + attacker_intent: "Proof-of-concept access — smoke test only." + business_impact: "Test scenario — no real impact." + mitre_final_tactic: "impact" + min_severity_score: 1.0 + +threat_actors: + - id: "TA-TEST" + name: "Test Attacker" + sophistication: "low" + objectives: ["OBJ-TEST"] + entry_preference: "internet-facing" + description: "Minimal low-sophistication attacker profile for smoke testing." + known_ttps: + - "T1190" + capable_tactics: + - "initial-access" + - "impact" + +risk_criteria: + acceptable_risk_score: 5.0 + max_hops: 2 + max_paths_per_objective: 1 + gdaf_min_technique_score: 0.8 diff --git a/threatModel_Template/GDAF_Debate_Smoke_Test/model.md b/threatModel_Template/GDAF_Debate_Smoke_Test/model.md new file mode 100644 index 0000000..0f982b0 --- /dev/null +++ b/threatModel_Template/GDAF_Debate_Smoke_Test/model.md @@ -0,0 +1,22 @@ +# System Model: GDAF Debate Smoke Test + +## Description +Minimal two-node system used to smoke-test GDAF attack path generation and the +Red/Blue debate feature with the smallest possible token footprint. Not a +realistic architecture — just enough for one attack path. + +## Context +gdaf_context = context/gdaf_test_context.yaml + +## Boundaries +- **Internet**: isTrusted=False, type=network-on-prem, traversal_difficulty=low, description="Untrusted public internet — attacker entry point." +- **Internal Network**: isTrusted=True, type=network-on-prem, traversal_difficulty=medium, description="Trusted internal network hosting the target server." + +## Actors +- **External Attacker**: boundary=Internet, authenticity=none, isTrusted=False, description="Unauthenticated external attacker with no prior access." + +## Servers +- **Target Server**: boundary="Internal Network", type=web-server, machine=virtual, internet_facing=True, mfa_enabled=False, credentials_stored=True, confidentiality=high, integrity=high, availability=medium, description="Single target asset exposed to the attacker for GDAF path testing." + +## Dataflows +- **AttackerToServer**: from="External Attacker", to="Target Server", protocol=HTTPS, is_encrypted=True, is_authenticated=False, description="Direct unauthenticated path from the attacker to the target." diff --git a/threatModel_Template/IP_Camera_Surveillance/model.md b/threatModel_Template/IP_Camera_Surveillance/model.md index ee77cab..cd7dc51 100644 --- a/threatModel_Template/IP_Camera_Surveillance/model.md +++ b/threatModel_Template/IP_Camera_Surveillance/model.md @@ -1,4 +1,4 @@ -# Threat Model: IP Camera Surveillance System +# System Model: IP Camera Surveillance System ## Description IP-based video surveillance system for a commercial building: fixed and PTZ cameras, diff --git a/threatModel_Template/IoT/model.md b/threatModel_Template/IoT/model.md index f4b8a2b..1386d30 100644 --- a/threatModel_Template/IoT/model.md +++ b/threatModel_Template/IoT/model.md @@ -1,4 +1,4 @@ -# Threat Model: IoT Architecture +# System Model: IoT Architecture ## Description IoT system with sensors, gateway, cloud platform, analytics, and device management. diff --git a/threatModel_Template/Kubernetes_Helm_Cluster/model.md b/threatModel_Template/Kubernetes_Helm_Cluster/model.md index 89c639e..ab78dd8 100644 --- a/threatModel_Template/Kubernetes_Helm_Cluster/model.md +++ b/threatModel_Template/Kubernetes_Helm_Cluster/model.md @@ -1,4 +1,4 @@ -# Threat Model: Kubernetes / Helm Cluster +# System Model: Kubernetes / Helm Cluster ## Description Production Kubernetes cluster deployed via Helm, with three application namespaces diff --git a/threatModel_Template/Microservices_Architecture/model.md b/threatModel_Template/Microservices_Architecture/model.md index e887723..086bae8 100644 --- a/threatModel_Template/Microservices_Architecture/model.md +++ b/threatModel_Template/Microservices_Architecture/model.md @@ -3,7 +3,7 @@ Copyright 2024 ellipse2v — Apache License, Version 2.0 https://www.apache.org/licenses/LICENSE-2.0 --> -# Threat Model: Microservices Architecture +# System Model: Microservices Architecture ## Description This threat model covers a microservices-based application with an API gateway, service mesh, message broker, and shared database. diff --git a/threatModel_Template/Mobile_Application/model.md b/threatModel_Template/Mobile_Application/model.md index 91117c2..bb87a58 100644 --- a/threatModel_Template/Mobile_Application/model.md +++ b/threatModel_Template/Mobile_Application/model.md @@ -1,4 +1,4 @@ -# Threat Model: Mobile Application +# System Model: Mobile Application ## Description Mobile application with backend API, third-party identity provider, and payment gateway integration. diff --git a/threatModel_Template/On-Prem_Enterprise_Network/model.md b/threatModel_Template/On-Prem_Enterprise_Network/model.md index 876e8c8..d8e11e1 100644 --- a/threatModel_Template/On-Prem_Enterprise_Network/model.md +++ b/threatModel_Template/On-Prem_Enterprise_Network/model.md @@ -1,4 +1,4 @@ -# Threat Model: On-Prem Enterprise Network +# System Model: On-Prem Enterprise Network ## Description This model describes a complex on-premises enterprise network for a large multinational manufacturing organization. It covers 9 security zones ranging from an untrusted internet perimeter to a highly restricted OT/SCADA environment. The architecture includes Active Directory infrastructure, a financial ERP system, a CI/CD development pipeline, and legacy industrial control systems. This model highlights threats related to lateral movement, credential theft, legacy protocols, and OT network security. diff --git a/threatModel_Template/Satellite_Ground_Segment/model.md b/threatModel_Template/Satellite_Ground_Segment/model.md index 5f6c32c..419d986 100644 --- a/threatModel_Template/Satellite_Ground_Segment/model.md +++ b/threatModel_Template/Satellite_Ground_Segment/model.md @@ -1,4 +1,4 @@ -# Threat Model: Satellite Ground Segment (Thales Demo Scenario) +# System Model: Satellite Ground Segment (Thales Demo Scenario) ## Description Low Earth Orbit satellite system modelling the attack surface demonstrated by Thales diff --git a/threatModel_Template/Serverless_AWS_Lambda/model.md b/threatModel_Template/Serverless_AWS_Lambda/model.md index 1ef2407..ed108b5 100644 --- a/threatModel_Template/Serverless_AWS_Lambda/model.md +++ b/threatModel_Template/Serverless_AWS_Lambda/model.md @@ -1,4 +1,4 @@ -# Threat Model: Serverless — AWS Lambda Event-Driven Architecture +# System Model: Serverless — AWS Lambda Event-Driven Architecture ## Description Event-driven serverless application on AWS. External clients reach the system via diff --git a/threatModel_Template/Simple_Monolithic_Web_Application/model.md b/threatModel_Template/Simple_Monolithic_Web_Application/model.md index 70c8efb..4809d57 100644 --- a/threatModel_Template/Simple_Monolithic_Web_Application/model.md +++ b/threatModel_Template/Simple_Monolithic_Web_Application/model.md @@ -1,4 +1,4 @@ -# Threat Model: Simple Monolithic Web Application +# System Model: Simple Monolithic Web Application ## Description A monolithic web application with a single application server handling all business logic. diff --git a/threatModel_Template/Three_Tier_Architecture/model.md b/threatModel_Template/Three_Tier_Architecture/model.md index 5a35a3d..5541c2e 100644 --- a/threatModel_Template/Three_Tier_Architecture/model.md +++ b/threatModel_Template/Three_Tier_Architecture/model.md @@ -3,7 +3,7 @@ Copyright 2024 ellipse2v — Apache License, Version 2.0 https://www.apache.org/licenses/LICENSE-2.0 --> -# Threat Model: Three-Tier Architecture +# System Model: Three-Tier Architecture ## Description Classic three-tier web architecture: presentation, business logic, and data. diff --git a/threatModel_Template/Traditional_Enterprise_Network/model.md b/threatModel_Template/Traditional_Enterprise_Network/model.md index b25f751..80b0369 100644 --- a/threatModel_Template/Traditional_Enterprise_Network/model.md +++ b/threatModel_Template/Traditional_Enterprise_Network/model.md @@ -1,4 +1,4 @@ -# Threat Model: Traditional Enterprise Network +# System Model: Traditional Enterprise Network ## Description Traditional enterprise with Active Directory, DMZ, file servers, and application servers. diff --git a/threatModel_Template/UAV_Drone_System/model.md b/threatModel_Template/UAV_Drone_System/model.md index 61c15ae..70d4c58 100644 --- a/threatModel_Template/UAV_Drone_System/model.md +++ b/threatModel_Template/UAV_Drone_System/model.md @@ -1,4 +1,4 @@ -# Threat Model: DJI Power Inspection Drone System +# System Model: DJI Power Inspection Drone System ## Description diff --git a/threatModel_Template/Zero_Trust/model.md b/threatModel_Template/Zero_Trust/model.md index 63e09c7..f7d3c88 100644 --- a/threatModel_Template/Zero_Trust/model.md +++ b/threatModel_Template/Zero_Trust/model.md @@ -3,7 +3,7 @@ Copyright 2024 ellipse2v — Apache License, Version 2.0 https://www.apache.org/licenses/LICENSE-2.0 --> -# Threat Model: Zero Trust Architecture +# System Model: Zero Trust Architecture ## Description This threat model describes a zero trust network architecture with identity-based access control. diff --git a/threatModel_Template/projects/example_3/api_gateway/model.md b/threatModel_Template/projects/example_3/api_gateway/model.md index 0df30d0..b5f6b3b 100644 --- a/threatModel_Template/projects/example_3/api_gateway/model.md +++ b/threatModel_Template/projects/example_3/api_gateway/model.md @@ -1,4 +1,4 @@ -# Threat Model: EcoShop — API Gateway Cluster +# System Model: EcoShop — API Gateway Cluster ## Description The API gateway cluster consists of a load balancer distributing requests across two Kong diff --git a/threatModel_Template/projects/example_3/backend/database/model.md b/threatModel_Template/projects/example_3/backend/database/model.md index fdfa32e..795c60f 100644 --- a/threatModel_Template/projects/example_3/backend/database/model.md +++ b/threatModel_Template/projects/example_3/backend/database/model.md @@ -1,4 +1,4 @@ -# Threat Model: EcoShop — Database Cluster +# System Model: EcoShop — Database Cluster ## Description The database cluster consists of a primary PostgreSQL instance and a streaming replica. diff --git a/threatModel_Template/projects/example_3/backend/model.md b/threatModel_Template/projects/example_3/backend/model.md index 7bc2276..394ff2e 100644 --- a/threatModel_Template/projects/example_3/backend/model.md +++ b/threatModel_Template/projects/example_3/backend/model.md @@ -1,4 +1,4 @@ -# Threat Model: EcoShop — Backend Services Tier +# System Model: EcoShop — Backend Services Tier ## Description The backend services tier hosts the core business microservices: an authentication service, diff --git a/threatModel_Template/projects/example_3/dmz/model.md b/threatModel_Template/projects/example_3/dmz/model.md index ac6b0db..d8afc4f 100644 --- a/threatModel_Template/projects/example_3/dmz/model.md +++ b/threatModel_Template/projects/example_3/dmz/model.md @@ -1,4 +1,4 @@ -# Threat Model: EcoShop — DMZ Tier +# System Model: EcoShop — DMZ Tier ## Description The DMZ contains two redundant hardware firewalls (active/passive HA pair) and an API diff --git a/threatModel_Template/projects/example_3/frontend/model.md b/threatModel_Template/projects/example_3/frontend/model.md index af6dda7..2fdf1da 100644 --- a/threatModel_Template/projects/example_3/frontend/model.md +++ b/threatModel_Template/projects/example_3/frontend/model.md @@ -1,4 +1,4 @@ -# Threat Model: EcoShop — Frontend Tier +# System Model: EcoShop — Frontend Tier ## Description The frontend tier consists of an nginx web server delivering the React SPA and a HAProxy diff --git a/threatModel_Template/projects/example_3/main.md b/threatModel_Template/projects/example_3/main.md index 2b9c2e4..8188c11 100644 --- a/threatModel_Template/projects/example_3/main.md +++ b/threatModel_Template/projects/example_3/main.md @@ -1,4 +1,4 @@ -# Threat Model: EcoShop — E-commerce Platform (Overview) +# System Model: EcoShop — E-commerce Platform (Overview) ## Description EcoShop is a B2C e-commerce platform serving EU and US customers. The architecture follows diff --git a/threatModel_Template/projects/example_3/order_service/model.md b/threatModel_Template/projects/example_3/order_service/model.md index e3ebac7..e81e72e 100644 --- a/threatModel_Template/projects/example_3/order_service/model.md +++ b/threatModel_Template/projects/example_3/order_service/model.md @@ -1,4 +1,4 @@ -# Threat Model: EcoShop — Order Service Cluster +# System Model: EcoShop — Order Service Cluster ## Description The Order Service is an event-driven microservices cluster responsible for the full order diff --git a/threatModel_Template/projects/example_4/card_vault/model.md b/threatModel_Template/projects/example_4/card_vault/model.md index 07fd409..00f215c 100644 --- a/threatModel_Template/projects/example_4/card_vault/model.md +++ b/threatModel_Template/projects/example_4/card_vault/model.md @@ -1,4 +1,4 @@ -# Threat Model: PayFlow — Card Data Vault +# System Model: PayFlow — Card Data Vault ## Description The card data vault is the most sensitive component in the PayFlow architecture. A TokenProxy diff --git a/threatModel_Template/projects/example_4/fraud_engine/model.md b/threatModel_Template/projects/example_4/fraud_engine/model.md index e952c02..8fa010a 100644 --- a/threatModel_Template/projects/example_4/fraud_engine/model.md +++ b/threatModel_Template/projects/example_4/fraud_engine/model.md @@ -1,4 +1,4 @@ -# Threat Model: PayFlow — Fraud Detection Engine +# System Model: PayFlow — Fraud Detection Engine ## Description The fraud detection engine provides real-time payment risk scoring with a latency SLA of diff --git a/threatModel_Template/projects/example_4/main.md b/threatModel_Template/projects/example_4/main.md index 7b603a1..bfcba9c 100644 --- a/threatModel_Template/projects/example_4/main.md +++ b/threatModel_Template/projects/example_4/main.md @@ -1,4 +1,4 @@ -# Threat Model: PayFlow — Cloud-Native Fintech Payment Platform +# System Model: PayFlow — Cloud-Native Fintech Payment Platform ## Description PayFlow is a PCI-DSS Level 1 certified payment processing platform serving EU and US markets. diff --git a/threatModel_Template/projects/example_4/payment_api/model.md b/threatModel_Template/projects/example_4/payment_api/model.md index 0390bb1..f81dd4b 100644 --- a/threatModel_Template/projects/example_4/payment_api/model.md +++ b/threatModel_Template/projects/example_4/payment_api/model.md @@ -1,4 +1,4 @@ -# Threat Model: PayFlow — Payment API Services +# System Model: PayFlow — Payment API Services ## Description The payment API tier contains the core payment processing microservices behind the Kong reverse diff --git a/threatModel_Template/threat_model.md b/threatModel_Template/threat_model.md index d3195dc..aa9542c 100644 --- a/threatModel_Template/threat_model.md +++ b/threatModel_Template/threat_model.md @@ -1,4 +1,4 @@ -# Threat Model: Advanced DMZ Architecture +# System Model: Advanced DMZ Architecture ## Description This model describes a network architecture with a Demilitarized Zone (DMZ), external and internal dataflows, a scalable gateway, and a potentially untrusted command zone. The goal is to identify STRIDE threats and map them to MITRE ATT&CK techniques. diff --git a/threat_analysis/__main__.py b/threat_analysis/__main__.py index dc59606..30cd67c 100644 --- a/threat_analysis/__main__.py +++ b/threat_analysis/__main__.py @@ -22,6 +22,7 @@ import os import sys import argparse +import asyncio import logging from typing import Dict, List, Any, Optional, Tuple from pathlib import Path @@ -31,6 +32,7 @@ from threat_analysis import config # Re-add config import from threat_analysis.utils import resolve_path, _validate_path_within_project, compare_threat_reports +from threat_analysis.generation.utils import get_enriched_threats # Import library modules # Lazy imports will be handled within methods @@ -207,6 +209,7 @@ def __init__( ai_config_path: Optional[Path] = None, context_path: Optional[Path] = None, cve_definitions_path: Optional[Path] = None, # New parameter + disable_rag: bool = False, ): """Initializes the analysis framework""" self.markdown_content = markdown_content @@ -245,6 +248,7 @@ def __init__( self._ai_config_path = ai_config_path self._context_path = context_path self._cve_definitions_path = cve_definitions_path + self._disable_rag = disable_rag self._initialize_components() logging.info(f"🚀 Analysis framework initialized: {self.model_name}") @@ -298,9 +302,7 @@ def _initialize_components(self) -> None: self.mitre_mapper = MitreMapping(threat_model_path=self.model_file_path) self.threat_model = self._load_and_validate_model(self.markdown_content) - self.severity_calculator = SeverityCalculator( - markdown_file_path=str(Path("threatModel_Template/threat_model.md")) - ) + self.severity_calculator = SeverityCalculator(markdown_content=self.markdown_content) self.report_generator = ReportGenerator( self.severity_calculator, self.mitre_mapper, implemented_mitigations_path=Path(self._implemented_mitigations_path) if self._implemented_mitigations_path else None, @@ -308,6 +310,7 @@ def _initialize_components(self) -> None: ai_config_path=self._ai_config_path, context_path=self._context_path, threat_model_ref=self.threat_model, + disable_rag=self._disable_rag, ) self.diagram_generator = DiagramGenerator() @@ -457,10 +460,7 @@ def generate_navigator_layer(self) -> Optional[str]: AttackNavigatorGenerator = attack_navigator_module.AttackNavigatorGenerator # We need all detailed threats, not just grouped ones. - if self.threat_model: - all_threats = self.threat_model.get_all_threats_details() - else: - all_threats = [] + all_threats = get_enriched_threats(self.threat_model) if self.threat_model else [] navigator_generator = AttackNavigatorGenerator( threat_model_name=self.model_name, @@ -501,8 +501,10 @@ def generate_and_save_attack_flow(threat_model: 'ThreatModel', output_dir: Path, attack_flow_module = importlib.import_module("threat_analysis.generation.attack_flow_generator") AttackFlowGenerator = attack_flow_module.AttackFlowGenerator - # The generator now expects the raw threat data to perform its own filtering. - raw_threats = threat_model.mitre_analysis_results.get("processed_threats", []) + # Prefers the AI-enriched list cached by generate_html_report() (already run by + # this point via framework.generate_reports()) — falls back to the raw pytm-only + # list otherwise, so this still works if HTML generation was skipped. + raw_threats = get_enriched_threats(threat_model) if not raw_threats: logging.warning("No raw threats found, skipping Attack Flow generation.") return @@ -518,6 +520,25 @@ def generate_and_save_attack_flow(threat_model: 'ThreatModel', output_dir: Path, logging.error(f"❌ Failed to generate Attack Flow files for {model_name}: {e}") traceback.print_exc() + +def _build_project_ai_service(ai_config_path: Path, force_disable_rag: bool = False): + """Builds and initializes an AIService for `secopstm --project` mode. + + Needed for cross-model RAG threat generation (Decision A2) — see the call site in + main(). Returns None (never raises) if AI init fails, matching every other AI + degrade-silently path in this codebase — cross-model RAG is additive, not required. + """ + try: + ai_service_module = importlib.import_module("threat_analysis.server.ai_service") + AIService = ai_service_module.AIService + ai_service = AIService(config_path=str(ai_config_path), force_disable_rag=force_disable_rag) + asyncio.run(ai_service.init_ai()) + return ai_service + except Exception as e: + logging.warning(f"AIService init failed for project mode (cross-model RAG skipped, non-fatal): {e}") + return None + + def load_iac_plugins() -> Dict[str, 'IaCPlugin']: """Dynamically loads IaC plugins from the iac_plugins directory. @@ -584,10 +605,24 @@ def __init__(self, loaded_plugins: Dict[str, 'IaCPlugin']): action="store_true", help="Force re-download even if the vector store already exists (use with --init-rag).", ) + common.add_argument( + "--no-rag", + action="store_true", + dest="no_rag", + help="Disable RAG (skip vector store load / chromadb+embeddings pre-warm) for this run, " + "even if rag.enabled: true in ai_config.yaml. Useful to test with/without RAG, or to " + "speed up iteration (RAG cold-start can take minutes).", + ) common.add_argument( "--server", action="store_true", - help="Launch the web editor (Monaco + graphical canvas).", + help="Launch the web editor (Monaco editor, Simple Mode).", + ) + common.add_argument( + "--graphical-editor", + action="store_true", + dest="graphical_editor", + help="Enable the experimental drag-and-drop Graphical Editor mode (disabled by default; use with --server).", ) common.add_argument( "--port", @@ -965,6 +1000,8 @@ def run_single_analysis(args: argparse.Namespace, loaded_iac_plugins: Dict[str, _ai_config_rel = args.ai_config_file if hasattr(args, "ai_config_file") else "config/ai_config.yaml" ai_config_path = Path(_ai_config_rel) if Path(_ai_config_rel).is_absolute() else Path.cwd() / _ai_config_rel + from threat_analysis.server.server import get_model_name # noqa: PLC0415 + framework = SecOpsTMFramework( markdown_content=markdown_content_for_analysis, model_name=get_model_name(markdown_content_for_analysis), # Derive from content @@ -975,6 +1012,7 @@ def run_single_analysis(args: argparse.Namespace, loaded_iac_plugins: Dict[str, cve_service=cve_service, ai_config_path=ai_config_path, cve_definitions_path=cve_definitions_path, + disable_rag=getattr(args, "no_rag", False), ) threats = framework.run_analysis() @@ -1146,7 +1184,11 @@ def main(): if getattr(args, "port", None): os.environ["FLASK_PORT"] = str(args.port) try: - run_server(model_filepath=args.model_file, project_path=args.project) + run_server( + model_filepath=args.model_file, + project_path=args.project, + enable_graphical_editor=args.graphical_editor, + ) except ImportError: logging.error( f"[{time.time() - _start_time_main:.4f}s] ❌ Flask is not installed. Please install it to use the web server: " @@ -1202,6 +1244,7 @@ def main(): # Lazy import ReportGenerator report_generator_module = importlib.import_module("threat_analysis.generation.report_generator") ReportGenerator = report_generator_module.ReportGenerator + _disable_rag = getattr(args, "no_rag", False) report_generator = ReportGenerator( severity_calculator, mitre_mapping, @@ -1209,9 +1252,17 @@ def main(): cve_service=cve_service, ai_config_path=ai_config_path, threat_model_ref=None, # Project reports are generated per model, ref is set internally + disable_rag=_disable_rag, ) - project_threat_model = report_generator.generate_project_reports(project_path, output_dir) + # AIService (distinct from ReportGenerator's own per-submodel ai_provider) is needed + # for cross-model RAG threat generation (Decision A2) — generate_project_reports() + # only runs it when given one. Without this, `secopstm --project` silently skips + # cross-boundary RAG threats even with RAG enabled and AI online; only the web UI + # (export_service.py) passed ai_service until now. + ai_service = _build_project_ai_service(ai_config_path, force_disable_rag=_disable_rag) + + project_threat_model = report_generator.generate_project_reports(project_path, output_dir, ai_service=ai_service) if args.navigator and project_threat_model: logging.info("🗺️ Generating ATT&CK Navigator layer for project...") @@ -1220,7 +1271,7 @@ def main(): attack_navigator_module = importlib.import_module("threat_analysis.generation.attack_navigator_generator") AttackNavigatorGenerator = attack_navigator_module.AttackNavigatorGenerator - all_threats = project_threat_model.get_all_threats_details() + all_threats = get_enriched_threats(project_threat_model) navigator_generator = AttackNavigatorGenerator( threat_model_name=str(project_threat_model.tm.name), all_detailed_threats=all_threats diff --git a/threat_analysis/ai_engine/prompts/attack_flow_prompts.py b/threat_analysis/ai_engine/prompts/attack_flow_prompts.py deleted file mode 100644 index 468ea13..0000000 --- a/threat_analysis/ai_engine/prompts/attack_flow_prompts.py +++ /dev/null @@ -1,51 +0,0 @@ -# Copyright 2025 ellipse2v -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -Attack Flow prompt helpers — all prompt text lives in config/prompts.yaml. -""" - -from typing import Dict -from threat_analysis.ai_engine.prompt_loader import get as _get - - -def __getattr__(name: str) -> str: - if name == "ATTACK_FLOW_SYSTEM_PROMPT": - return _get("attack_flow", "system") - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") - - -def build_attack_flow_prompt(threat: Dict, component: Dict, context: Dict) -> str: - """Builds a STIX 2.1 Attack Flow generation prompt. - - Reads the template from ``config/prompts.yaml`` - (``attack_flow.component_template``) and injects threat + component data. - """ - mitre_techniques = threat.get("mitre_techniques", []) - threat_category = threat.get("category", "Unknown") - - return _get( - "attack_flow", - "component_template", - threat_category=threat_category, - threat_category_lower=threat_category.lower().replace(" ", "-"), - threat_title=threat.get("title", "Unknown Threat"), - threat_description=threat.get("description", ""), - attack_scenario=threat.get("attack_scenario", ""), - mitre_techniques=", ".join(mitre_techniques) if mitre_techniques else "None identified", - component_type=component.get("type", "Unknown"), - component_name=component.get("name", "Unknown"), - component_description=component.get("description", ""), - system_context=context.get("system_description", "No additional context provided"), - ) diff --git a/threat_analysis/ai_engine/providers/base_provider.py b/threat_analysis/ai_engine/providers/base_provider.py index 4588361..058e5b2 100644 --- a/threat_analysis/ai_engine/providers/base_provider.py +++ b/threat_analysis/ai_engine/providers/base_provider.py @@ -31,16 +31,6 @@ async def generate_threats( """Generates STRIDE threats for a component.""" pass - @abstractmethod - async def generate_attack_flow( - self, - threat: Dict, - component: Dict, - context: Dict, - ) -> Dict: - """Generates an Attack Flow STIX 2.1 for a threat.""" - pass - @abstractmethod async def generate_markdown( self, @@ -97,3 +87,51 @@ async def generate_soc_analysis( ``siem_rules``, ``iocs``. """ return [] + + async def generate_debate_turn( + self, + prompt: str, + system_prompt: str, + ) -> Dict: + """Generates one Red or Blue turn in an adversarial debate round. + + Optional — providers that do not override this return an empty dict, + which causes RedBlueDebateEngine to skip the turn gracefully. + + Args: + prompt: Role-specific user prompt (grounding + prior turn summary). + system_prompt: Red or Blue persona system instruction from prompts.yaml. + + Returns: + A dict with keys: ``viability_score`` and role-specific keys + (``techniques_attempted``/``failed_alternatives`` for Red, + ``techniques_blocked``/``detection_gaps`` for Blue), plus + ``rationale`` and ``evidence``. Returns ``{}`` on failure or when + not overridden. + """ + return {} + + async def generate_attack_path_narrative( + self, + prompt: str, + system_prompt: str, + ) -> Dict: + """Generates a short grounded narrative for an already-computed discovered + attack path (AttackFlowGenerator.get_paths_summary — hops/techniques/targets + are fixed before this call, the model explains them, it does not design them). + + Optional — providers that do not override this return an empty dict, which + causes the narrative to be omitted for that path. + + Args: + prompt: Path-specific user prompt (grounding facts for one path). + system_prompt: Narrative persona system instruction from prompts.yaml. + + Returns: + A dict with keys ``narrative`` and ``business_impact``. Returns ``{}`` + on failure or when not overridden. Callers must reject any response + containing an ID pattern (T####, CVE-, CAPEC-, D3-) — the persona is + instructed never to emit one, so any that appear are a grounding + violation, not a fact to trust. + """ + return {} diff --git a/threat_analysis/ai_engine/providers/litellm_client.py b/threat_analysis/ai_engine/providers/litellm_client.py index adf1cfd..56a9666 100644 --- a/threat_analysis/ai_engine/providers/litellm_client.py +++ b/threat_analysis/ai_engine/providers/litellm_client.py @@ -224,8 +224,9 @@ async def generate_content(self, prompt: str, system_prompt: str, stream: Option "timeout": effective_timeout, "num_retries": int(self.ai_config.get("threat_generation", {}).get("num_retries", 3)), "api_base": self.api_base, - "enable_thinking": bool(self.provider_config.get("enable_thinking", False)), } + if "enable_thinking" in self.provider_config: + completion_params["enable_thinking"] = bool(self.provider_config["enable_thinking"]) if "top_p" in self.provider_config: completion_params["top_p"] = self.provider_config["top_p"] @@ -255,15 +256,26 @@ async def generate_content(self, prompt: str, system_prompt: str, stream: Option logging.debug(f"[LLM Response Time: {llm_call_end_time - llm_call_start_time:.4f}s] Model: {self.model_name}") full_response_content = "" + finish_reason = None if use_stream: async for chunk in response: if chunk.choices and chunk.choices[0].delta.content: content_chunk = chunk.choices[0].delta.content full_response_content += content_chunk yield content_chunk + if chunk.choices and getattr(chunk.choices[0], "finish_reason", None): + finish_reason = chunk.choices[0].finish_reason else: full_response_content = response.choices[0].message.content - + finish_reason = getattr(response.choices[0], "finish_reason", None) + + if finish_reason in ("length", "max_tokens"): + logging.warning( + "LLM response truncated by max_tokens (finish_reason=%s) — model=%s max_tokens=%s. " + "Raise max_tokens or batch_tokens_per_component in config.", + finish_reason, self.model_name, completion_params.get("max_tokens"), + ) + if output_format == "json": cleaned_content = extract_json_from_llm_response(full_response_content) if cleaned_content: @@ -273,7 +285,10 @@ async def generate_content(self, prompt: str, system_prompt: str, stream: Option logging.error(f"Failed to decode JSON from AI response: {e}") yield f"Error: Failed to decode JSON from AI response. Raw output: {full_response_content}" else: - yield f"Error: No valid JSON found in AI response. Raw output: {full_response_content}" + yield ( + f"Error: No valid JSON found in AI response " + f"(finish_reason={finish_reason}). Raw output: {full_response_content}" + ) else: yield full_response_content diff --git a/threat_analysis/ai_engine/providers/litellm_provider.py b/threat_analysis/ai_engine/providers/litellm_provider.py index c3d3a56..35f39ac 100644 --- a/threat_analysis/ai_engine/providers/litellm_provider.py +++ b/threat_analysis/ai_engine/providers/litellm_provider.py @@ -17,7 +17,6 @@ from .base_provider import BaseLLMProvider from .litellm_client import LiteLLMClient from ..prompts.stride_prompts import build_component_prompt, build_batch_prompt -from ..prompts.attack_flow_prompts import build_attack_flow_prompt from threat_analysis.ai_engine.prompt_loader import get as _get_prompt import json @@ -155,23 +154,6 @@ def _parse_batch(raw) -> Dict[str, List[Dict]]: logging.error("Error in batch threat generation: %s", e) return {} - async def generate_attack_flow(self, threat: Dict, component: Dict, context: Dict) -> Dict: - client = await self._get_client() - prompt = build_attack_flow_prompt(threat, component, context) - - try: - async for chunk in client.generate_content( - prompt=prompt, - system_prompt=_get_prompt("attack_flow", "system"), - output_format="json" - ): - if isinstance(chunk, dict): - return chunk - return {} - except Exception as e: - logging.error(f"Error generating attack flow via LiteLLM: {e}") - return {} - async def generate_ciso_triage(self, prompt: str, system_prompt: str) -> Dict: """Calls the LLM with the CISO persona and returns the parsed briefing.""" client = await self._get_client() @@ -211,11 +193,56 @@ async def generate_soc_analysis(self, batch_prompt: str, system_prompt: str) -> if isinstance(chunk, dict): # Some providers may wrap the array in a dict return chunk.get("results", chunk.get("threats", [])) + if isinstance(chunk, str) and chunk.startswith("Error:"): + # JSON extraction/parsing failed (e.g. truncated response) — + # previously fell through to `return []` silently, with no log at + # all, making SOC enrichment failures indistinguishable from "no + # threats to analyze" or "provider offline". + logging.warning("generate_soc_analysis: LLM returned error string: %s", chunk[:200]) + return [] + logging.warning("generate_soc_analysis: no parseable chunk received — returning []") return [] except Exception as exc: logging.error("SOC analysis generation failed: %s", exc) return [] + async def generate_debate_turn(self, prompt: str, system_prompt: str) -> Dict: + """Calls the LLM with a Red or Blue debate persona and returns the parsed turn.""" + client = await self._get_client() + try: + async for chunk in client.generate_content( + prompt=prompt, + system_prompt=system_prompt, + output_format="json", + ): + if isinstance(chunk, dict): + return chunk + return {} + except Exception as exc: + logging.error("Red/Blue debate turn generation failed: %s", exc) + return {} + + async def generate_attack_path_narrative(self, prompt: str, system_prompt: str) -> Dict: + """Calls the LLM with the attack-path-narrative persona and returns the parsed result. + + Grounding/ID-leakage rejection is the caller's responsibility (report_generator.py) — + this method only handles the LLM call and JSON parsing, same as the other optional + provider methods. + """ + client = await self._get_client() + try: + async for chunk in client.generate_content( + prompt=prompt, + system_prompt=system_prompt, + output_format="json", + ): + if isinstance(chunk, dict): + return chunk + return {} + except Exception as exc: + logging.error("Attack path narrative generation failed: %s", exc) + return {} + async def generate_markdown( self, prompt: str, diff --git a/threat_analysis/ai_engine/providers/ollama_provider.py b/threat_analysis/ai_engine/providers/ollama_provider.py index ccee63e..db40653 100644 --- a/threat_analysis/ai_engine/providers/ollama_provider.py +++ b/threat_analysis/ai_engine/providers/ollama_provider.py @@ -19,7 +19,6 @@ from typing import Dict, List from .base_provider import BaseLLMProvider from ..prompts.stride_prompts import build_component_prompt -from ..prompts.attack_flow_prompts import build_attack_flow_prompt from threat_analysis.ai_engine.prompt_loader import get as _get_prompt class OllamaProvider(BaseLLMProvider): @@ -29,7 +28,6 @@ def __init__(self, config: Dict): self.host = config.get("host", "http://localhost:11434") self.model = config.get("model", "mistral") self.temperature = config.get("temperature", 0.3) - self.attack_flow_temperature = config.get("attack_flow_temperature", 0.2) self.num_ctx = config.get("num_ctx", 4096) self.num_predict = config.get("num_predict", 4096) self.timeout = aiohttp.ClientTimeout(total=config.get("timeout", 120)) @@ -83,44 +81,3 @@ async def generate_threats(self, component: Dict, context: Dict) -> List[Dict]: logging.error(f"Received text: {result_text}") return [] return [] - - async def generate_attack_flow(self, threat: Dict, component: Dict, context: Dict) -> Dict: - prompt = build_attack_flow_prompt(threat, component, context) - result_text = "" - - max_retries = 3 - for attempt in range(max_retries): - async with aiohttp.ClientSession(timeout=self.timeout) as session: - try: - async with session.post( - f"{self.host}/api/generate", - json={ - "model": self.model, - "prompt": f"{_get_prompt('attack_flow', 'system')}\n\n{prompt}", - "format": "json", - "stream": False, - "options": { - "temperature": self.attack_flow_temperature, - "num_ctx": self.num_ctx, - "num_predict": self.num_predict - } - } - ) as response: - response.raise_for_status() - result_text = await response.text() - result = json.loads(result_text) - # The actual response is a JSON string inside the 'response' key - return json.loads(result.get('response', '{}')) - except aiohttp.ClientError as e: - if attempt < max_retries - 1: - wait_time = 2 ** attempt - logging.warning(f"Error connecting to Ollama for attack flow generation (attempt {attempt + 1}/{max_retries}): {e}. Retrying in {wait_time}s...") - await asyncio.sleep(wait_time) - continue - logging.error(f"Error connecting to Ollama for attack flow generation after {max_retries} attempts: {e}") - return {} - except json.JSONDecodeError as e: - logging.error(f"Error decoding JSON from Ollama for attack flow generation: {e}") - logging.error(f"Received text: {result_text}") - return {} - return {} diff --git a/threat_analysis/ai_engine/rag_service.py b/threat_analysis/ai_engine/rag_service.py index 9f17c86..fa83289 100644 --- a/threat_analysis/ai_engine/rag_service.py +++ b/threat_analysis/ai_engine/rag_service.py @@ -161,6 +161,13 @@ def _initialize_components(self): prefix = "ollama" if name == "ollama" else name llm_model = f"{prefix}/{config.get('model')}" llm_params['temperature'] = config.get('temperature', 0.5) + # Without this, litellm falls back to its own (often much lower) + # per-model default and RAG responses risk the same JSON truncation + # that the provider's own max_tokens tuning (ai_config.yaml) already + # fixed for component-level generation — RAG calls litellm.completion() + # directly, bypassing LiteLLMClient entirely, so it never saw that value. + if config.get('max_tokens'): + llm_params['max_tokens'] = config['max_tokens'] if name == "ollama": llm_params['api_base'] = config.get('host', 'http://localhost:11434') elif config.get('api_base'): @@ -276,11 +283,20 @@ def _esc(s: str) -> str: if user_threat_intelligence: ctx_sections.append(f"## User Threat Intelligence\n{_esc(user_threat_intelligence)}") - human_message = self._rag_human_template.format( - optional_context="\n\n".join(ctx_sections), - threat_model_markdown=_esc(threat_model_markdown), - context=_esc(context_text), - ) + try: + human_message = self._rag_human_template.format( + optional_context="\n\n".join(ctx_sections), + threat_model_markdown=_esc(threat_model_markdown), + context=_esc(context_text), + ) + except (KeyError, IndexError, ValueError) as e: + # A stray unescaped { or } in the template (e.g. an unescaped JSON example) + # makes str.format() raise instead of substituting — this must degrade like + # every other AI failure (return no RAG threats), not abort the whole + # enrichment pipeline that called this. + logger.error("RAG human_template formatting failed (bad template — check for " + "unescaped {{ / }} in prompts.yaml rag.human_template): %s", e) + return [] messages = [ {"role": "system", "content": self._rag_system_prompt}, {"role": "user", "content": human_message}, diff --git a/threat_analysis/config/prompts.yaml b/threat_analysis/config/prompts.yaml index c5709ba..7c9de88 100644 --- a/threat_analysis/config/prompts.yaml +++ b/threat_analysis/config/prompts.yaml @@ -16,11 +16,11 @@ dsl_generation: system: | You are an expert cybersecurity architect specializing in STRIDE threat modeling. - Your task is to produce a complete, valid threat model in the SecOpsTM Markdown DSL. + Your task is to produce a complete, valid system model in the SecOpsTM Markdown DSL. ## DSL Structure - # Threat Model: [Name] + # System Model: [Name] ## Description [Concise description of the system and its security context] @@ -272,118 +272,6 @@ stride_analysis: } ] -# --------------------------------------------------------------------------- -# 3. ATTACK FLOW GENERATION -# Used by: LiteLLMProvider.generate_attack_flow → build_attack_flow_prompt() -# Purpose: Generate a STIX 2.1 Attack Flow for a specific threat. -# --------------------------------------------------------------------------- -attack_flow: - system: | - You are an expert in cyber attack chain analysis and MITRE ATT&CK framework v14+. - You produce detailed Attack Flow diagrams in MITRE Attack Flow v3.0 / STIX 2.1 format. - - Your expertise: - - Full kill chain modeling (Initial Access → Impact) - - Lateral movement and pivot analysis - - Detection engineering (Sigma rules, MITRE D3FEND, data source mapping) - - STIX 2.1 objects and relationship graph - - Quality requirements: - - All ATT&CK technique IDs must be real and current (v14+). Do not invent IDs. - - Model the most realistic attack path, not the worst-case theoretical one. - - Every action must have at least one detection opportunity. - - Include a realistic failure/blocked path for the first action. - - 5 to 7 main actions — detailed enough to be actionable, concise enough to be readable. - - component_template: | - # Attack Flow — <> - - ## Threat to Model - | Field | Value | - |---|---| - | STRIDE Category | <> | - | Title | <> | - | Description | <> | - | Known MITRE Techniques | <> | - - ## Basic Scenario - <> - - ## Target Component - | Field | Value | - |---|---| - | Type | <> | - | Name | <> | - | Description | <> | - - ## System Context - <> - - ## Your Task - Generate a detailed Attack Flow for this threat in MITRE Attack Flow v3.0 / STIX 2.1. - - For each action include: - - ATT&CK Tactic (ID + name) - - ATT&CK Technique (ID + name + sub-technique if applicable) - - Realistic description of the adversary action - - Success path (next action ref) - - Failure/blocked path (detection node ref) - - Specific detection data sources (Windows Event IDs, Linux audit, cloud trail, etc.) - - Return ONLY valid JSON: - - { - "type": "attack-flow", - "spec_version": "3.0.0", - "id": "attack-flow--<>-{uuid4}", - "name": "<>", - "description": "Attack flow for <> on <>", - "scope": "incident", - "start_refs": ["action--1"], - "actions": [ - { - "type": "action", - "id": "action--1", - "name": "Action name", - "tactic": {"id": "TA0001", "name": "Initial Access"}, - "technique": { - "id": "T1190", - "name": "Exploit Public-Facing Application", - "subtechnique": null - }, - "description": "Detailed adversary action…", - "confidence": 85, - "success_refs": ["action--2"], - "failure_refs": ["detection--1"] - } - ], - "conditions": [ - { - "type": "condition", - "id": "condition--1", - "description": "Prerequisite (e.g. internet-facing service with no WAF)" - } - ], - "assets": [ - { - "type": "asset", - "id": "asset--1", - "name": "<>", - "description": "<>" - } - ], - "detection_points": [ - { - "type": "detection", - "id": "detection--1", - "name": "Detection method", - "description": "How blue team detects this action", - "data_sources": ["Windows Security Event 4625", "Cloudtrail: UnauthorizedOperation"], - "blocks_action": "action--1" - } - ] - } - # --------------------------------------------------------------------------- # 4. RAG SYSTEM-LEVEL THREAT GENERATION # Used by: RAGThreatGenerator (LangChain ChatPromptTemplate) @@ -409,10 +297,23 @@ rag: Prioritize threats that are SPECIFIC to the architecture described, grounded in the retrieved CVE/CAPEC patterns and your system-level understanding. + STRICT GROUNDING RULE — the retrieved CVE/CAPEC patterns below come from a large, + generic security knowledge base and frequently describe technology (Kubernetes, + cloud storage, specific databases, etc.) that is NOT part of this architecture. + - affected_components MUST only name components that literally appear in the + "## System Model (Architecture)" section below — never invent a component. + - Use retrieved patterns only for the ATTACK TECHNIQUE or PATTERN they illustrate + (e.g. "credential exfiltration via misconfigured access control"), not for the + specific technology named in them, unless that exact technology is also present + in the architecture below. + - If nothing in the retrieved context applies to components actually present in + this architecture, rely on your own system-level understanding instead — do not + force-fit an irrelevant retrieved pattern. + human_template: | {optional_context} - ## Threat Model (Architecture) + ## System Model (Architecture) {threat_model_markdown} ## Retrieved Security Knowledge (CVE / CAPEC patterns) @@ -439,7 +340,7 @@ rag: Format your response as a JSON array: [ - { + {{ "name": "Threat name referencing specific components", "description": "Multi-hop attack description…", "affected_components": ["ComponentA", "ComponentB"], @@ -449,7 +350,7 @@ rag: "impact": "high", "source": "LLM", "confidence": 0.75 - } + }} ] # --------------------------------------------------------------------------- @@ -538,21 +439,29 @@ ciso_triage: - Translating technical vulnerabilities into actionable risk posture statements - Identifying quick wins that reduce exposure with minimal engineering effort - Your task is to synthesise a ranked list of STRIDE threats identified in a security - assessment and produce a concise, actionable CISO-level risk briefing. + Your task is to synthesise a ranked list of STRIDE threats — and, when available, the + goal-driven attack scenarios (GDAF) and Red/Blue debate outcomes below them — into a + concise, actionable CISO-level risk briefing. Rules: - posture_score is a float 0–10 (10 = critical risk, 0 = no material risk). Weight it toward the severity and volume of unmitigated CRITICAL/HIGH threats. + A GDAF scenario the debate found "still viable after Blue's defences" should push + posture_score up; one "largely blocked by Blue" should not. - posture_label must be one of: "CRITICAL", "HIGH", "ELEVATED", "MODERATE", "LOW". - top_findings are the 5 most impactful risk areas (not individual threats — themes). - Each finding must include a "title", "detail" (2–3 sentences), and "threat_ids" - (list of threat IDs from the input that support it). + Each finding must include a "title", "detail" (2–3 sentences), "threat_ids" + (list of STRIDE threat IDs from the input that support it, e.g. "T-0001"), and + "scenario_ids" (list of GDAF scenario IDs that support it, e.g. "GDAF-..." — empty + list if the finding is not grounded in a GDAF scenario). When a finding is driven by + a GDAF scenario, mention in "detail" whether the Red/Blue debate confirmed it is + still exploitable or found it largely blocked. - quick_wins are 3–5 remediations with high impact-to-effort ratio. "impact" and "effort" must each be "HIGH", "MEDIUM", or "LOW". - - narrative is a 3–5 sentence executive summary suitable for a board slide. - - Be specific and factual — base findings on the provided threat data only. - - Do NOT invent threats not present in the input. + - narrative is a 3–5 sentence executive summary suitable for a board slide. If a GDAF + scenario remained viable after the debate, the narrative should say so explicitly. + - Be specific and factual — base findings on the provided threat and scenario data only. + - Do NOT invent threats or scenarios not present in the input. template: | ## Threat Summary for CISO Briefing @@ -565,6 +474,10 @@ ciso_triage: <> + ## Goal-Driven Attack Scenarios (GDAF) and Red/Blue Debate Outcomes + + <> + --- Produce a CISO briefing as a single JSON object — no markdown fence, no preamble: @@ -576,7 +489,8 @@ ciso_triage: "rank": 1, "title": "{concise risk theme title}", "detail": "{2-3 sentence explanation of the risk and its business impact}", - "threat_ids": ["{T-NNNN}", ...] + "threat_ids": ["{T-NNNN}", ...], + "scenario_ids": ["{GDAF-...}", ...] } ], "quick_wins": [ @@ -589,3 +503,154 @@ ciso_triage: ], "narrative": "{3-5 sentence executive summary for board presentation}" } + +# --------------------------------------------------------------------------- +# 7. RED/BLUE ADVERSARIAL DEBATE +# Used by: RedBlueDebateEngine / LiteLLMProvider.generate_debate_turn +# Purpose: Multi-round attack/defence debate over top GDAF attack scenarios. +# Output (red): viability_score, techniques_attempted, failed_alternatives, +# rationale, evidence[] +# Output (blue): viability_score, techniques_blocked, detection_gaps[], +# rationale, evidence[] +# --------------------------------------------------------------------------- +red_blue_debate: + red_system: | + You are a senior red team operator with deep expertise in MITRE ATT&CK techniques, + exploitation of misconfigurations, and adversarial path-finding. + + Rules: + - Only use facts explicitly listed in the grounding block you are given — never invent + a CVE, port, protocol, or misconfiguration that is not present there. + - Every claim you make must cite a specific piece of grounding evidence (CVE, misconfig, + port/service, or privilege boundary). + - If Blue's prior response blocks your previous approach, try a genuinely different + technique or hop. If none works, record the attempt under failed_alternatives — + this is valuable signal, not a failure to hide. + - viability_score reflects how likely THIS round's attempt is to succeed end-to-end, + from 0.0 (fully blocked) to 1.0 (fully viable). + + blue_system: | + You are a senior blue team / detection engineer with deep expertise in SIEM correlation + rules, EDR behavioral detection, IDS signatures, and MITRE D3FEND countermeasures. + + Rules: + - Only use facts explicitly listed in the grounding block you are given — never invent + a control that is not implied by it. + - For each Red technique, decide whether it is blocked/detected or not, and name the + specific control family: "SIEM", "EDR", or "IDS". + - detection_gaps lists steps you CANNOT cover, with covered=false and a specific reason. + - Every claim you make must cite a specific piece of grounding evidence. + - viability_score reflects your own estimate of the residual attack viability after your + response, from 0.0 (fully blocked) to 1.0 (fully viable). + + red_round_user: | + ## Attack Scenario + Objective: <> + Entry point: <> + Target asset: <> + + ## Grounding facts (use these — do not invent facts not listed here) + <> + + ## Blue team's prior response + <> + + --- + Attempt to compromise the target through this path. Return ONLY a JSON object, + no markdown fence, no preamble: + { + "viability_score": 0.7, + "techniques_attempted": ["T1190", "T1078"], + "failed_alternatives": ["Attempted T1110 password spraying — blocked by account lockout policy"], + "rationale": "short justification", + "evidence": [ + { + "claim": "SSH endpoint accepts password auth", + "evidence_type": "misconfig", + "evidence_ref": "hop 2 dataflow: protocol=ssh, is_authenticated=false", + "confidence": "high" + } + ] + } + + blue_round_user: | + ## Grounding facts (use these — do not invent facts not listed here) + <> + + ## Red team's attack this round + <> + + --- + Assess which of Red's techniques you can detect or block, and identify detection gaps + for anything you cannot cover. Return ONLY a JSON object, no markdown fence, no preamble: + { + "viability_score": 0.4, + "techniques_blocked": ["T1078"], + "detection_gaps": [ + { + "step": "hop 2 SSH access", + "control_family": "SIEM", + "covered": false, + "detail": "No correlation rule for password-auth SSH from external IP ranges", + "confidence": "medium" + } + ], + "rationale": "short justification", + "evidence": [ + { + "claim": "EDR blocks known T1078 credential reuse pattern on this host", + "evidence_type": "privilege_boundary", + "evidence_ref": "hop 2 asset has mfa_enabled=true per grounding", + "confidence": "high" + } + ] + } + +# --------------------------------------------------------------------------- +# 8. DISCOVERED ATTACK PATH NARRATIVE +# Used by: ReportGenerator._generate_path_narratives / LiteLLMProvider.generate_attack_path_narrative +# Purpose: Short grounded narrative explaining an already-computed attack path +# (AttackFlowGenerator.get_paths_summary — pytm+AI+LLM threats, no GDAF +# context needed). The path's hops/techniques/targets are fixed before +# this call — the model explains them, it does not design them. +# --------------------------------------------------------------------------- +attack_path_narrative: + system: | + You are a senior security analyst writing a short, grounded explanation of an + already-computed attack path for a technical threat report. + + STRICT GROUNDING RULE: + - The path's hops, techniques, and targets are fixed — listed in the grounding + block you are given. You are explaining them, not designing new ones. Do not + propose, rename, or add any hop, technique, CVE, CAPEC, or ATT&CK ID that is + not explicitly present in the grounding block. + - NEVER output an ID of any kind — no "T1234"-style technique IDs, no + "CVE-YYYY-NNNNN", no "CAPEC-NNN", no "D3-XXXX". Refer to techniques by name + only. The report already shows the exact IDs next to each hop; your job is + the narrative, not the taxonomy. Any output containing an ID pattern is + discarded entirely and logged as a grounding violation — so leaving IDs out + is not optional, it is the only way your response gets used. + - If you are not confident why a hop matters, say so plainly rather than + inventing a plausible-sounding justification. + + Output exactly two short paragraphs (2-4 sentences each): a narrative + explaining how an attacker would realistically move through this specific + path, and a business_impact paragraph explaining the consequence if it + succeeds. Nothing else — no preamble, no extra sections. + + template: | + ## Discovered Attack Path — <> + + Path score: <> | Hops: <> + Path: <> + + ## Grounding facts (use these — do not invent hops, techniques, or facts not listed here) + <> + + --- + Write the narrative and business_impact as instructed. Return ONLY a JSON + object, no markdown fence, no preamble: + { + "narrative": "...", + "business_impact": "..." + } diff --git a/threat_analysis/core/attack_id_validator.py b/threat_analysis/core/attack_id_validator.py index c0f93b1..c1f9506 100644 --- a/threat_analysis/core/attack_id_validator.py +++ b/threat_analysis/core/attack_id_validator.py @@ -13,20 +13,28 @@ # limitations under the License. """ -AttackIdValidator — validates ATT&CK technique IDs in generated threat reports. +AttackIdValidator — validates ATT&CK/CAPEC/D3FEND IDs in generated threat reports. -Checks every ``mitre_techniques[].id`` in the threat list against the local -``external_data/enterprise-attack.json`` corpus. Identifies three classes of -problems: +Checks every ``mitre_techniques[].id``, ``capecs[].capec_id``, and +``mitre_techniques[].defend_mitigations[].id`` in the threat list against the +local ``external_data/`` corpus. Identifies three classes of problems for +ATT&CK techniques (CAPEC/D3FEND only support "invalid" — the local corpora +have no revoked/deprecated concept): - **invalid** — ID not present in the corpus at all (likely an AI hallucination). - **revoked** — ID present but marked ``revoked: true`` (superseded by another technique). - **deprecated** — ID present but marked ``x_mitre_deprecated: true``. -The validator is fully offline: it reads the pre-committed STIX bundle and -never makes network calls. The JSON is parsed once per process and cached -at class level. +CWE IDs are intentionally NOT validated here: the pipeline never uses the raw +``cwe_ids`` field an LLM may return (see prompts.yaml) — CWEs that actually +reach a report always come from ``CVEService.get_cwes_for_cve()``, which reads +the committed CVE JSONL corpus directly. There is nothing LLM-supplied to +ground-truth for CWE. + +The validator is fully offline: it reads pre-committed data files and never +makes network calls. Each corpus is parsed once per process and cached at +class level. """ import json @@ -35,6 +43,8 @@ from pathlib import Path from typing import Dict, FrozenSet, List, Optional, Set, Tuple +from threat_analysis.core import data_loader + logger = logging.getLogger(__name__) # Path to the enterprise-attack STIX bundle (relative to this file) @@ -69,8 +79,14 @@ def label(self) -> str: @property def attack_url(self) -> str: - """Best-effort ATT&CK URL for a valid/revoked/deprecated ID.""" + """Best-effort reference URL for a valid/revoked/deprecated ID — + dispatches by ID namespace (ATT&CK technique, CAPEC, or D3FEND).""" tid = self.technique_id + if tid.upper().startswith("CAPEC-"): + num = tid.split("-", 1)[1] if "-" in tid else tid + return f"https://capec.mitre.org/data/definitions/{num}.html" + if tid.upper().startswith("D3-"): + return f"https://d3fend.mitre.org/technique/d3f:{tid}/" if "." in tid: parts = tid.split(".", 1) return f"https://attack.mitre.org/techniques/{parts[0]}/{parts[1]}/" @@ -107,14 +123,30 @@ def n_deprecated(self) -> int: class AttackIdValidator: - """Validates ATT&CK technique IDs against the local STIX bundle. + """Validates ATT&CK technique, CAPEC, and D3FEND IDs against local corpora. - The index is built lazily and cached at class level so it is only + Each index is built lazily and cached at class level so it is only parsed once per process, even when multiple report generators run. """ # Class-level cache: (valid_ids, revoked_ids, deprecated_ids) _index: Optional[Tuple[FrozenSet[str], FrozenSet[str], FrozenSet[str]]] = None + _capec_index: Optional[FrozenSet[str]] = None + _d3fend_index: Optional[FrozenSet[str]] = None + + @classmethod + def _load_capec_index(cls) -> FrozenSet[str]: + """Build (or return cached) set of known CAPEC IDs.""" + if cls._capec_index is None: + cls._capec_index = data_loader.load_capec_catalog_ids() + return cls._capec_index + + @classmethod + def _load_d3fend_index(cls) -> FrozenSet[str]: + """Build (or return cached) set of known D3FEND IDs.""" + if cls._d3fend_index is None: + cls._d3fend_index = frozenset(data_loader.load_d3fend_mapping().keys()) + return cls._d3fend_index @classmethod def _load_index(cls, stix_path: Path = _ENTERPRISE_ATTACK_PATH) -> Tuple[ @@ -172,13 +204,16 @@ def _load_index(cls, stix_path: Path = _ENTERPRISE_ATTACK_PATH) -> Tuple[ def _reset_cache(cls) -> None: """Reset class-level cache (for testing).""" cls._index = None + cls._capec_index = None + cls._d3fend_index = None def validate_all( self, all_threats: List[Dict], stix_path: Path = _ENTERPRISE_ATTACK_PATH, ) -> ValidationReport: - """Validate all ATT&CK technique IDs across the full threat list. + """Validate all ATT&CK technique, CAPEC, and D3FEND IDs across the full + threat list. Args: all_threats: List of threat dicts as returned by @@ -189,17 +224,15 @@ def validate_all( A :class:`ValidationReport` with per-issue details. """ valid_ids, revoked_ids, deprecated_ids = self._load_index(stix_path) - - # Empty index = corpus not available; skip silently - if not valid_ids: - return ValidationReport(total_techniques_checked=0) + capec_ids = self._load_capec_index() + d3fend_ids = self._load_d3fend_index() total = 0 invalid_issues: List[IdIssue] = [] revoked_issues: List[IdIssue] = [] deprecated_issues: List[IdIssue] = [] - # Deduplicate per (threat_id, technique_id) to avoid double-counting + # Deduplicate per (threat_id, id) to avoid double-counting seen: Set[Tuple[str, str]] = set() for threat in all_threats: @@ -207,33 +240,60 @@ def validate_all( threat_name = threat.get("name") or threat.get("description", "?") threat_target = threat.get("target", "?") - for tech in threat.get("mitre_techniques", []): - raw_id = (tech.get("id") or "").strip() + def _make_issue(raw_id: str) -> Optional[IdIssue]: + nonlocal total + raw_id = raw_id.strip() if not raw_id: - continue + return None total += 1 key = (threat_id, raw_id) if key in seen: - continue + return None seen.add(key) - - issue = IdIssue( + return IdIssue( technique_id=raw_id, - issue_type="", # filled below + issue_type="", # filled by caller threat_id=threat_id, threat_name=str(threat_name)[:80], threat_target=str(threat_target)[:60], ) - if raw_id not in valid_ids: - issue.issue_type = INVALID - invalid_issues.append(issue) - elif raw_id in revoked_ids: - issue.issue_type = REVOKED - revoked_issues.append(issue) - elif raw_id in deprecated_ids: - issue.issue_type = DEPRECATED - deprecated_issues.append(issue) + # ATT&CK techniques — valid_ids empty means the corpus isn't available. + if valid_ids: + for tech in threat.get("mitre_techniques", []): + issue = _make_issue(tech.get("id") or "") + if issue is None: + continue + if issue.technique_id not in valid_ids: + issue.issue_type = INVALID + invalid_issues.append(issue) + elif issue.technique_id in revoked_ids: + issue.issue_type = REVOKED + revoked_issues.append(issue) + elif issue.technique_id in deprecated_ids: + issue.issue_type = DEPRECATED + deprecated_issues.append(issue) + + # CAPEC patterns — no revoked/deprecated concept in the local catalog. + if capec_ids: + for capec in threat.get("capecs", []): + issue = _make_issue(capec.get("capec_id") or "") + if issue is None: + continue + if issue.technique_id not in capec_ids: + issue.issue_type = INVALID + invalid_issues.append(issue) + + # D3FEND mitigations attached to each mapped technique. + if d3fend_ids: + for tech in threat.get("mitre_techniques", []): + for mitigation in tech.get("defend_mitigations", []): + issue = _make_issue(mitigation.get("id") or "") + if issue is None: + continue + if issue.technique_id not in d3fend_ids: + issue.issue_type = INVALID + invalid_issues.append(issue) report = ValidationReport( total_techniques_checked=total, diff --git a/threat_analysis/core/data_loader.py b/threat_analysis/core/data_loader.py index d555bdb..4d95325 100644 --- a/threat_analysis/core/data_loader.py +++ b/threat_analysis/core/data_loader.py @@ -93,6 +93,50 @@ def load_capec_to_mitre_mapping() -> Dict[str, List[Dict[str, Any]]]: return capec_to_mitre +def load_capec_catalog_ids() -> "frozenset[str]": + """Loads the broadest known set of CAPEC IDs (e.g. "CAPEC-66") from every + committed CAPEC source — the ground truth used to validate LLM-cited CAPEC + IDs before they appear in a report or STIX export. + + Unions three sources because none alone is the full CAPEC catalog: + - CAPEC_VIEW_ATT&CK_Related_Patterns.csv — only patterns with a known ATT&CK + relationship (177 entries) — a narrow cross-reference view, NOT the master + catalog. Using this alone previously flagged legitimate pytm-derived CAPEC + IDs (sourced from stride_to_capec.json) as "invalid" — a false positive, + not a real hallucination or corpus drift. + - stride_to_capec.json — the STRIDE→CAPEC mapping pytm-derived threats + actually draw from (530 distinct IDs, the broadest single source here). + - capec_to_mitre_structured_mapping.json — CAPEC→ATT&CK technique mapping + (181 distinct IDs, mostly but not fully overlapping with the above). + """ + import csv + csv_path = Path(__file__).parent.parent / 'external_data' / 'CAPEC_VIEW_ATT&CK_Related_Patterns.csv' + ids = set() + try: + with open(csv_path, 'r', newline='', encoding='utf-8') as f: + reader = csv.reader(f) + next(reader, None) # header row — ID is always the first column + for row in reader: + if row and row[0].strip(): + ids.add(f"CAPEC-{row[0].strip()}") + except FileNotFoundError: + logging.error( + "CAPEC catalog CSV not found at %s. " + "Run: secopstm download-data", csv_path + ) + except Exception as e: + logging.error(f"Error processing CAPEC catalog CSV: {e}") + + for cat_entries in load_stride_to_capec_map().values(): + for entry in cat_entries: + capec_id = entry.get('capec_id') + if capec_id: + ids.add(capec_id) + + ids.update(load_capec_to_mitre_mapping().keys()) + + return frozenset(ids) + def load_stride_to_capec_map() -> Dict[str, List[Dict[str, str]]]: """Loads the STRIDE to CAPEC mapping from the JSON file.""" diff --git a/threat_analysis/core/debate_engine.py b/threat_analysis/core/debate_engine.py new file mode 100644 index 0000000..a8519ab --- /dev/null +++ b/threat_analysis/core/debate_engine.py @@ -0,0 +1,370 @@ +# Copyright 2025 ellipse2v +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +RedBlueDebateEngine — adversarial Red/Blue debate over GDAF attack scenarios. + +A single configured LLM provider plays two personas (Red: attack, Blue: defence) +across multiple rounds on the top-N highest-viability AttackScenario objects. +The debate is grounded on facts already present in the model (hop protocol/ +auth/encryption, BOM CVEs) — no network calls beyond the existing LLM provider. + +The final viability adjusts the scenario's path_score/risk_level in place +(never the STRIDE threat table). All processing degrades to an empty result +when the provider is unavailable or returns malformed data — never raises. +""" + +import logging +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional + +from threat_analysis.ai_engine.prompt_loader import get as _get_prompt +from threat_analysis.core.gdaf_engine import GDAFEngine, compute_risk_level + +logger = logging.getLogger(__name__) + + +@dataclass +class Evidence: + claim: str + evidence_type: str # "cve" | "misconfig" | "port_service" | "privilege_boundary" | "none" + evidence_ref: str + confidence: str # "high" | "medium" | "low" + verified: bool + + +@dataclass +class DetectionGap: + step: str + control_family: str # "SIEM" | "EDR" | "IDS" | "NONE" + covered: bool + detail: str + confidence: str + + +@dataclass +class DebateTurn: + role: str # "red" | "blue" + round_index: int + viability_score: float + techniques_attempted: List[str] = field(default_factory=list) + techniques_blocked: List[str] = field(default_factory=list) + detection_gaps: List[DetectionGap] = field(default_factory=list) + failed_alternatives: List[str] = field(default_factory=list) + evidence: List[Evidence] = field(default_factory=list) + rationale: str = "" + + +@dataclass +class DebateResult: + scenario_id: str + objective_name: str + entry_point: str + target_asset: str + rounds: List[DebateTurn] + blocked_paths: List[str] + residual_detection_gaps: List[DetectionGap] + red_failed_attempts: List[str] + round_count: int + convergence_delta: float + converged: bool + final_viability: float + residual_path_viable: bool + debate_factor: float = 1.0 + + +def _clamp(value: float, low: float = 0.0, high: float = 1.0) -> float: + return max(low, min(high, value)) + + +class RedBlueDebateEngine: + """Runs the Red/Blue debate loop over GDAF AttackScenario objects. + + ``provider`` must implement ``async generate_debate_turn(prompt, system_prompt) -> Dict`` + (see BaseLLMProvider.generate_debate_turn). ``bom_directory`` is optional — grounding + falls back to hop protocol/auth/encryption facts only when no BOM directory is given. + """ + + def __init__( + self, + provider: Any, + config: Optional[Dict] = None, + bom_directory: Optional[str] = None, + ): + self.provider = provider + cfg = config or {} + self.top_n: int = int(cfg.get("top_n", 5)) + self.min_viability_threshold: float = float(cfg.get("min_viability_threshold", 0.5)) + self.max_rounds: int = int(cfg.get("max_rounds", 3)) + self.viability_delta_threshold: float = float(cfg.get("viability_delta_threshold", 0.1)) + self.debate_factor_min: float = float(cfg.get("debate_factor_min", 0.5)) + self.debate_factor_max: float = float(cfg.get("debate_factor_max", 1.5)) + self._bom_directory = bom_directory + + async def run(self, scenarios: List[Any]) -> List[DebateResult]: + """Debates the top-N eligible scenarios and mutates them in place with the outcome.""" + selected = self._select_scenarios(scenarios) + results: List[DebateResult] = [] + for scenario in selected: + result = await self._debate_scenario(scenario) + if result is None: + continue + self._reinject_score(scenario, result) + results.append(result) + return results + + # ------------------------------------------------------------------ + # Selection + # ------------------------------------------------------------------ + + def _select_scenarios(self, scenarios: List[Any]) -> List[Any]: + critical = GDAFEngine.get_risk_thresholds().get("CRITICAL", 4.0) + eligible = [ + s for s in scenarios + if self._initial_viability(s, critical) >= self.min_viability_threshold + ] + eligible.sort(key=lambda s: s.path_score, reverse=True) + return eligible[: self.top_n] + + @staticmethod + def _initial_viability(scenario: Any, critical_threshold: float) -> float: + detection = float(getattr(scenario, "detection_coverage", 0.0)) + raw = (scenario.path_score / critical_threshold) * (1.0 - 0.5 * detection) + return _clamp(raw) + + # ------------------------------------------------------------------ + # Grounding (evidence-first — facts already in the model, no new lookups) + # ------------------------------------------------------------------ + + def _build_grounding(self, scenario: Any) -> str: + bom_loader = None + if self._bom_directory: + from threat_analysis.core.bom_loader import BOMLoader # lazy: avoid cost when unused + bom_loader = BOMLoader(self._bom_directory) + + lines: List[str] = [] + for hop in getattr(scenario, "hops", []): + auth = "authenticated" if hop.is_authenticated else "unauthenticated" + enc = "encrypted" if hop.is_encrypted else "unencrypted" + techniques = ", ".join(t.id for t in hop.techniques[:3]) if hop.techniques else "none mapped" + cve_str = "no known CVEs on file" + if bom_loader: + bom_data = bom_loader.get(hop.asset_name) + cves = bom_data.get("known_cves") or [] + if cves: + cve_str = f"known CVEs: {', '.join(cves)}" + lines.append( + f"- {hop.asset_name} [{hop.hop_position}]: protocol={hop.protocol or 'unknown'}, " + f"{auth}, {enc}, candidate techniques: {techniques}; {cve_str}" + ) + return "\n".join(lines) if lines else "No hop data available." + + # ------------------------------------------------------------------ + # Debate loop + # ------------------------------------------------------------------ + + async def _debate_scenario(self, scenario: Any) -> Optional[DebateResult]: + grounding = self._build_grounding(scenario) + critical = GDAFEngine.get_risk_thresholds().get("CRITICAL", 4.0) + prev_viability = self._initial_viability(scenario, critical) + final_viability = prev_viability + + rounds: List[DebateTurn] = [] + blue_prior = "No prior Blue response yet — this is the first round." + converged = False + convergence_delta = 0.0 + + for round_index in range(self.max_rounds): + red_turn = await self._run_turn(scenario, "red", round_index, grounding, blue_prior) + if red_turn is None: + break + rounds.append(red_turn) + + blue_turn = await self._run_turn( + scenario, "blue", round_index, grounding, self._summarize_turn(red_turn) + ) + if blue_turn is None: + break + rounds.append(blue_turn) + + viability = _clamp(red_turn.viability_score - 0.15 * len(blue_turn.techniques_blocked)) + convergence_delta = abs(viability - prev_viability) + blue_prior = self._summarize_turn(blue_turn) + final_viability = viability + + if convergence_delta < self.viability_delta_threshold: + converged = True + prev_viability = viability + break + prev_viability = viability + + if not rounds: + return None + + return self._summarize_result(scenario, rounds, final_viability, converged, convergence_delta) + + async def _run_turn( + self, + scenario: Any, + role: str, + round_index: int, + grounding: str, + prior_summary: str, + ) -> Optional[DebateTurn]: + try: + if role == "red": + system_prompt = _get_prompt("red_blue_debate", "red_system") + user_prompt = _get_prompt( + "red_blue_debate", "red_round_user", + objective=scenario.objective_name, + entry_point=scenario.entry_point, + target_asset=scenario.target_asset, + grounding=grounding, + blue_prior=prior_summary, + ) + else: + system_prompt = _get_prompt("red_blue_debate", "blue_system") + user_prompt = _get_prompt( + "red_blue_debate", "blue_round_user", + grounding=grounding, + red_attack=prior_summary, + ) + except KeyError as exc: + logger.warning("Red/Blue debate: prompt key missing (%s) — skipping.", exc) + return None + + try: + raw = await self.provider.generate_debate_turn(user_prompt, system_prompt) + except Exception as exc: + logger.warning( + "Red/Blue debate: %s turn failed for scenario %s: %s", + role, getattr(scenario, "scenario_id", "?"), exc, + ) + return None + + if not isinstance(raw, dict) or "viability_score" not in raw: + logger.warning( + "Red/Blue debate: malformed %s turn for scenario %s — skipping", + role, getattr(scenario, "scenario_id", "?"), + ) + return None + + try: + return DebateTurn( + role=role, + round_index=round_index, + viability_score=_clamp(float(raw.get("viability_score", 0.0))), + techniques_attempted=list(raw.get("techniques_attempted") or []), + techniques_blocked=list(raw.get("techniques_blocked") or []), + detection_gaps=self._parse_detection_gaps(raw.get("detection_gaps")), + failed_alternatives=list(raw.get("failed_alternatives") or []), + evidence=self._parse_evidence(raw.get("evidence")), + rationale=str(raw.get("rationale", "")), + ) + except (TypeError, ValueError) as exc: + logger.warning("Red/Blue debate: could not parse %s turn: %s", role, exc) + return None + + @staticmethod + def _parse_evidence(raw_list: Any) -> List[Evidence]: + result: List[Evidence] = [] + for item in raw_list or []: + if not isinstance(item, dict): + continue + etype = str(item.get("evidence_type", "none")).lower() + ref = str(item.get("evidence_ref", "")) + result.append(Evidence( + claim=str(item.get("claim", "")), + evidence_type=etype, + evidence_ref=ref, + confidence=str(item.get("confidence", "low")).lower(), + verified=bool(etype != "none" and ref), + )) + return result + + @staticmethod + def _parse_detection_gaps(raw_list: Any) -> List[DetectionGap]: + result: List[DetectionGap] = [] + for item in raw_list or []: + if not isinstance(item, dict): + continue + result.append(DetectionGap( + step=str(item.get("step", "")), + control_family=str(item.get("control_family", "none")).upper(), + covered=bool(item.get("covered", False)), + detail=str(item.get("detail", "")), + confidence=str(item.get("confidence", "low")).lower(), + )) + return result + + @staticmethod + def _summarize_turn(turn: DebateTurn) -> str: + if turn.role == "red": + techs = ", ".join(turn.techniques_attempted) or "none" + return f"Red attempted: {techs}. Rationale: {turn.rationale}" + blocked = ", ".join(turn.techniques_blocked) or "none" + return f"Blue blocked: {blocked}. Rationale: {turn.rationale}" + + def _summarize_result( + self, + scenario: Any, + rounds: List[DebateTurn], + final_viability: float, + converged: bool, + convergence_delta: float, + ) -> DebateResult: + blocked_paths: List[str] = [] + residual_gaps: List[DetectionGap] = [] + failed_attempts: List[str] = [] + for turn in rounds: + if turn.role == "red": + failed_attempts.extend(turn.failed_alternatives) + else: + blocked_paths.extend(turn.techniques_blocked) + residual_gaps.extend(g for g in turn.detection_gaps if not g.covered) + + residual_path_viable = final_viability >= self.min_viability_threshold + + return DebateResult( + scenario_id=getattr(scenario, "scenario_id", ""), + objective_name=getattr(scenario, "objective_name", ""), + entry_point=getattr(scenario, "entry_point", ""), + target_asset=getattr(scenario, "target_asset", ""), + rounds=rounds, + blocked_paths=blocked_paths, + residual_detection_gaps=residual_gaps, + red_failed_attempts=failed_attempts, + round_count=sum(1 for t in rounds if t.role == "red"), + convergence_delta=round(convergence_delta, 3), + converged=converged, + final_viability=round(final_viability, 3), + residual_path_viable=residual_path_viable, + debate_factor=1.0, + ) + + # ------------------------------------------------------------------ + # Scoring reinjection (GDAF scenario only — never the STRIDE threat table) + # ------------------------------------------------------------------ + + def _reinject_score(self, scenario: Any, result: DebateResult) -> None: + span = self.debate_factor_max - self.debate_factor_min + base = self.debate_factor_min + span * result.final_viability + factor = base if result.residual_path_viable else min(base, 1.0) + + scenario.path_score_pre_debate = scenario.path_score + scenario.path_score = round(scenario.path_score * factor, 2) + thresholds = GDAFEngine.get_risk_thresholds() + scenario.risk_level = compute_risk_level(scenario.path_score, thresholds) + scenario.debate_factor = round(factor, 2) + result.debate_factor = scenario.debate_factor diff --git a/threat_analysis/core/gdaf_engine.py b/threat_analysis/core/gdaf_engine.py index 10d2b79..7d4d7bf 100644 --- a/threat_analysis/core/gdaf_engine.py +++ b/threat_analysis/core/gdaf_engine.py @@ -55,6 +55,17 @@ _DETECTION_COVERAGE = {"none": 0.0, "low": 0.2, "medium": 0.5, "high": 0.8} +def compute_risk_level(path_score: float, thresholds: Dict[str, float]) -> str: + """Classify a GDAF path_score into CRITICAL/HIGH/MEDIUM/LOW using calibrated thresholds.""" + if path_score >= thresholds.get("CRITICAL", 4.0): + return "CRITICAL" + if path_score >= thresholds.get("HIGH", 2.8): + return "HIGH" + if path_score >= thresholds.get("MEDIUM", 1.8): + return "MEDIUM" + return "LOW" + + @dataclass class AttackHop: asset_name: str @@ -87,6 +98,8 @@ class AttackScenario: detection_coverage: float # 0.0 (no detection controls mapped) unacceptable_risk: bool min_technique_score: float = 0.8 # threshold for OR-branch rendering in .afb + path_score_pre_debate: Optional[float] = None # set by RedBlueDebateEngine; None if never debated + debate_factor: float = 1.0 # multiplier applied to path_score by RedBlueDebateEngine class GDAFEngine: @@ -113,6 +126,17 @@ def _load_scoring_config(cls) -> Dict: cls._scoring_config = {} return cls._scoring_config + @classmethod + def get_risk_thresholds(cls) -> Dict[str, float]: + """Return merged risk_thresholds (scoring_config.yaml overrides over built-in defaults).""" + cfg = cls._load_scoring_config().get("gdaf", {}) + rt = cfg.get("risk_thresholds", {}) + return { + "CRITICAL": float(rt.get("CRITICAL", 4.0)), + "HIGH": float(rt.get("HIGH", 2.8)), + "MEDIUM": float(rt.get("MEDIUM", 1.8)), + } + def __init__( self, threat_model: Any, @@ -684,12 +708,7 @@ def _build_scenario( # Thresholds calibrated for avg_tech_score × hop_weight scoring: # hop_weight ~ 1.0–2.0, avg_tech_score ~ 1.0–2.5 → hop_score ~ 1.0–5.0 # path_score (average) ~ 1.0–5.0 + CIA bonus 0–0.5 - risk_level = ( - "CRITICAL" if path_score >= self._risk_thresholds["CRITICAL"] else - "HIGH" if path_score >= self._risk_thresholds["HIGH"] else - "MEDIUM" if path_score >= self._risk_thresholds["MEDIUM"] else - "LOW" - ) + risk_level = compute_risk_level(path_score, self._risk_thresholds) target_name = path[-1][0] diff --git a/threat_analysis/core/mitre_mapping_module.py b/threat_analysis/core/mitre_mapping_module.py index 174ba09..922b332 100644 --- a/threat_analysis/core/mitre_mapping_module.py +++ b/threat_analysis/core/mitre_mapping_module.py @@ -168,6 +168,36 @@ def _get_d3fend_mitigations_for_mitre_id(self, mitigation_id: str) -> List[Dict[ break return d3fend_mitigations + _CAPEC_RELEVANCE_MAX = 10 # cap on STRIDE-category-derived CAPECs per threat (see _rank_capecs_by_relevance) + + @staticmethod + def _rank_capecs_by_relevance(threat_description: str, capec_list: List[Dict[str, Any]], + max_results: int = _CAPEC_RELEVANCE_MAX) -> List[Dict[str, Any]]: + """Narrows a STRIDE-category CAPEC bucket (up to ~150 entries) to the patterns whose + own description shares vocabulary with the threat's description. + + Without this, every threat lacking explicit ``capec_ids`` (i.e. every plain pytm + threat — only SecOpsTM's custom_threats.py rules set capec_ids) gets the *entire* + STRIDE-category bucket verbatim, most of it unrelated to the actual threat (e.g. a + network tampering threat tagged with "Buffer Overflow via Environment Variables"). + Offline word-overlap ranking, same approach as ThreatConsolidator's dedup — no new + dependency, deterministic. + """ + if not capec_list: + return capec_list + if not threat_description or len(capec_list) <= max_results: + return capec_list + desc_words = set(re.findall(r"[a-z]{4,}", threat_description.lower())) + if not desc_words: + return capec_list[:max_results] + scored = [ + (len(desc_words & set(re.findall(r"[a-z]{4,}", c.get('description', '').lower()))), c) + for c in capec_list + ] + scored.sort(key=lambda pair: pair[0], reverse=True) + top = [c for score, c in scored[:max_results] if score > 0] + return top if top else capec_list[:max_results] + def map_threat_to_mitre(self, threat: Dict[str, Any]) -> Dict[str, List[Dict[str, Any]]]: """ Maps a threat to MITRE ATT&CK techniques and CAPEC patterns. @@ -179,8 +209,11 @@ def map_threat_to_mitre(self, threat: Dict[str, Any]) -> Dict[str, List[Dict[str explicit_capec_ids = threat.get("capec_ids", []) if not explicit_capec_ids: - # No CAPEC IDs provided — derive from STRIDE category - capec_list = self.stride_to_capec.get(stride_category, []) + # No CAPEC IDs provided — derive from STRIDE category, narrowed to the patterns + # relevant to this specific threat (see _rank_capecs_by_relevance). + capec_list = self._rank_capecs_by_relevance( + threat.get("description", ""), self.stride_to_capec.get(stride_category, []) + ) direct_capec_ids = [c['capec_id'] for c in capec_list] for capec_info in capec_list: if capec_info['capec_id'] not in found_capecs: @@ -227,7 +260,9 @@ def map_threat_to_mitre(self, threat: Dict[str, Any]) -> Dict[str, List[Dict[str # Fallback: if explicit CAPEC IDs were provided but none mapped to ATT&CK techniques, # fall back to the stride_category lookup so reports are never empty. if explicit_capec_ids and not found_techniques and stride_category: - capec_list = self.stride_to_capec.get(stride_category, []) + capec_list = self._rank_capecs_by_relevance( + threat.get("description", ""), self.stride_to_capec.get(stride_category, []) + ) for capec_info in capec_list: cid = capec_info['capec_id'] if cid not in found_capecs: diff --git a/threat_analysis/core/models_module.py b/threat_analysis/core/models_module.py index 8dd097b..31c35e3 100644 --- a/threat_analysis/core/models_module.py +++ b/threat_analysis/core/models_module.py @@ -372,28 +372,6 @@ def _apply_custom_threats(self) -> List[Tuple[CustomThreat, Any]]: return generated_custom_threats - def _expand_class_targets(self, threats: List[Any]) -> List[Tuple[Any, Any]]: - """ - Expands threats that target a class (e.g., Server, Actor) into separate threats - for each instance of that class in the model. - This method is now generic and uses the _component_collections registry. - """ - expanded_threats = [] - for threat in threats: - target = getattr(threat, 'target', None) - - if isinstance(target, type) and target in self._component_collections: - collection = self._component_collections[target] - for item_info in collection: - instance = item_info['object'] - new_threat = threat.__class__(**threat.__dict__) - new_threat.target = instance - expanded_threats.append((new_threat, instance)) - else: - expanded_threats.append((threat, target)) - - return expanded_threats - # Valid STRIDE categories — threats outside this set are excluded from grouping. _VALID_STRIDE_CATEGORIES: frozenset = _STRIDE_CATEGORIES_CONST @@ -475,6 +453,10 @@ def get_statistics(self) -> Dict[str, int]: def add_severity_multiplier(self, element_name: str, multiplier: float): """Adds a severity multiplier for a given element.""" self.severity_multipliers[element_name] = multiplier + # Keep this model's own SeverityCalculator (used for custom-rule threat scoring, + # see _apply_custom_threats) in sync — it starts empty since it has no file/content + # to load from at construction time (ThreatModel.__init__ runs before parsing). + self.severity_calculator.target_multipliers[element_name] = multiplier diff --git a/threat_analysis/core/report_serializer.py b/threat_analysis/core/report_serializer.py index 7f4ae80..aea87c8 100644 --- a/threat_analysis/core/report_serializer.py +++ b/threat_analysis/core/report_serializer.py @@ -100,8 +100,13 @@ def _serialize_model(threat_model: Any) -> Dict: @staticmethod def _serialize_threat(t: Dict, index: int) -> Dict: severity = t.get("severity") or {} + # Reuse a stable id pre-assigned upstream (e.g. by generate_html_report(), so the + # CISO triage prompt can cite ids that actually match the final report) — otherwise + # fall back to positional numbering as before. + existing_id = t.get("id") + stable_id = existing_id if isinstance(existing_id, str) and existing_id.startswith("T-") else f"T-{index + 1:04d}" return { - "id": f"T-{index + 1:04d}", + "id": stable_id, "description": str(t.get("description", "")), "source": t.get("source", "pytm"), "stride_category": t.get("stride_category", ""), diff --git a/threat_analysis/generation/attack_flow_generator.py b/threat_analysis/generation/attack_flow_generator.py index 6cb60ac..8da05ca 100644 --- a/threat_analysis/generation/attack_flow_generator.py +++ b/threat_analysis/generation/attack_flow_generator.py @@ -28,7 +28,22 @@ class AttackFlowGenerator: Each path shows the sequence of actions and the assets they target. """ - def __init__(self, threats, model_name="Attack Flow"): + _DEFAULT_ALLOWED_CATEGORIES = { + "Spoofing", + "Tampering", + "Repudiation", + "Information Disclosure", + "Denial of Service", + "Elevation of Privilege", + } + + def __init__(self, threats, model_name="Attack Flow", allowed_categories=None): + """Args: + allowed_categories: STRIDE categories to build paths for — defaults to all + six. Callers may pass config/ai_config.yaml's + attack_flows.generate_for_categories to narrow this (e.g. to skip + Spoofing/Denial of Service/Repudiation paths as noise). + """ self.model_name = model_name # Filter out generic/class-based threats as requested @@ -41,14 +56,7 @@ def __init__(self, threats, model_name="Attack Flow"): continue filtered_threats.append(threat) - allowed_categories = { - "Spoofing", - "Tampering", - "Repudiation", - "Information Disclosure", - "Denial of Service", - "Elevation of Privilege" - } + allowed_categories = set(allowed_categories) if allowed_categories else self._DEFAULT_ALLOWED_CATEGORIES self.threats = [ threat for threat in filtered_threats if threat.get("stride_category") in allowed_categories @@ -131,28 +139,30 @@ def find_paths_recursive(current_path, phase_idx): find_paths_recursive([], 0) return all_paths - def generate_and_save_flows(self, output_dir): - afb_output_dir = os.path.join(output_dir, "afb") - os.makedirs(afb_output_dir, exist_ok=True) - + def _compute_best_paths(self): + """Finds the highest-severity attack path per STRIDE category. + + Shared by generate_and_save_flows() (writes .afb files) and + get_paths_summary() (compact data for HTML rendering) — same selection, + two different output formats. + """ attack_paths = self._find_attack_paths() if not attack_paths: - print("INFO: No logical attack paths found based on tactic progression.") - return + return {} paths_by_objective = defaultdict(list) final_objectives = {"Spoofing", "Tampering", "Repudiation", "Information Disclosure", "Denial of Service", "Elevation of Privilege"} - + for path in attack_paths: if not path: continue - + path_objectives = set() for _, threat in path: objective = threat.get('stride_category') if objective in final_objectives: path_objectives.add(objective) - + for objective in path_objectives: paths_by_objective[objective].append(path) @@ -160,26 +170,71 @@ def generate_and_save_flows(self, output_dir): for objective, path_list in paths_by_objective.items(): best_path_for_objective = None max_score = -1 - + for path in path_list: - current_score = sum(threat.get('severity', {}).get('score', 0) for _, threat in path) - + current_score = sum(threat.get('severity', {}).get('score') or 0 for _, threat in path) + if current_score > max_score: max_score = current_score best_path_for_objective = path - + if best_path_for_objective: best_paths_data[objective] = {"path": best_path_for_objective, "objectives": [objective]} + return best_paths_data + + def get_paths_summary(self): + """Returns a compact, JSON-serialisable summary of the best attack path per + STRIDE category — for HTML rendering, without writing any .afb file. + """ + best_paths_data = self._compute_best_paths() + summary = [] + for category, path_data in sorted(best_paths_data.items()): + path = path_data["path"] + hops, threat_ids, total_score = [], [], 0.0 + for tech_id, threat in path: + tech_data = self.techniques[tech_id] + target = self._get_target_name(threat.get('target')) + tid = threat.get('id') + hops.append({ + "technique_id": tech_id, + "technique_name": tech_data['name'], + "tactic": tech_data['tactics'][0] if tech_data['tactics'] else "", + "target": target, + # Per-hop threat id — do not derive this from threat_ids[i] by + # position, that list skips threats with no id and would misalign. + "threat_id": tid, + "threat_description": str(threat.get('description', ''))[:200], + }) + if tid: + threat_ids.append(tid) + total_score += threat.get('severity', {}).get('score') or 0 + if not hops: + continue + summary.append({ + "stride_category": category, + "path_score": round(total_score, 2), + "hop_count": len(hops), + "hops": hops, + "threat_ids": threat_ids, + "path_label": " → ".join(h["target"] for h in hops), + }) + return summary + + def generate_and_save_flows(self, output_dir): + afb_output_dir = os.path.join(output_dir, "afb") + os.makedirs(afb_output_dir, exist_ok=True) + + best_paths_data = self._compute_best_paths() if not best_paths_data: - print("INFO: No valid attack paths found after optimization.") + print("INFO: No logical attack paths found based on tactic progression.") return print(f"INFO: Found {len(best_paths_data)} unique attack paths after optimization.") i = 0 for objective, path_data in best_paths_data.items(): path = path_data["path"] - + flow_data = self._generate_single_path_flow(path, i + 1) if flow_data: file_path = os.path.join(afb_output_dir, f"optimized_path_{objective}.afb") diff --git a/threat_analysis/generation/diagram_generator.py b/threat_analysis/generation/diagram_generator.py index 598ef77..67b2ab4 100644 --- a/threat_analysis/generation/diagram_generator.py +++ b/threat_analysis/generation/diagram_generator.py @@ -22,6 +22,7 @@ import logging import json import datetime +import tempfile import xml.etree.ElementTree as ET from urllib.parse import urlparse # Added line from typing import Callable, Dict, List, Optional @@ -124,9 +125,16 @@ def generate_diagram_from_dot(self, dot_code: str, output_file: str, format: str return None except subprocess.CalledProcessError as e: - with open("/tmp/graphviz_error.log", "w") as f: - f.write(e.stderr) - logging.error(f"❌ Graphviz error: {e.stderr}") + # tempfile.mkstemp() (not a fixed /tmp path): concurrent requests would + # otherwise overwrite each other's error dump, and a fixed world-writable + # path is vulnerable to a symlink attack on a shared host. + try: + fd, error_log_path = tempfile.mkstemp(prefix="graphviz_error_", suffix=".log") + with os.fdopen(fd, "w") as f: + f.write(e.stderr) + logging.error(f"❌ Graphviz error (details in {error_log_path}): {e.stderr}") + except OSError as log_exc: + logging.error(f"❌ Graphviz error (could not write error log: {log_exc}): {e.stderr}") print(f"Graphviz error: {e.stderr}") logging.error(f"DOT code preview: {cleaned_dot[:200]}...") print(f"DOT code preview: {cleaned_dot[:200]}...") @@ -481,7 +489,8 @@ def _get_element_name(self, element) -> Optional[str]: # Try to convert to string as last resort try: return str(element) - except: + except Exception as exc: + logging.warning("Could not convert element to string for name extraction: %s", exc) return None def _extract_data_info(self, dataflow) -> Optional[str]: diff --git a/threat_analysis/generation/report_generator.py b/threat_analysis/generation/report_generator.py index 401f92b..1fdfe9f 100644 --- a/threat_analysis/generation/report_generator.py +++ b/threat_analysis/generation/report_generator.py @@ -21,6 +21,7 @@ import json import logging import sys +import importlib from typing import Dict, List, Any, Optional, Set from datetime import datetime import webbrowser @@ -33,7 +34,7 @@ from threat_analysis.core.cve_service import CVEService from threat_analysis.core.threat_ranker import rank_and_trim from threat_analysis.core.accepted_risks import AcceptedRiskLoader, compute_threat_key -from .utils import extract_name_from_object, get_target_name +from .utils import extract_name_from_object, get_target_name, get_enriched_threats import yaml import asyncio @@ -57,6 +58,21 @@ from threat_analysis.core.report_serializer import ReportSerializer from threat_analysis.severity_calculator_module import RiskContext +# Attack-path narrative persona (config/prompts.yaml → attack_path_narrative) is +# instructed to never emit an ID — the path's hops/techniques are already fixed by +# AttackFlowGenerator before the LLM sees them, so any ID in free text is a grounding +# violation, not a fact worth trusting. Matched independently of the persona's +# cooperation — this is enforced in code, not just requested in the prompt. +_ID_LEAKAGE_PATTERN = re.compile( + r"\bT\d{4}(?:\.\d{3})?\b|\bCVE-\d{4}-\d+\b|\bCAPEC-\d+\b|\bD3-[A-Z]+\b" +) + + +def _narrative_has_id_leakage(text: str) -> bool: + """True if free text contains anything shaped like an ATT&CK/CVE/CAPEC/D3FEND ID.""" + return bool(text) and bool(_ID_LEAKAGE_PATTERN.search(text)) + + def _resolve_active_cves( component_name: str, vex_loader: Optional[Any], @@ -271,7 +287,8 @@ def __init__(self, severity_calculator, mitre_mapping, cve_service: Optional[CVEService] = None, ai_config_path: Optional[Path] = None, context_path: Optional[Path] = None, # kept for backwards compat, unused - threat_model_ref: Optional[ThreatModel] = None): + threat_model_ref: Optional[ThreatModel] = None, + disable_rag: bool = False): self.severity_calculator = severity_calculator self.mitre_mapping = mitre_mapping self.env = Environment(loader=FileSystemLoader(Path(__file__).parent.parent / 'templates'), extensions=['jinja2.ext.do']) @@ -284,6 +301,9 @@ def __init__(self, severity_calculator, mitre_mapping, self.cve_service = cve_service if cve_service else CVEService(project_root, project_root / "cve_definitions.yml") self.ai_provider = None self.ai_context = None + self._ai_config_path = ai_config_path + self._ai_service = None # lazily built by _get_ai_service() + self._disable_rag = disable_rag self.threat_model_ref = threat_model_ref # Store the reference # Threat ranking / volume control — defaults (overridden from ai_config below) self._ranking_max_total: int = 0 @@ -291,6 +311,8 @@ def __init__(self, severity_calculator, mitre_mapping, self._ranking_weights: Dict[str, float] = {} self._enrich_batch_size: int = 5 self._enrich_max_concurrent: int = 3 + self._debate_config: Dict = {} + self._attack_flows_config: Dict = {} if ai_config_path and ai_config_path.exists(): with open(ai_config_path, "r", encoding="utf-8") as f: @@ -325,12 +347,23 @@ def __init__(self, severity_calculator, mitre_mapping, self._ranking_weights = {k: float(v) for k, v in rw.items() if isinstance(v, (int, float))} self._enrich_batch_size: int = int(tg.get("batch_size", 5)) self._enrich_max_concurrent: int = int(tg.get("max_concurrent_ai_requests", 3)) + self._debate_config = ai_config.get("debate", {}) or {} + self._attack_flows_config = ai_config.get("attack_flows", {}) or {} - async def _run_ciso_triage(self, all_threats: List[Dict]) -> Dict: + async def _run_ciso_triage( + self, + all_threats: List[Dict], + gdaf_scenarios: Optional[List[Any]] = None, + debate_results: Optional[List[Any]] = None, + ) -> Dict: """Generates a CISO-level risk briefing via the AI provider. - Returns an empty dict when AI is unavailable or the provider does not - implement ``generate_ciso_triage``. Never raises. + ``gdaf_scenarios``/``debate_results`` are optional — when provided (post-debate, + since the caller runs the debate before this), the briefing can reference + goal-driven attack paths and how the Red/Blue debate changed their risk level, + not just the flat STRIDE threat list. Returns an empty dict when AI is + unavailable or the provider does not implement ``generate_ciso_triage``. + Never raises. """ if not self.ai_provider: return {} @@ -383,6 +416,28 @@ async def _run_ciso_triage(self, all_threats: List[Dict]) -> Dict: name = t.get("name") or t.get("description", "?") threats_summary_lines.append(f"- [{sev}] {tid} | {stride} | {target} | {name[:80]}") + # Build a compact GDAF attack-scenario summary, noting debate outcomes when a + # scenario's score was adjusted by the Red/Blue debate. + debate_by_scenario = {getattr(r, "scenario_id", None): r for r in (debate_results or [])} + gdaf_lines: List[str] = [] + for scenario in (gdaf_scenarios or [])[:10]: + sid = getattr(scenario, "scenario_id", "?") + obj = getattr(scenario, "objective_name", "?") + actor = getattr(scenario, "actor_name", "?") + risk = getattr(scenario, "risk_level", "?") + score = getattr(scenario, "path_score", 0.0) or 0.0 + pre_score = getattr(scenario, "path_score_pre_debate", None) + line = f"- {sid} | Objective: {obj} | Actor: {actor} | Risk: {risk} | Score: {score:.2f}" + debate_result = debate_by_scenario.get(sid) + if debate_result is not None and pre_score is not None: + outcome = "path still viable after Blue's defences" if getattr(debate_result, "residual_path_viable", False) else "largely blocked by Blue" + line += ( + f" (was {pre_score:.2f} pre-debate — Red/Blue debate over " + f"{getattr(debate_result, 'round_count', 0)} round(s): {outcome})" + ) + gdaf_lines.append(line) + gdaf_summary = "\n".join(gdaf_lines) if gdaf_lines else "No GDAF attack scenarios were generated for this model." + stride_breakdown = ", ".join(f"{cat}={cnt}" for cat, cnt in sorted(stride_counts.items())) prompt = ( template @@ -393,6 +448,7 @@ async def _run_ciso_triage(self, all_threats: List[Dict]) -> Dict: .replace("<>", str(counts.get("LOW", 0))) .replace("<>", stride_breakdown) .replace("<>", "\n".join(threats_summary_lines)) + .replace("<>", gdaf_summary) ) try: @@ -418,194 +474,177 @@ async def _run_ciso_triage(self, all_threats: List[Dict]) -> Dict: ) return result - @staticmethod - def _build_ai_context_from_model(threat_model: "ThreatModel") -> Dict[str, Any]: - """Build AI context dict from DSL ## Context keys on the threat model. - - Priority: DSL context_config keys > context/*.yaml (loaded by AIService). - The report_generator path only uses DSL keys (context/*.yaml is handled - by AIService in server mode). + async def _generate_path_narratives(self, discovered_attack_paths: List[Dict]) -> None: + """Adds a grounded narrative + business_impact to each discovered attack path, + in place. One LLM call per path — naturally capped at one path per STRIDE + category (at most 6), since that's all AttackFlowGenerator.get_paths_summary() + ever returns. Never raises; skips silently when AI is unavailable, disabled via + config, or the provider does not implement the persona. + + Grounding: each path's hops/techniques/targets are already fixed by + AttackFlowGenerator before this call — the persona explains them, it does not + design them, and is instructed to never emit an ID. Any response that does + anyway (T-numbers, CVE, CAPEC, D3-) is discarded entirely, not partially + trusted — the code has no way to tell a correct ID from a hallucinated one in + ungrounded free text, so the only safe move is to drop the whole response. """ - ctx_cfg = getattr(threat_model, "context_config", {}) - _AI_KEYS = { - "system_description", "sector", "deployment_environment", - "data_sensitivity", "internet_facing", "user_base", - "compliance_requirements", "integrations", - } - ctx: Dict[str, Any] = {k: v for k, v in ctx_cfg.items() if k in _AI_KEYS} - ctx.setdefault("system_description", getattr(getattr(threat_model, "tm", None), "description", "") or "") - ctx.setdefault("data_sensitivity", "High") - return ctx + if not discovered_attack_paths or not self.ai_provider: + return + if not self._attack_flows_config.get("include_narrative", True): + return + try: + _client = await self.ai_provider._get_client() + if not _client.ai_online: + return + except Exception: + try: + if not await self.ai_provider.check_connection(): + return + except Exception: + pass # proceed; generate_attack_path_narrative() will fail safely if offline - async def _enrich_threats_with_ai(self, threat_model: ThreatModel, all_threats: List[Dict], progress_callback = None) -> List[Dict]: - if not self.ai_provider: - logging.warning("AI enrichment skipped: No AI provider initialized.") - return all_threats + from threat_analysis.ai_engine.prompt_loader import get as _get_prompt + try: + system_prompt = _get_prompt("attack_path_narrative", "system") + template = _get_prompt("attack_path_narrative", "template") + except KeyError as exc: + logging.warning("Attack path narrative: prompt key missing (%s) — skipping.", exc) + return + + for path in discovered_attack_paths: + try: + hops_lines = [] + for i, hop in enumerate(path.get("hops", []), start=1): + line = f"{i}. {hop.get('target', '?')} — technique: {hop.get('technique_name', '?')}" + if hop.get("tactic"): + line += f" (tactic: {hop['tactic']})" + threat_id = hop.get("threat_id") + threat_desc = hop.get("threat_description") + if threat_id or threat_desc: + line += f"\n Related threat: {threat_id or '?'} — {threat_desc or ''}" + hops_lines.append(line) + hops_detail = "\n".join(hops_lines) if hops_lines else "No hop data available." + + prompt = ( + template + .replace("<>", str(path.get("stride_category", "?"))) + .replace("<>", str(path.get("path_score", 0))) + .replace("<>", str(path.get("hop_count", 0))) + .replace("<>", str(path.get("path_label", "?"))) + .replace("<>", hops_detail) + ) + + result = await self.ai_provider.generate_attack_path_narrative(prompt, system_prompt) + if not isinstance(result, dict): + continue - # Build context from DSL ## Context keys (replaces config/context.yaml) - self.ai_context = self._build_ai_context_from_model(threat_model) - if not self.ai_context.get("system_description") and not self.ai_context.get("sector"): - logging.info("AI enrichment: no AI context keys in ## Context — threats will be generic.") + narrative = result.get("narrative") + business_impact = result.get("business_impact") + if not narrative and not business_impact: + continue + + if _narrative_has_id_leakage(narrative or "") or _narrative_has_id_leakage(business_impact or ""): + logging.warning( + "Attack path narrative for '%s' discarded — response contained an " + "ID pattern, which the persona is grounded never to emit.", + path.get("stride_category", "?"), + ) + continue + + path["narrative"] = narrative + path["business_impact"] = business_impact + except Exception as exc: + logging.warning( + "Attack path narrative generation failed for '%s' (non-fatal): %s", + path.get("stride_category", "?"), exc, + ) + + async def _run_debate(self, threat_model: Any) -> List[Any]: + """Runs the Red/Blue adversarial debate over GDAF attack scenarios. + Returns an empty list when debate is disabled, AI is unavailable, or no + GDAF scenarios exist. Never raises — mirrors _run_ciso_triage's + degrade-silently contract. + """ + if not self._debate_config.get("enabled", False): + return [] + if not self.ai_provider: + return [] - logging.info(f"Enriching threats with AI using provider: {type(self.ai_provider).__name__}") + scenarios = getattr(threat_model, "gdaf_scenarios", None) or [] + if not scenarios: + return [] - # Use the cached ai_online flag from the underlying client — avoids a second - # network round-trip (the client already ran check_connection() during create()). try: _client = await self.ai_provider._get_client() if not _client.ai_online: - logging.warning( - "AI enrichment skipped: Provider %s is offline (cached state).", - type(self.ai_provider).__name__, - ) - return all_threats + return [] except Exception: - pass # If we can't introspect the client, proceed and let generate_threats() fail naturally - - # Combine all components to enrich: servers, actors, and boundaries - components_to_enrich = [] - for server in threat_model.servers: - components_to_enrich.append({"name": server.get("name"), "type": "Server", "description": server.get("description", ""), "business_value": server.get("business_value")}) - for actor in threat_model.actors: - components_to_enrich.append({"name": actor.get("name"), "type": "Actor", "description": actor.get("description", ""), "business_value": actor.get("business_value")}) - for b_name, b_info in threat_model.boundaries.items(): - components_to_enrich.append({"name": b_name, "type": "Boundary", "description": b_info.get("description", ""), "business_value": b_info.get("business_value")}) + try: + if not await self.ai_provider.check_connection(): + return [] + except Exception: + pass # proceed; generate_debate_turn() will fail safely if offline - total_components = len(components_to_enrich) - if total_components == 0: - return all_threats + from threat_analysis.core.debate_engine import RedBlueDebateEngine # noqa: PLC0415 + bom_dir = resolve_bom_directory(threat_model) - if progress_callback: - progress_callback(f"Starting AI enrichment for {total_components} components...") - - ai_threats: List[Dict] = [] - processed_count = [0] - semaphore = asyncio.Semaphore(self._enrich_max_concurrent) - - def _build_threat_dict(threat: Dict, component: Dict) -> Dict: - """Convert a raw LLM threat dict into a normalized threat record.""" - stride_category = threat.get("category", "InformationDisclosure") - target_name = component.get("name") - severity_info = self.severity_calculator.get_severity_info( - stride_category, - target_name, - impact=threat.get("business_impact", {}).get("impact_score") if isinstance(threat.get("business_impact"), dict) else None, - likelihood=threat.get("business_impact", {}).get("likelihood_score") if isinstance(threat.get("business_impact"), dict) else None, - ) - raw_capecs = [ - c for c in threat.get("capec_ids", []) - if isinstance(c, str) and c.upper().startswith("CAPEC-") - ] - mapping = self.mitre_mapping.map_threat_to_mitre({ - "stride_category": stride_category, - "capec_ids": raw_capecs, - "description": threat.get("description", ""), - }) - return { - "type": stride_category, - "description": threat.get("description"), - "target": target_name, - "severity": severity_info, - "mitre_techniques": mapping.get("techniques", []), - "stride_category": stride_category, - "capecs": mapping.get("capecs", []), - "cve": [], # CVEs come exclusively from CVEService, never from LLM output - "business_value": component.get("business_value"), - "confidence": threat.get("confidence", 0.8), - "source": "AI", - } + engine = RedBlueDebateEngine(self.ai_provider, config=self._debate_config, bom_directory=bom_dir) + try: + results = await engine.run(scenarios) + except Exception as exc: + logging.warning("Red/Blue debate failed (non-fatal): %s", exc) + return [] - def _flush_batch_results(results: Dict[str, List[Dict]], batch: List[Dict]) -> None: - """Distribute batch results into ai_threats and update progress.""" - comp_by_name = {c["name"]: c for c in batch} - for comp_name, threats_json in results.items(): - component = comp_by_name.get(comp_name) - if component is None: - continue - for threat in threats_json or []: - if not isinstance(threat, dict): - continue - try: - ai_threats.append(_build_threat_dict(threat, component)) - except Exception as exc: - logging.warning("Skipping malformed batch threat for %s: %s", comp_name, exc) - # Update progress for all components in this batch - processed_count[0] += len(batch) - if progress_callback: - names = ", ".join(c["name"] for c in batch) - progress_callback( - f"AI Enrichment: {processed_count[0]}/{total_components} components processed ({names})" - ) + logging.info("Red/Blue debate: produced %d results", len(results)) + return results - async def process_batch(batch: List[Dict]) -> None: - async with semaphore: - try: - results = await self.ai_provider.generate_threats_batch(batch, self.ai_context) - _flush_batch_results(results, batch) - except Exception as e: - logging.error("Batch AI enrichment failed (%s): %s — falling back to individual calls", [c["name"] for c in batch], e) - # Fall back: call generate_threats per component sequentially - for component in batch: - try: - threats_json = await self.ai_provider.generate_threats(component, self.ai_context) - _flush_batch_results({component["name"]: threats_json}, [component]) - except Exception as exc: - logging.error("Error enriching component %s: %s", component.get("name"), exc) - processed_count[0] += 1 - - async def process_component_individual(component: Dict) -> None: - async with semaphore: - try: - generated_threats = await self.ai_provider.generate_threats(component, self.ai_context) - processed_count[0] += 1 - if progress_callback: - progress_callback(f"AI Enrichment: {processed_count[0]}/{total_components} components processed ({component.get('name')})") - for threat in generated_threats: - if not isinstance(threat, dict): - continue - try: - ai_threats.append(_build_threat_dict(threat, component)) - except Exception as exc: - logging.warning("Skipping malformed threat for %s: %s", component.get("name"), exc) - except Exception as e: - logging.error("Error enriching component %s: %s", component.get("name"), e) - processed_count[0] += 1 - - use_batch = ( - self._enrich_batch_size > 1 - and hasattr(self.ai_provider, "generate_threats_batch") - ) - if use_batch: - batch_size = self._enrich_batch_size - batches = [ - components_to_enrich[i: i + batch_size] - for i in range(0, len(components_to_enrich), batch_size) - ] - logging.info( - "AI report enrichment (batch): %d components → %d batch(es) of up to %d", - total_components, len(batches), batch_size, - ) - await asyncio.gather(*[process_batch(b) for b in batches]) - else: - logging.info("AI report enrichment (individual): %d components", total_components) - tasks = [process_component_individual(c) for c in components_to_enrich] - await asyncio.gather(*tasks) - - # Simple deduplication: identify unique AI threats based on category and description - existing_threats_signatures = set() - for threat in all_threats: - signature = (threat.get("stride_category"), threat.get("description")) - existing_threats_signatures.add(signature) - - unique_ai_threats = [] - for ai_threat in ai_threats: - signature = (ai_threat.get("stride_category"), ai_threat.get("description")) - if signature not in existing_threats_signatures: - unique_ai_threats.append(ai_threat) - existing_threats_signatures.add(signature) - - return all_threats + unique_ai_threats + async def _get_ai_service(self): + """Lazily constructs and initializes an AIService for the full AI-enrichment + pipeline (per-component AIThreatCache, boundary-trust-aware prompts, SOC + analysis pass, system-level RAG threats) — see server/ai_service.py. + + Cached on self for the ReportGenerator's lifetime so repeated calls (e.g. the + per-submodel loop in project mode) don't re-run the connection check / RAG + pre-warm every time. Returns None (never raises) if AIService can't be built — + mirrors every other AI degrade-silently path in this codebase. + """ + if self._ai_service is not None: + return self._ai_service + if not self._ai_config_path: + return None + try: + ai_service_module = importlib.import_module("threat_analysis.server.ai_service") + AIService = ai_service_module.AIService + service = AIService(config_path=str(self._ai_config_path), force_disable_rag=self._disable_rag) + await service.init_ai() + self._ai_service = service + return service + except Exception as exc: + logging.warning("AIService init failed for report enrichment (non-fatal): %s", exc) + return None + + async def _run_ai_enrichment(self, threat_model: ThreatModel, progress_callback=None) -> None: + """Runs the full AI enrichment pipeline via AIService, mutating threat_model in + place (appends to each element's `.threats` and to `tm.global_threats_llm`). + + Must be called BEFORE _get_all_threats_with_mitre_info() — that method's + existing collectors read exactly these two mutation points to build the + AI/LLM-sourced entries in the final threat list; this replaces the previous, + weaker duplicate enrichment (no cache, no SOC analysis) that used to run after. + """ + if not self.ai_provider: + logging.warning("AI enrichment skipped: No AI provider initialized.") + return + service = await self._get_ai_service() + if service is None: + return + if progress_callback: + progress_callback("Starting AI enrichment...") + try: + await service._enrich_with_ai_threats(threat_model) + except Exception as exc: + logging.warning("AI enrichment failed (non-fatal): %s", exc) def generate_html_report(self, threat_model, grouped_threats: Dict[str, List], output_file: Path = Path("stride_mitre_report.html"), @@ -621,14 +660,17 @@ def generate_html_report(self, threat_model, grouped_threats: Dict[str, List], total_threats_analyzed = threat_model.mitre_analysis_results.get('total_threats', 0) total_mitre_techniques_mapped = threat_model.mitre_analysis_results.get('mitre_techniques_count', 0) + # AI enrichment must run BEFORE _get_all_threats_with_mitre_info(): it mutates + # threat_model in place (element.threats, tm.global_threats_llm), and that + # method's own collectors are what turn those mutations into AI/LLM entries + # in all_detailed_threats below. + if self.ai_provider: + asyncio.run(self._run_ai_enrichment(threat_model, progress_callback=progress_callback)) + if all_detailed_threats is None: all_detailed_threats = self._get_all_threats_with_mitre_info(grouped_threats, threat_model) - - if self.ai_provider: - original_count = len(all_detailed_threats) - all_detailed_threats = asyncio.run(self._enrich_threats_with_ai(threat_model, all_detailed_threats, progress_callback=progress_callback)) - ai_added = len(all_detailed_threats) - original_count - logging.info(f"AI enrichment complete. Added {ai_added} new threats.") + _ai_count = sum(1 for t in all_detailed_threats if t.get('source') in ('AI', 'LLM')) + logging.info(f"AI enrichment complete. {_ai_count} AI/LLM threat(s) included.") # Recompute STRIDE distribution from the full threat list (pytm + AI + LLM) stride_distribution: Dict[str, int] = {} @@ -637,6 +679,13 @@ def generate_html_report(self, threat_model, grouped_threats: Dict[str, List], if cat in self._VALID_STRIDE: stride_distribution[cat] = stride_distribution.get(cat, 0) + 1 + # Assign stable T-0001-style ids now (same scheme as ReportSerializer) so + # downstream consumers built from this same list — CISO triage, threat graph — + # cite ids that actually match the final JSON/HTML report, instead of CISO + # triage inventing its own reference format because no id existed yet. + for _i, _t in enumerate(all_detailed_threats): + _t.setdefault("id", f"T-{_i + 1:04d}") + self.all_detailed_threats = all_detailed_threats # Cache the final enriched threat list on the model so generate_global_project_report # can include AI-enriched threats without re-running the enrichment pipeline. @@ -657,20 +706,57 @@ def generate_html_report(self, threat_model, grouped_threats: Dict[str, List], ) completeness = _score_model_completeness(threat_model) - threat_graph = self._build_threat_graph_data(threat_model, all_detailed_threats) attack_id_validation = AttackIdValidator().validate_all(all_detailed_threats) - # CISO triage pass — runs after full ranked threat list is available + threat_model._completeness = completeness + threat_model._attack_id_validation = attack_id_validation + + # Red/Blue adversarial debate — mutates gdaf_scenarios in place (score/risk_level only). + # Runs BEFORE CISO triage and the threat graph so both can cite debate-adjusted + # GDAF outcomes instead of pre-debate scores. + debate_results = [] + if self._debate_config.get("enabled") and self.ai_provider and getattr(threat_model, "gdaf_scenarios", None): + try: + debate_results = asyncio.run(self._run_debate(threat_model)) + except Exception as exc: + logging.warning("Red/Blue debate failed (non-fatal): %s", exc) + threat_model.debate_results = debate_results + + # Re-write the .afb Attack Flow files if the debate changed any scenario score. + # run_gdaf_engine() (called earlier in the pipeline, before this method) already + # wrote them once from the pre-debate scenarios — without this, the .afb files + # would permanently disagree with the debate-adjusted scores/risk_levels shown + # in this same HTML report. + if debate_results: + try: + from threat_analysis.generation.attack_flow_builder import AttackFlowBuilder # noqa: PLC0415 + AttackFlowBuilder( + threat_model.gdaf_scenarios, model_name=str(threat_model.tm.name) + ).generate_and_save(str(Path(output_file).parent)) + except Exception as exc: + logging.warning("Failed to re-write .afb files after debate (non-fatal): %s", exc) + + # Built after the debate so GDAF path overlays reflect debate-adjusted scores. + threat_graph = self._build_threat_graph_data( + threat_model, all_detailed_threats, + gdaf_scenarios=getattr(threat_model, "gdaf_scenarios", None), + debate_results=debate_results, + ) + + # CISO triage pass — runs after the full ranked threat list AND the debate are + # available, so the briefing can reference GDAF attack scenarios and debate outcomes. ciso_triage = {} if self.ai_provider and all_detailed_threats: try: - ciso_triage = asyncio.run(self._run_ciso_triage(all_detailed_threats)) + ciso_triage = asyncio.run(self._run_ciso_triage( + all_detailed_threats, + gdaf_scenarios=getattr(threat_model, "gdaf_scenarios", None), + debate_results=debate_results, + )) except Exception as exc: logging.warning("CISO triage failed: %s", exc) # Cache on the model so generate_json_export can include it without re-running. threat_model._ciso_triage = ciso_triage if ciso_triage else None - threat_model._completeness = completeness - threat_model._attack_id_validation = attack_id_validation # Build a serialised summary of GDAF scenarios for the HTML template. # Uses getattr for safety — works even if gdaf_scenarios was never populated. @@ -710,6 +796,10 @@ def generate_html_report(self, threat_model, grouped_threats: Dict[str, List], "target_asset": getattr(scenario, 'target_asset', ''), "path": " → ".join(h["node"] for h in hops_summary), "score": round(float(getattr(scenario, 'path_score', 0)), 2), + "pre_debate_score": ( + round(float(getattr(scenario, 'path_score_pre_debate')), 2) + if getattr(scenario, 'path_score_pre_debate', None) is not None else None + ), "risk_level": getattr(scenario, 'risk_level', 'LOW'), "hop_count": len(hops_summary), "hops": hops_summary, @@ -717,6 +807,78 @@ def generate_html_report(self, threat_model, grouped_threats: Dict[str, List], "unacceptable_risk": bool(getattr(scenario, 'unacceptable_risk', False)), }) + # Build a serialised summary of debate results for the HTML template. + debate_data = [] + for result in debate_results: + rounds_data = [] + for turn in result.rounds: + rounds_data.append({ + "role": turn.role, + "round_index": turn.round_index, + "viability_score": turn.viability_score, + "techniques_attempted": turn.techniques_attempted, + "techniques_blocked": turn.techniques_blocked, + "failed_alternatives": turn.failed_alternatives, + "rationale": turn.rationale, + "detection_gaps": [ + { + "step": g.step, "control_family": g.control_family, + "covered": g.covered, "detail": g.detail, "confidence": g.confidence, + } + for g in turn.detection_gaps + ], + "evidence": [ + { + "claim": e.claim, "evidence_type": e.evidence_type, + "evidence_ref": e.evidence_ref, "confidence": e.confidence, + "verified": e.verified, + } + for e in turn.evidence + ], + }) + debate_data.append({ + "scenario_id": result.scenario_id, + "objective_name": result.objective_name, + "entry_point": result.entry_point, + "target_asset": result.target_asset, + "blocked_paths": result.blocked_paths, + "residual_detection_gaps": [ + { + "step": g.step, "control_family": g.control_family, + "covered": g.covered, "detail": g.detail, "confidence": g.confidence, + } + for g in result.residual_detection_gaps + ], + "red_failed_attempts": result.red_failed_attempts, + "round_count": result.round_count, + "convergence_delta": result.convergence_delta, + "converged": result.converged, + "final_viability": result.final_viability, + "residual_path_viable": result.residual_path_viable, + "debate_factor": result.debate_factor, + "rounds": rounds_data, + }) + + # Automatically-discovered attack paths — best path per STRIDE category through + # the threats' own MITRE techniques, found independently of GDAF (no analyst- + # defined objectives/actors needed). Complements GDAF's analyst-driven scenarios + # with paths the threat data itself suggests. + discovered_attack_paths = [] + if self._attack_flows_config.get("enabled", True): + try: + discovered_attack_paths = AttackFlowGenerator( + all_detailed_threats, model_name=str(threat_model.tm.name), + allowed_categories=self._attack_flows_config.get("generate_for_categories"), + ).get_paths_summary() + except Exception as exc: + logging.warning("Discovered attack path summary failed (non-fatal): %s", exc) + discovered_attack_paths = [] + if discovered_attack_paths: + try: + asyncio.run(self._generate_path_narratives(discovered_attack_paths)) + except Exception as exc: + logging.warning("Attack path narrative pass failed (non-fatal): %s", exc) + template = self.env.get_template('report_template.html') html = template.render( title="STRIDE & MITRE ATT&CK Report", @@ -733,6 +895,8 @@ def generate_html_report(self, threat_model, grouped_threats: Dict[str, List], implemented_mitigation_ids=self.implemented_mitigations, attack_chains=attack_chains, gdaf_scenarios=gdaf_data, + debate_results=debate_data, + discovered_attack_paths=discovered_attack_paths, ciso_triage=ciso_triage, completeness=completeness, attack_id_validation=attack_id_validation, @@ -754,7 +918,14 @@ def generate_json_export(self, threat_model, grouped_threats: Dict[str, List], self.threat_model_ref = threat_model try: - all_detailed_threats = self._get_all_threats_with_mitre_info(grouped_threats, threat_model) + # Reuse the AI-enriched list cached by generate_html_report() if available — + # otherwise this recomputes pytm-only threats and silently drops AI/LLM threats + # (generate_html_report must run first in the same process for the cache to exist). + cached_threats = getattr(threat_model, "_report_all_detailed_threats", None) + if cached_threats: + all_detailed_threats = cached_threats + else: + all_detailed_threats = self._get_all_threats_with_mitre_info(grouped_threats, threat_model) export_data = ReportSerializer.serialize(threat_model, all_detailed_threats) # Include cached CISO triage if available (set by generate_html_report) @@ -892,7 +1063,13 @@ def generate_stix_export(self, threat_model, grouped_threats: Dict[str, List], """Generates a STIX export of the analysis data""" output_dir.mkdir(parents=True, exist_ok=True) - all_detailed_threats = self._get_all_threats_with_mitre_info(grouped_threats, threat_model) + # Reuse the AI-enriched list cached by generate_html_report() if available — see + # generate_json_export() for why this matters. + cached_threats = getattr(threat_model, "_report_all_detailed_threats", None) + if cached_threats: + all_detailed_threats = cached_threats + else: + all_detailed_threats = self._get_all_threats_with_mitre_info(grouped_threats, threat_model) stix_generator = StixGenerator(threat_model, all_detailed_threats) stix_bundle = stix_generator.generate_stix_bundle() @@ -918,6 +1095,17 @@ def _export_detailed_threats(self, grouped_threats: Dict[str, List], threat_mode def _get_all_threats_with_mitre_info(self, grouped_threats: Dict[str, List], threat_model: ThreatModel) -> List[Dict[str, Any]]: """Gathers detailed information for all threats, including MITRE ATT&CK mapping and severity.""" + # Sync severity multipliers from the actual model's '## Severity Multipliers' DSL + # section (already parsed onto threat_model.severity_multipliers by ModelParser). + # self.severity_calculator is injected at construction time and may be a long-lived + # singleton (server mode) reused across different models — without this, it would + # score every model using whichever multipliers happened to be loaded first (or + # none at all). Always sync (even with an empty dict) so a model with no + # multipliers of its own correctly clears any left over from a previous one. + self.severity_calculator.update_target_multipliers( + getattr(threat_model, "severity_multipliers", {}) or {} + ) + pytm_threat_dicts = [] # Load VEX and BOM once — VEX takes priority for CVE scoring, BOM is fallback @@ -1066,6 +1254,14 @@ def _get_all_threats_with_mitre_info(self, grouped_threats: Dict[str, List], thr # Dataflows are stored as pytm objects directly (not dicts) for df in threat_model.dataflows: enriched_elements.append((df, getattr(df, 'name', ''), None)) + # Boundaries are also an AI-enrichment target (A3, see decisions.md) — omitting + # them here silently dropped every AI threat attached to a boundary from the + # final report, even though AIService._enrich_with_ai_threats() and the SOC + # analysis pass both process boundary elements. + for b_name, b_info in threat_model.boundaries.items(): + boundary_obj = b_info.get('boundary') + if boundary_obj is not None: + enriched_elements.append((boundary_obj, b_name, b_info.get('business_value'))) for element_obj, element_name, business_value in enriched_elements: if element_obj is None: @@ -1145,7 +1341,13 @@ def _get_all_threats_with_mitre_info(self, grouped_threats: Dict[str, List], thr # Process global RAG threats if hasattr(threat_model.tm, 'global_threats_llm'): # Access via threat_model for threat in threat_model.tm.global_threats_llm: - target_name = "Threat Model (Global)" # RAG threats are system-level + # RAG threats are system-level (cross-component) by design, but the LLM + # is explicitly asked for affected_components — use it for a target more + # specific than a generic label when available (also lets severity + # multipliers on those components apply, same as any other threat). + ai_details = getattr(threat, 'ai_details', None) or {} + affected = [c for c in (ai_details.get('affected_components') or []) if c] + target_name = " → ".join(affected) if affected else "Threat Model (Global)" threat_description = getattr(threat, 'description', 'RAG-generated global threat') stride_category = getattr(threat, 'category', 'Generic RAG Threat') threat_source = getattr(threat, 'source', 'LLM') @@ -1352,6 +1554,8 @@ def _build_threat_graph_data( self, threat_model: "ThreatModel", all_threats: List[Dict], + gdaf_scenarios: Optional[List[Any]] = None, + debate_results: Optional[List[Any]] = None, ) -> Dict: """Build a JSON-serialisable graph for the interactive threat visualization. @@ -1360,7 +1564,12 @@ def _build_threat_graph_data( The ``threats_by_node`` dict maps node ids to a compact threat list so the template JS can render a click panel without extra requests. - Returns a dict with keys ``nodes``, ``edges``, ``threats_by_node``. + ``gdaf_scenarios``/``debate_results`` are optional — when provided (post-debate), + a ``gdaf_paths`` list is included so the template JS can let the user highlight a + GDAF attack path (entry actor → hop assets → target) over the graph, showing + whether the Red/Blue debate found that path still viable. + + Returns a dict with keys ``nodes``, ``edges``, ``threats_by_node``, ``gdaf_paths``. An empty dict is returned when the model has no components. """ # --- build per-component threat index --------------------------------- @@ -1458,10 +1667,34 @@ def _extract_boundary(item) -> str: for k, v in threats_by_node.items() } + # --- GDAF attack paths (post-debate, if available) -------------------- + debate_by_scenario = {getattr(r, "scenario_id", None): r for r in (debate_results or [])} + gdaf_paths: List[Dict] = [] + for scenario in (gdaf_scenarios or [])[:10]: + entry = getattr(scenario, "entry_point", "") or "" + hop_names = [getattr(h, "asset_name", "") for h in getattr(scenario, "hops", [])] + node_sequence = [n for n in ([entry] + hop_names) if n] + # dedupe consecutive repeats (entry_point sometimes equals the first hop) + node_sequence = [n for i, n in enumerate(node_sequence) if i == 0 or n != node_sequence[i - 1]] + if len(node_sequence) < 2: + continue + debate_result = debate_by_scenario.get(getattr(scenario, "scenario_id", None)) + gdaf_paths.append({ + "scenario_id": getattr(scenario, "scenario_id", ""), + "objective_name": getattr(scenario, "objective_name", ""), + "actor_name": getattr(scenario, "actor_name", ""), + "risk_level": getattr(scenario, "risk_level", "LOW"), + "path_score": round(float(getattr(scenario, "path_score", 0) or 0), 2), + "debated": debate_result is not None, + "residual_path_viable": bool(getattr(debate_result, "residual_path_viable", True)) if debate_result is not None else None, + "nodes": node_sequence, + }) + return { "nodes": nodes, "edges": edges, "threats_by_node": compact_threats, + "gdaf_paths": gdaf_paths, } def _get_all_business_values(self, threat_model: ThreatModel) -> List[str]: @@ -1829,7 +2062,7 @@ def _recursively_generate_reports(self, model_path: Path, project_path: Path, ou if progress_callback: progress_callback(f"Generating STIX report: {model_name}...") try: stix_output_file = output_dir / f"{model_name}_stix_report.json" - all_detailed_threats = threat_model.get_all_threats_details() + all_detailed_threats = get_enriched_threats(threat_model) stix_generator_instance = StixGenerator( threat_model=threat_model, all_detailed_threats=all_detailed_threats @@ -1844,7 +2077,7 @@ def _recursively_generate_reports(self, model_path: Path, project_path: Path, ou if progress_callback: progress_callback(f"Generating ATT&CK Navigator: {model_name}...") try: navigator_output_file = output_dir / f"{model_name}_attack_navigator_layer.json" - all_detailed_threats = threat_model.get_all_threats_details() + all_detailed_threats = get_enriched_threats(threat_model) navigator_generator = AttackNavigatorGenerator( threat_model_name=threat_model.tm.name, all_detailed_threats=all_detailed_threats @@ -1855,7 +2088,7 @@ def _recursively_generate_reports(self, model_path: Path, project_path: Path, ou logging.error(f"❌ Failed to generate ATT&CK Navigator layer for {model_name}: {e}") try: - all_detailed_threats = threat_model.get_all_threats_details() + all_detailed_threats = get_enriched_threats(threat_model) attack_flow_gen = AttackFlowGenerator( threats=all_detailed_threats, model_name=threat_model.tm.name, diff --git a/threat_analysis/generation/svg_generator.py b/threat_analysis/generation/svg_generator.py index a1f9a3d..afe83c3 100644 --- a/threat_analysis/generation/svg_generator.py +++ b/threat_analysis/generation/svg_generator.py @@ -254,7 +254,9 @@ def _generate_node_svg(self, node: Dict) -> List[str]: pos_str = node.get('pos', '0,0') try: node_x, node_y = map(float, pos_str.split(',')) - except: + except (ValueError, TypeError) as exc: + logging.warning("Could not parse node position '%s' for node '%s' (%s) — defaulting to (0, 0).", + pos_str, node.get('name', '?'), exc) node_x, node_y = 0, 0 # Get node width (in inches, need to convert to points: 1 inch = 72 points) diff --git a/threat_analysis/generation/utils.py b/threat_analysis/generation/utils.py index ece79a0..ac59296 100644 --- a/threat_analysis/generation/utils.py +++ b/threat_analysis/generation/utils.py @@ -15,7 +15,22 @@ """ Utility functions for threat analysis generation modules """ -from typing import Any, Union +from typing import Any, List, Union + + +def get_enriched_threats(threat_model: Any) -> List[dict]: + """Returns the fullest available threat list for a model. + + Prefers ``threat_model._report_all_detailed_threats`` — the AI-enriched, debate-aware + list cached by ``ReportGenerator.generate_html_report()`` — over + ``threat_model.get_all_threats_details()``, which only knows about pytm/custom + threats. Callers that run after ``generate_html_report()`` (JSON/STIX/Navigator/ + Attack Flow exports) get AI/LLM-sourced threats this way instead of silently + dropping them; callers that run standalone (no HTML report generated) fall back to + the pytm-only list, unchanged from prior behaviour. + """ + cached = getattr(threat_model, "_report_all_detailed_threats", None) + return cached if cached else threat_model.get_all_threats_details() def extract_name_from_object(obj: Any) -> str: diff --git a/threat_analysis/server/ai_service.py b/threat_analysis/server/ai_service.py index 67d6139..22eb3d1 100644 --- a/threat_analysis/server/ai_service.py +++ b/threat_analysis/server/ai_service.py @@ -63,18 +63,34 @@ class AIService: _sync_loop: Optional[asyncio.AbstractEventLoop] = None _sync_loop_lock = threading.Lock() - def __init__(self, config_path: str, ai_status_event_queue: Optional[queue.Queue] = None): + def __init__(self, config_path: str, ai_status_event_queue: Optional[queue.Queue] = None, + force_disable_rag: bool = False): self.provider: Optional[BaseLLMProvider] = None self.rag_generator = None self.ai_online = False self.ai_config = self._load_ai_config(config_path) self.ai_status_event_queue = ai_status_event_queue + # CLI --no-rag override — skips the RAG pre-warm (chromadb+embeddings, can take + # minutes cold) regardless of rag.enabled in ai_config.yaml. Component-level AI + # enrichment is unaffected — only system-level RAG threats are skipped. + self._force_disable_rag = force_disable_rag self.rate_limit_sleep: float = self.ai_config.get( "threat_generation", {} ).get("rate_limit_sleep", 0.0) self.max_concurrent: int = self.ai_config.get( "threat_generation", {} ).get("max_concurrent_ai_requests", 1) + # 0 = no limit + self.max_threats_per_component: int = int(self.ai_config.get( + "threat_generation", {} + ).get("max_threats_per_component", 0)) + self.confidence_threshold: float = float(self.ai_config.get( + "threat_generation", {} + ).get("confidence_threshold", 0.0)) + # 0 = no threshold (always run RAG when otherwise enabled). + self.rag_min_components: int = int(self.ai_config.get( + "rag", {} + ).get("min_components", 0)) _st = self.ai_config.get("threat_generation", {}).get("sync_timeouts", {}) self._rag_prewarm_timeout: float = float(_st.get("rag_prewarm", 300)) self._rag_sync_timeout: float = float(_st.get("rag_sync", 120)) @@ -279,7 +295,9 @@ async def init_ai(self): # of litellm being available (~64s cold import). By the time check_connection() returns, # RAG components will already be loaded — net savings ≈ 26s on first startup. # run_in_executor returns an asyncio.Future that can be awaited without blocking the loop. - rag_enabled = self.ai_config.get("rag", {}).get("enabled", False) + rag_enabled = (not self._force_disable_rag) and self.ai_config.get("rag", {}).get("enabled", False) + if self._force_disable_rag: + logging.info("RAG disabled by --no-rag.") rag_task: Optional[asyncio.Future] = None if rag_enabled: logging.info("RAG pre-warm started (parallel to AI connection check).") @@ -404,6 +422,15 @@ def generate_markdown_from_prompt_sync(self, prompt: str, markdown: Optional[str return iter(["Error: AI generation failed or timed out."]) return iter(chunks) + @staticmethod + def _count_components(threat_model) -> int: + """Counts actors + servers across the main model and any sub-models — + matches what _generate_rag_threats actually sends as RAG context.""" + count = len(threat_model.actors) + len(threat_model.servers) + for sub in getattr(threat_model, "sub_models", []): + count += len(sub.actors) + len(sub.servers) + return count + async def _generate_rag_threats(self, threat_model) -> List[ExtendedThreat]: # Return List[ExtendedThreat] """ Generates system-level threats using the RAG service. @@ -414,12 +441,20 @@ async def _generate_rag_threats(self, threat_model) -> List[ExtendedThreat]: # R if not self.rag_generator: return [] + component_count = self._count_components(threat_model) + if self.rag_min_components > 0 and component_count < self.rag_min_components: + logging.info( + "Skipping system-level RAG (%d components < rag.min_components=%d).", + component_count, self.rag_min_components, + ) + return [] + logging.debug("Generating system-level threats using RAG...") def _model_markdown(tm, label: str = "") -> str: """Render a single ThreatModel as Markdown for RAG context.""" header = label or tm.tm.name - md = f"# Threat Model: {header}\n\n" + md = f"# System Model: {header}\n\n" md += f"## Description\n\n{tm.tm.description}\n\n" md += "## Components\n\n" # Actors and servers may be stored as dicts (with 'object' key) or as objects @@ -451,7 +486,8 @@ def _to_obj(item): except Exception as exc: logging.warning("Could not append sub-model to RAG context: %s", exc) - rag_generated_threats_json = self.rag_generator.generate_threats(tm_markdown_content) + _top_k = int(self.ai_config.get("rag", {}).get("top_k_examples", 5)) + rag_generated_threats_json = self.rag_generator.generate_threats(tm_markdown_content, k=_top_k) pytm_rag_threats = [] severity_map = {"critical": 5, "high": 4, "medium": 3, "low": 2, "info": 1} @@ -498,9 +534,15 @@ async def _enrich_with_ai_threats(self, threat_model, ai_status_event_queue: Opt — falls back to individual calls when provider lacks generate_threats_batch 5. SOC analysis pass """ - # 1. System-level RAG threats + # 1. System-level RAG threats — isolated in its own try/except so a RAG-side bug + # (e.g. a malformed prompt template) degrades to "no RAG threats" instead of + # aborting the per-component enrichment below, which is the main value driver. if self.ai_online and self.rag_generator: - system_rag_threats = await self._generate_rag_threats(threat_model) + try: + system_rag_threats = await self._generate_rag_threats(threat_model) + except Exception as exc: + logging.warning("System-level RAG threat generation failed (non-fatal): %s", exc) + system_rag_threats = [] if not hasattr(threat_model.tm, 'global_threats_llm'): threat_model.tm.global_threats_llm = [] threat_model.tm.global_threats_llm.extend(system_rag_threats) @@ -726,10 +768,25 @@ def _apply_threats_to_element(element, ai_threats_json: List[Dict]) -> None: """Convert a list of threat dicts to ExtendedThreat objects on element.""" severity_map = {"critical": 5, "high": 4, "medium": 3, "low": 2, "info": 1} likelihood_map = {"high": 5, "medium": 3, "low": 1} - for threat_json in (ai_threats_json or []): - if not isinstance(threat_json, dict): - continue - threat_desc = f"(AI) {threat_json.get('title', 'N/A')}: {threat_json.get('description', '')}" + + candidates = [t for t in (ai_threats_json or []) if isinstance(t, dict)] + if self.confidence_threshold > 0: + candidates = [ + t for t in candidates + if float(t.get("confidence", 1.0) or 1.0) >= self.confidence_threshold + ] + if self.max_threats_per_component > 0: + candidates = sorted( + candidates, key=lambda t: float(t.get("confidence", 0.0) or 0.0), reverse=True + )[: self.max_threats_per_component] + + for threat_json in candidates: + # Computed once so the ExtendedThreat and the log line below always agree — + # the LLM occasionally drops 'title' under batch token pressure, and + # logging the raw (missing) field previously printed "Added AI threat + # 'None' to ..." even though the threat itself got a sensible default. + title = threat_json.get('title') or 'Unknown AI Threat' + threat_desc = f"(AI) {title}: {threat_json.get('description', '')}" business_impact = threat_json.get('business_impact', {}) sev_raw = business_impact.get('severity', 'medium') if isinstance(business_impact, dict) else 'medium' severity = severity_map.get(str(sev_raw).lower(), 3) @@ -737,7 +794,7 @@ def _apply_threats_to_element(element, ai_threats_json: List[Dict]) -> None: str(threat_json.get('likelihood', 'medium')).lower(), 3 ) new_threat = ExtendedThreat( - SID=threat_json.get('title', 'Unknown AI Threat'), + SID=title, description=threat_desc, category=threat_json.get('category', 'Unknown'), likelihood=likelihood, @@ -750,7 +807,7 @@ def _apply_threats_to_element(element, ai_threats_json: List[Dict]) -> None: ] new_threat.ai_details = threat_json element.threats.append(new_threat) - logging.info("Added AI threat '%s' to %s", threat_json.get('title'), element.name) + logging.info("Added AI threat '%s' to %s", title, element.name) async def _update_progress(name: str) -> None: nonlocal processed_elements diff --git a/threat_analysis/server/export_service.py b/threat_analysis/server/export_service.py index a99411f..4d32e6e 100644 --- a/threat_analysis/server/export_service.py +++ b/threat_analysis/server/export_service.py @@ -30,6 +30,7 @@ from threat_analysis.generation.attack_navigator_generator import AttackNavigatorGenerator from threat_analysis.generation.stix_generator import StixGenerator from threat_analysis.generation.attack_flow_generator import AttackFlowGenerator +from threat_analysis.generation.utils import get_enriched_threats from threat_analysis.utils import resolve_gdaf_context, resolve_bom_directory TIMESTAMP_FORMAT = "%Y-%m-%d_%H-%M-%S" @@ -98,7 +99,7 @@ def export_files_logic(self, markdown_content: str, export_format: str, self.report_generator.generate_html_report(threat_model, grouped_threats, output_path) return str(output_path), output_filename elif export_format == "markdown": - output_filename = "threat_model.md" + output_filename = "system_model.md" output_path = output_dir / output_filename with open(output_path, "w", encoding="utf-8") as f: f.write(markdown_content) @@ -170,7 +171,7 @@ def generate_full_project_export(self, markdown_content: str, export_path: Path, if errors := validator.validate(): raise ValueError("Validation failed: " + ", ".join(errors)) - (export_path / "threat_model.md").write_text(markdown_content, encoding="utf-8") + (export_path / "system_model.md").write_text(markdown_content, encoding="utf-8") # Process threats first so severity_map is available for the diagram HTML def single_file_progress_cb(message, is_new_model=False): @@ -182,14 +183,36 @@ def single_file_progress_cb(message, is_new_model=False): dot_code = self.diagram_generator._generate_manual_dot(threat_model) svg_filepath = export_path / "tm_diagram.svg" self.diagram_generator.generate_diagram_from_dot(dot_code, str(svg_filepath), format="svg") + + # GDAF: Goal-Driven Attack Flow (objective-based, requires context with attack_objectives) + # Run BEFORE the HTML report so that gdaf_scenarios are available in the report. + try: + from threat_analysis.utils import run_gdaf_engine + _scenarios = run_gdaf_engine(threat_model, export_path=export_path) + if _scenarios: + logging.info("GDAF: generated %d attack scenarios in %s/gdaf", len(_scenarios), export_path) + except Exception as e: + logging.warning("GDAF generation skipped (non-fatal): %s", e) + + # HTML report generated before JSON/STIX/Navigator/Attack Flow below: it runs AI + # enrichment, the Red/Blue debate, and caches the resulting threat list on + # threat_model._report_all_detailed_threats — get_enriched_threats() picks that + # cache up so these exports include AI/LLM threats instead of only pytm's. + report_path = export_path / "stride_mitre_report.html" + self.report_generator.generate_html_report(threat_model, grouped_threats, report_path, progress_callback=single_file_progress_cb) + + # Diagram severity heat map (B2) built after the HTML report so it reflects + # AI-enriched threats too, not just pytm's — same reasoning as the AI-threat + # reordering above. generate_html_report() attaches AI threats to + # element.threats, which _compute_severity_map() also reads. graph_metadata = self.diagram_service._extract_graph_metadata_for_frontend(threat_model) severity_map = self.report_generator._compute_severity_map(threat_model) - html_diagram_path = export_path / "tm_diagram.html" self.diagram_generator._generate_html_with_legend( svg_filepath, html_diagram_path, threat_model, graph_metadata, severity_map, report_url="stride_mitre_report.html", ) + json_report_path = export_path / "mitre_analysis.json" self.report_generator.generate_json_export(threat_model, grouped_threats, json_report_path) @@ -200,7 +223,7 @@ def single_file_progress_cb(message, is_new_model=False): except Exception as e: logging.warning(f"Could not generate remediation checklist: {e}") - all_detailed_threats = threat_model.get_all_threats_details() + all_detailed_threats = get_enriched_threats(threat_model) navigator_generator = AttackNavigatorGenerator(threat_model_name=str(threat_model.tm.name), all_detailed_threats=all_detailed_threats) navigator_filename = JSON_NAVIGATOR_FILENAME_TPL.format(timestamp=timestamp) navigator_generator.save_layer_to_file(str(export_path / navigator_filename)) @@ -220,20 +243,6 @@ def single_file_progress_cb(message, is_new_model=False): except Exception as e: logging.error("Failed to generate Attack Flow: %s", e) - # GDAF: Goal-Driven Attack Flow (objective-based, requires context with attack_objectives) - # Run BEFORE the HTML report so that gdaf_scenarios are available in the report. - try: - from threat_analysis.utils import run_gdaf_engine - _scenarios = run_gdaf_engine(threat_model, export_path=export_path) - if _scenarios: - logging.info("GDAF: generated %d attack scenarios in %s/gdaf", len(_scenarios), export_path) - except Exception as e: - logging.warning("GDAF generation skipped (non-fatal): %s", e) - - # HTML report generated last so it includes GDAF scenarios - report_path = export_path / "stride_mitre_report.html" - self.report_generator.generate_html_report(threat_model, grouped_threats, report_path, progress_callback=single_file_progress_cb) - result["reports"] = { "html": "stride_mitre_report.html", "json": "mitre_analysis.json", @@ -274,7 +283,7 @@ def export_all_files_logic(self, markdown_content: str, submodels: list | None = metadata = { "version": "1.0", "version_id": version_id, "last_updated": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), - "model_file": "threat_model.md", "positions": element_positions + "model_file": "system_model.md", "positions": element_positions } (export_path / "element_positions.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8") @@ -315,7 +324,12 @@ def export_navigator_stix_logic(self, markdown_content: str, submodels: list | N if errors := validator.validate(): raise ValueError("Validation failed: " + ", ".join(errors)) - all_detailed_threats = threat_model.get_all_threats_details() + # Intentionally pytm-only, no AI enrichment: this is the fast dedicated + # Navigator+STIX export (no generate_html_report() call, no AI/GDAF/debate) — + # unlike "Export All", which runs the full enrichment pipeline and takes + # noticeably longer. get_enriched_threats() would fall back to + # get_all_threats_details() anyway since this threat_model is freshly built. + all_detailed_threats = get_enriched_threats(threat_model) navigator_generator = AttackNavigatorGenerator(threat_model_name=str(threat_model.tm.name), all_detailed_threats=all_detailed_threats) navigator_filepath = output_dir / JSON_NAVIGATOR_FILENAME_TPL.format(timestamp=timestamp) navigator_generator.save_layer_to_file(str(navigator_filepath)) @@ -350,7 +364,9 @@ def export_attack_flow_logic(self, markdown_content: str, raise RuntimeError("Failed to create or validate threat model") threat_model.process_threats() - all_detailed_threats = threat_model.get_all_threats_details() + # Intentionally pytm-only, no AI enrichment — see export_navigator_stix_logic + # above for why (fast dedicated export, no generate_html_report() call). + all_detailed_threats = get_enriched_threats(threat_model) attack_flow_generator = AttackFlowGenerator(threats=all_detailed_threats, model_name=threat_model.tm.name) attack_flow_generator.generate_and_save_flows(temp_export_dir) diff --git a/threat_analysis/server/server.py b/threat_analysis/server/server.py index 9d0c687..170ca37 100644 --- a/threat_analysis/server/server.py +++ b/threat_analysis/server/server.py @@ -17,6 +17,7 @@ import base64 import logging import re +import secrets import datetime import json import glob @@ -24,7 +25,7 @@ import time from pathlib import Path from typing import Dict, List, Any, Optional -from flask import Flask, render_template, request, jsonify, send_from_directory, send_file, make_response, Response, stream_with_context, g +from flask import Flask, render_template, request, jsonify, send_from_directory, send_file, make_response, Response, stream_with_context, g, abort, session from threat_analysis import config import asyncio from threat_analysis.server.events import ai_status_event_queue @@ -86,6 +87,11 @@ def get_threat_model_service(): template_folder=os.path.join(server_dir, "templates"), static_folder=os.path.join(server_dir, "static"), static_url_path="/static") +# Regenerated on every process start — this session cookie is used only to isolate +# per-browser-session state (active project/model file path, see +# _get_active_model_file_path below), not for authentication. Sessions not +# surviving a server restart is an acceptable trade-off for this local tool. +app.secret_key = secrets.token_hex(32) @app.before_request def before_request(): @@ -151,8 +157,10 @@ async def init_and_log(): try: from threat_analysis.server.events import ai_status_event_queue ai_status_event_queue.put(f"event: ai_status\ndata: {json.dumps(data)}\n\n") - except Exception: - pass + except Exception as exc: + # If this queue is broken, no further AI status updates reach the + # client via this path — worth a log line instead of a silent no-op. + logging.warning("Could not push ai_status to event queue: %s", exc) except Exception as e: logging.error(f"Error during background AI initialization: {e}", exc_info=True) @@ -215,10 +223,30 @@ def generate_config_files(): print("Warning: Could not verify config.js generation") initial_markdown_content = "" +# Startup-only defaults, set once in run_server() and never mutated afterward — safe +# to read as plain globals. The per-session ACTIVE values (which /api/set_project_path +# can override at runtime) live in the Flask session instead — see the two helpers +# below. Without this split, two browser tabs/windows pointed at different projects +# would clobber each other's active-project state (server ran threaded=True). initial_project_path = None initial_model_file_path: Optional[str] = None +graphical_editor_enabled = False + + +def _get_active_model_file_path() -> Optional[str]: + """The model file path for the CURRENT browser session — /api/set_project_path's + override if this session has called it, else the server's startup value (the file + given on the command line). Must be called from within a request context. + """ + return session.get('model_file_path', initial_model_file_path) + -DEFAULT_EMPTY_MARKDOWN = """# Threat Model: New Model +def _get_active_project_path() -> Optional[str]: + """The project directory for the CURRENT browser session — see + _get_active_model_file_path for the same session-first, startup-fallback logic.""" + return session.get('project_path', initial_project_path) + +DEFAULT_EMPTY_MARKDOWN = """# System Model: New Model ## Description A new threat model. Describe your system here. @@ -246,13 +274,23 @@ def generate_config_files(): def get_model_name(markdown_content: str) -> str: - match = re.search(r"^# Threat Model: (.*)$", markdown_content, re.MULTILINE) + match = re.search(r"^# (Threat|System) Model: (.*)$", markdown_content, re.MULTILINE) if match: - return match.group(1).strip() + if match.group(1) == "Threat": + logging.warning( + "Model title uses the deprecated '# Threat Model: ...' tag — use " + "'# System Model: ...' instead (the DSL file describes the system model; " + "the threat model is what SecOpsTM generates from it)." + ) + return match.group(2).strip() return "Untitled Model" -def run_server(model_filepath: Optional[str] = None, project_path: Optional[str] = None): +def run_server( + model_filepath: Optional[str] = None, + project_path: Optional[str] = None, + enable_graphical_editor: bool = False, +): """ This function is the main entry point for the web server. It launches the Flask application on a single port and serves a menu @@ -264,6 +302,8 @@ def run_server(model_filepath: Optional[str] = None, project_path: Optional[str] global initial_markdown_content global initial_project_path global initial_model_file_path + global graphical_editor_enabled + graphical_editor_enabled = enable_graphical_editor effective_model_path = None if project_path and os.path.isdir(project_path): @@ -317,6 +357,18 @@ def run_server(model_filepath: Optional[str] = None, project_path: Optional[str] if _debug and not _host.startswith("127."): logging.warning("FLASK_DEBUG ignored — debug mode is only allowed on loopback addresses.") _debug = False + if not _host.startswith("127.") and _host != "localhost": + # No authentication on any route (single-user tool, see architecture.md). + # FLASK_HOST=0.0.0.0 is what the Docker image sets by default so `docker run + # -p 5000:5000` works at all — but it also means anyone who can reach this + # host and port has the exact same access as the person running it. + logging.warning( + "⚠️ Server is binding to %s — reachable from OTHER machines, not just this one. " + "There is NO authentication on any route. Do not expose this port to an " + "untrusted network. For local-only access, set FLASK_HOST=127.0.0.1 " + "(or run the Docker image with `-p 127.0.0.1:%d:5000`).", + _host, _port, + ) app.run(debug=_debug, port=_port, host=_host, threaded=True) @@ -337,7 +389,7 @@ def get_data_dictionary(): @app.route("/") def index(): """Serves the main menu.""" - return render_template("index.html") + return render_template("index.html", graphical_editor_enabled=graphical_editor_enabled) @app.route("/simple") def simple_mode(): @@ -347,8 +399,9 @@ def simple_mode(): model_name = get_model_name(initial_markdown_content) # Prepare initial_models for the template - if initial_project_path: - initial_models_data = get_threat_model_service().load_project(initial_project_path) + active_project_path = _get_active_project_path() + if active_project_path: + initial_models_data = get_threat_model_service().load_project(active_project_path) else: initial_models_data = [{ "path": "main.md", @@ -373,6 +426,8 @@ def graphical_editor(): """ Serves the main web interface. """ + if not graphical_editor_enabled: + abort(404) return render_template( "graphical_editor.html", ai_online=get_threat_model_service().ai_online # Pass AI status to template @@ -473,7 +528,8 @@ async def generate_markdown_from_prompt(): else: logging.warning(f"AI response did not contain a clear markdown block. Raw response preview: {full_response[:200]}") # Fallback: find the start of the threat model and clean up common AI chatter. - model_start_index = full_response.find("# Threat Model:") + model_start_match = re.search(r"# (?:Threat|System) Model:", full_response) + model_start_index = model_start_match.start() if model_start_match else -1 if model_start_index != -1: extracted_markdown = full_response[model_start_index:].strip() end_block_index = extracted_markdown.rfind("```") @@ -512,7 +568,7 @@ def update_diagram(): try: result = get_threat_model_service().update_diagram_logic( - markdown_content, submodels=submodels, model_file_path=initial_model_file_path + markdown_content, submodels=submodels, model_file_path=_get_active_model_file_path() ) g.processing_time_ms = result.get("processing_time_ms", 0) g.generation_time_ms = result.get("generation_time_ms", 0) @@ -542,7 +598,7 @@ def _format_properties(item: dict, props_to_include: list) -> str: def convert_json_to_markdown(data: dict) -> str: """Converts JSON from the graphical editor to Markdown DSL.""" - markdown_lines = ["# Threat Model: Graphical Editor"] + markdown_lines = ["# System Model: Graphical Editor"] boundaries = data.get("boundaries", []) actors = data.get("actors", []) @@ -594,6 +650,8 @@ def graphical_update(): """ Receives JSON graph data, converts it to Markdown, and returns the analysis. """ + if not graphical_editor_enabled: + abort(404) logging.debug("Entering graphical_update function.") json_data = request.json if not json_data: @@ -639,7 +697,7 @@ def export_files(): try: output_path, output_filename = get_threat_model_service().export_files_logic( markdown_content=markdown_content, export_format=export_format, - model_file_path=initial_model_file_path, + model_file_path=_get_active_model_file_path(), ) absolute_output_directory = os.path.join(project_root, os.path.dirname(output_path)) @@ -678,7 +736,7 @@ def export_all_files(): submodels = request.json.get("submodels", []) zip_buffer, timestamp = get_threat_model_service().export_all_files_logic( markdown_content=markdown_content, submodels=submodels, - model_file_path=initial_model_file_path, + model_file_path=_get_active_model_file_path(), ) return send_file( zip_buffer, @@ -721,7 +779,7 @@ def export_navigator_stix_files(): try: submodels = request.json.get("submodels", []) zip_buffer, timestamp = get_threat_model_service().export_navigator_stix_logic( - markdown_content, submodels=submodels, model_file_path=initial_model_file_path, + markdown_content, submodels=submodels, model_file_path=_get_active_model_file_path(), ) if not zip_buffer: return jsonify({"error": "Failed to generate navigator and STIX files."}), 500 @@ -784,7 +842,7 @@ def export_attack_flow(): markdown_content = convert_json_to_markdown(json_data) zip_buffer, timestamp = get_threat_model_service().export_attack_flow_logic( - markdown_content, model_file_path=initial_model_file_path, + markdown_content, model_file_path=_get_active_model_file_path(), ) @@ -1020,12 +1078,14 @@ def generate_all(): if not browser_managed_files: def _find_src_root() -> Optional[Path]: candidates = [] - if initial_model_file_path: - p = Path(initial_model_file_path).parent + active_model_file_path = _get_active_model_file_path() + if active_model_file_path: + p = Path(active_model_file_path).parent if p.is_dir(): candidates.append(p) - if initial_project_path: - candidates.append(Path(initial_project_path)) + active_project_path = _get_active_project_path() + if active_project_path: + candidates.append(Path(active_project_path)) for c in candidates: if (c / "main.md").exists() or (c / "model.md").exists(): return c @@ -1286,7 +1346,7 @@ async def export_json(): model_description="JSON export from web interface", cve_service=get_threat_model_service().cve_service, validate=True, - model_file_path=initial_model_file_path, + model_file_path=_get_active_model_file_path(), ) if not threat_model: return jsonify({"error": "Failed to create threat model"}), 400 @@ -1383,9 +1443,12 @@ async def validate_markdown(): @app.route("/api/set_project_path", methods=["POST"]) async def set_project_path(): - """Set the project base path for BOM/context auto-discovery.""" - global initial_model_file_path, initial_project_path + """Set the project base path for BOM/context auto-discovery. + Stored in this browser session only (not the process-wide globals) — otherwise + one tab/window setting a project would silently redirect every other concurrent + session's exports to it too (server runs threaded=True). + """ data = await request.get_json() if not data: return jsonify({"success": False, "error": "Missing request body"}), 400 @@ -1406,18 +1469,18 @@ async def set_project_path(): bom_available = (path_obj / "BOM").is_dir() context_available = (path_obj / "context").is_dir() - # Point initial_model_file_path at a plausible model file within the directory so + # Point model_file_path at a plausible model file within the directory so # that ExportService auto-discovery of BOM/ and context/ is activated. for candidate in ("main.md", "model.md"): candidate_path = path_obj / candidate if candidate_path.exists(): - initial_model_file_path = str(candidate_path.resolve()) + session['model_file_path'] = str(candidate_path.resolve()) break else: # No .md found — use a synthetic path so the parent dir is set correctly - initial_model_file_path = str(path_obj.resolve() / "main.md") + session['model_file_path'] = str(path_obj.resolve() / "main.md") - initial_project_path = resolved + session['project_path'] = resolved logging.info("Project path set to: %s (bom=%s, context=%s)", resolved, bom_available, context_available) return jsonify({ "success": True, diff --git a/threat_analysis/server/static/js/ThreatModelGenerator.js b/threat_analysis/server/static/js/ThreatModelGenerator.js index 2938aeb..278d933 100644 --- a/threat_analysis/server/static/js/ThreatModelGenerator.js +++ b/threat_analysis/server/static/js/ThreatModelGenerator.js @@ -275,7 +275,7 @@ class ThreatModelGenerator { } convertJsonToMarkdown(data) { - const markdown_lines = ["# Threat Model: Graphical Editor"]; + const markdown_lines = ["# System Model: Graphical Editor"]; const boundaries = data.boundaries || []; const actors = data.actors || []; @@ -360,7 +360,7 @@ class ThreatModelGenerator { } getModelName(markdownContent) { - const match = markdownContent.match(/^# Threat Model: (.*)$/m); + const match = markdownContent.match(/^# (?:Threat|System) Model: (.*)$/m); return match ? match[1].trim() : "Untitled Model"; } } \ No newline at end of file diff --git a/threat_analysis/server/static/js/config.js b/threat_analysis/server/static/js/config.js index 472cd72..7ec574c 100644 --- a/threat_analysis/server/static/js/config.js +++ b/threat_analysis/server/static/js/config.js @@ -1,5 +1,5 @@ // Threat Model Configuration - Web UI -// Generated on: 2026-05-18 21:30:14 +// Generated on: 2026-07-12 20:24:38 // This file contains configuration for the threat model web interface const ThreatModelConfig = { diff --git a/threat_analysis/server/templates/graphical_editor.html b/threat_analysis/server/templates/graphical_editor.html index 7afbad3..78a49e4 100644 --- a/threat_analysis/server/templates/graphical_editor.html +++ b/threat_analysis/server/templates/graphical_editor.html @@ -108,7 +108,7 @@

Legend