From 39f0e20e7535ad0cde173bc12b848dd98597b399 Mon Sep 17 00:00:00 2001 From: melonattacker <41631269+melonattacker@users.noreply.github.com> Date: Tue, 21 Apr 2026 10:17:42 +0900 Subject: [PATCH 1/3] feat: Enhance DFD generation from system descriptions --- README.md | 14 +- docs/cli.md | 13 ++ docs/tutorials.md | 18 ++ src/threat_thinker/business_context.py | 256 +++++++++++++++++++++++++ src/threat_thinker/constants.py | 49 +++++ src/threat_thinker/llm/inference.py | 144 ++++++++++++++ src/threat_thinker/main.py | 245 ++++++++++++++++++----- src/threat_thinker/webui.py | 213 +++++++++++++------- tests/test_business_context.py | 80 ++++++++ tests/test_cli_outputs.py | 23 +++ tests/test_llm_context_prompt.py | 86 +++++++++ tests/test_webui_helpers.py | 14 ++ 12 files changed, 1032 insertions(+), 123 deletions(-) create mode 100644 src/threat_thinker/business_context.py create mode 100644 tests/test_business_context.py diff --git a/README.md b/README.md index 5539450..3661175 100644 --- a/README.md +++ b/README.md @@ -14,9 +14,10 @@ AI-powered threat modeling that turns architecture diagrams and business context ## What is Threat Thinker? -Threat Thinker is an open-source tool that turns architecture diagrams and business context into threat models automatically. Provide a DFD or architecture diagram as the system shape, add Business Context for scope and assumptions, and optionally use RAG to bring in supporting standards or internal guidance. +Threat Thinker is an open-source tool that turns system descriptions, architecture diagrams, and business context into threat models automatically. Provide a natural-language system description or a DFD/architecture diagram as the system shape, add Business Context for scope and assumptions, and optionally use RAG to bring in supporting standards or internal guidance. Key Features: +- **Description-to-DFD**: Generates an intermediate Graph IR DFD from a natural-language system description when no diagram is available. - **Diagram coverage**: Ingests Mermaid, draw.io, Threat Dragon JSON, native Graph IR JSON, and images. - **Business Context**: Injects scope, actors, assets, assumptions, and constraints from PDF, Markdown, or text files. - **Attribute inference**: Uses LLMs to enrich components, data flows, and trust boundaries. @@ -26,7 +27,7 @@ Key Features: ## Key Features ### Diagram-to-threat reasoning -- Drop in a diagram via CLI (`--diagram` or format-specific flags) or Web UI and get threats without manual modeling. +- Provide `--description` when you do not have a diagram, or drop in a diagram via CLI (`--diagram` or format-specific flags) or Web UI. - Supports Mermaid, draw.io, Threat Dragon JSON, native Graph IR JSON, and image-based diagrams. - Deterministic parsing plus LLM reasoning fills missing labels, trust boundaries, and protocols. - Outputs prioritized threats with short rationales and OWASP ASVS/CWE references for quick review. @@ -38,6 +39,7 @@ Key Features:

### Business Context as first-class input +- Use `--description` for the system description that can generate a DFD when no diagram is provided. - Use `--context` to add required business context that is not visible in the DFD or architecture diagram. - Include scope, actors, sensitive assets, workflows, regulatory assumptions, availability needs, and audit expectations. - Threat Thinker injects the full extracted text from PDF, Markdown, or text files into the threat prompt. @@ -144,6 +146,14 @@ Here is an example of command using CLI mode. ```bash +# Think: Generate a DFD from a system description, then analyze threats +threat-thinker think \ + --description "Customers use a web app to manage orders. The app stores customer PII in Postgres and sends email through a third-party provider." \ + --topn 5 \ + --llm-api openai \ + --llm-model gpt-4.1 \ + --out-dir reports/ + # Think: Analyze a diagram threat-thinker think \ --diagram examples/diagrams/web/system.mmd \ diff --git a/docs/cli.md b/docs/cli.md index 12990c2..01e6a40 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -10,10 +10,13 @@ threat-thinker version | Flag | Purpose | Notes | | --- | --- | --- | +| `--description ` | Provide a natural-language system description | Used to generate a DFD when no diagram is supplied. Also injected into the threat prompt. Repeat to append multiple blocks. | +| `--description-file ` | Load the system description from a file | Supports PDF, Markdown, and text files via the context loader. Repeat for multiple files. | | `--mermaid / --drawio / --threat-dragon / --ir / --image / --diagram` | Choose input format | Mermaid `.mmd/.mermaid`, Draw.io `.xml`, Threat Dragon v2 `.json`, native Graph IR `.json`, image files, or generic `--diagram` autodetect (recognizes Threat Dragon JSON when version is 2.x). | | `--drawio-page ` | Select Draw.io page to parse | Optional; supports page id, page name, or 0-based index for multi-page `.drawio` files. | | `--infer-hints` | Ask LLM to infer node/edge attributes | Useful when diagrams omit component roles, protocols, or data sensitivity. | | `--context ` | Inject business context into the threat prompt | Repeat for multiple PDF, Markdown, or text files. Unlike RAG, each file's extracted full text is included directly. | +| `--context-file ` | Alias for `--context` | Added for clarity when scripts already use `--description-file`. | | `--prompt-token-limit ` | Fail before analysis if the assembled prompt is too large | Applies to graph, context documents, RAG snippets, and instructions. No truncation is performed. | | `--rag --kb ` | Enable local KB retrieval | Requires a built KB; pairs with `--rag-topk`. | | `--rag-topk ` | Set number of KB chunks to inject | Typical 5–10. | @@ -30,12 +33,22 @@ threat-thinker version | `--out-name ` | Override base filename | Affects `*_report.{json,md,html}` and diff outputs. | Notes: +- If no diagram input is provided, `--description` or `--description-file` is required. Threat Thinker generates an intermediate Graph IR DFD and writes it as `_report_dfd.json` next to the reports. +- If a diagram is provided, `--description` is not used to generate a DFD; it is included as additional threat-analysis context. +- Use `--description` for the system shape. Use `--context` for supplemental business rules, assumptions, policies, or constraints that should inform threat inference. - Ollama backend does not support image inputs; use Mermaid/Draw.io/Threat Dragon files with `--llm-api ollama`. - Native IR JSON is explicit-only in v1; use `--ir` or API/UI `type=ir`, not `--diagram`. - RAG requires OpenAI embeddings; set `OPENAI_API_KEY` when using `--rag`. - Use `--context` for scope, actors, assets, and business assumptions that should always be visible to the LLM. Use `--rag` for optional supporting references retrieved from larger KBs. They can be combined: ```bash +# Description-only analysis +threat-thinker think \ + --description-file examples/diagrams/web/business-context.md \ + --llm-api openai --llm-model gpt-4.1 \ + --out-dir reports/ + +# Diagram plus supplemental context threat-thinker think \ --mermaid examples/diagrams/web/system.mmd \ --context examples/diagrams/web/business-context.md \ diff --git a/docs/tutorials.md b/docs/tutorials.md index 2fe2e72..c7fb60b 100644 --- a/docs/tutorials.md +++ b/docs/tutorials.md @@ -1,4 +1,22 @@ # Tutorials +## Tutorial 0: Start from a system description +Use this flow when you do not have a DFD yet. Threat Thinker generates an intermediate Graph IR DFD from the description, runs threat inference, and writes the normal Markdown/JSON/HTML reports. + +### Command + +```bash +threat-thinker think \ + --description "Customers use a web app to manage orders. The frontend calls an API hosted on AWS. The API stores customer PII and order history in Postgres and sends transactional email through a third-party provider." \ + --topn 5 \ + --llm-api openai \ + --llm-model gpt-4.1 \ + --out-dir reports/ +``` + +The generated DFD is written next to the reports as `description_report_dfd.json`. If the description is too vague, Threat Thinker returns clarifying questions instead of guessing a large speculative graph. + +When you already have a diagram, keep using `--mermaid`, `--drawio`, `--threat-dragon`, `--ir`, `--image`, or `--diagram`. In that mode `--description` is optional extra context for threat inference, and `--context` remains the file-based supplemental context input. + ## Tutorial 1: Analyze simple web application In this example, we will analyze the architecture diagram of a simple web application. written in mermaid and identify potential threats. diff --git a/src/threat_thinker/business_context.py b/src/threat_thinker/business_context.py new file mode 100644 index 0000000..a34c521 --- /dev/null +++ b/src/threat_thinker/business_context.py @@ -0,0 +1,256 @@ +"""Utilities for turning system descriptions into Threat Thinker Graph IR.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional + +from threat_thinker.models import Edge, Graph, ImportMetrics, Node, Zone +from threat_thinker.zone_utils import ( + representative_zone_name, + sort_zone_ids_by_hierarchy, +) + + +CONFIDENCE_VALUES = {"stated", "implied", "assumed"} + + +@dataclass +class BusinessContextDfdResult: + summary: str + graph: Graph + metrics: ImportMetrics + assumptions: List[str] = field(default_factory=list) + clarifying_questions: List[str] = field(default_factory=list) + element_confidence: Dict[str, Dict[str, str]] = field(default_factory=dict) + + +def dfd_result_from_payload(payload: Dict[str, Any]) -> BusinessContextDfdResult: + """Convert a validated LLM DFD payload into Graph IR plus sidecar metadata.""" + graph_payload = payload.get("graph") or {} + graph = Graph(source_format="description") + metrics = ImportMetrics() + element_confidence: Dict[str, Dict[str, str]] = { + "nodes": {}, + "edges": {}, + "zones": {}, + } + + zones_payload = _dict_payload(graph_payload.get("zones")) + graph.zones = _parse_zones(zones_payload, element_confidence["zones"]) + + nodes_payload = _dict_payload(graph_payload.get("nodes")) + metrics.node_label_candidates = len(nodes_payload) + graph.nodes = _parse_nodes(nodes_payload, graph.zones, element_confidence["nodes"]) + metrics.node_labels_parsed = len(graph.nodes) + + edges_payload = graph_payload.get("edges") or [] + metrics.edge_candidates = ( + len(edges_payload) if isinstance(edges_payload, list) else 0 + ) + graph.edges = _parse_edges(edges_payload, graph.nodes, element_confidence["edges"]) + metrics.edges_parsed = len(graph.edges) + + return BusinessContextDfdResult( + summary=str(payload.get("summary") or "").strip(), + graph=graph, + metrics=metrics, + assumptions=_string_list(payload.get("assumptions")), + clarifying_questions=_string_list(payload.get("clarifying_questions")), + element_confidence=element_confidence, + ) + + +def graph_to_native_ir_dict(graph: Graph) -> Dict[str, Any]: + """Serialize Graph to the native IR shape accepted by the IR parser.""" + return { + "nodes": { + node_id: { + "id": node.id, + "label": node.label, + "zone": node.zone, + "zones": node.zones, + "type": node.type, + "data": node.data, + "auth": node.auth, + "notes": node.notes, + } + for node_id, node in graph.nodes.items() + }, + "edges": [ + { + "src": edge.src, + "dst": edge.dst, + "label": edge.label, + "protocol": edge.protocol, + "data": edge.data, + "id": edge.id, + } + for edge in graph.edges + ], + "zones": { + zone_id: { + "id": zone.id, + "name": zone.name, + "parent_id": zone.parent_id, + } + for zone_id, zone in graph.zones.items() + }, + } + + +def dfd_result_to_sidecar_dict(result: BusinessContextDfdResult) -> Dict[str, Any]: + """Return the JSON-serializable sidecar payload for generated DFDs.""" + return { + "summary": result.summary, + "graph": graph_to_native_ir_dict(result.graph), + "assumptions": result.assumptions, + "clarifying_questions": result.clarifying_questions, + "element_confidence": result.element_confidence, + "import_metrics": { + "total_lines": result.metrics.total_lines, + "edge_candidates": result.metrics.edge_candidates, + "edges_parsed": result.metrics.edges_parsed, + "node_label_candidates": result.metrics.node_label_candidates, + "node_labels_parsed": result.metrics.node_labels_parsed, + "import_success_rate": result.metrics.import_success_rate, + }, + } + + +def dfd_result_to_sidecar_json(result: BusinessContextDfdResult) -> str: + return json.dumps(dfd_result_to_sidecar_dict(result), ensure_ascii=False, indent=2) + + +def _parse_zones( + zones_payload: Dict[str, Any], confidence_out: Dict[str, str] +) -> Dict[str, Zone]: + zones: Dict[str, Zone] = {} + for zone_key, zone_value in zones_payload.items(): + if not isinstance(zone_value, dict): + continue + zone_id = str(zone_value.get("id") or zone_key).strip() + zone_name = str(zone_value.get("name") or zone_id).strip() + if not zone_id or not zone_name: + continue + parent_id = _strip_optional_str(zone_value.get("parent_id")) + confidence = _confidence(zone_value.get("confidence")) + confidence_out[zone_id] = confidence + zones[zone_id] = Zone(id=zone_id, name=zone_name, parent_id=parent_id) + + for zone in zones.values(): + if zone.parent_id not in zones: + zone.parent_id = None + return zones + + +def _parse_nodes( + nodes_payload: Dict[str, Any], + zones: Dict[str, Zone], + confidence_out: Dict[str, str], +) -> Dict[str, Node]: + nodes: Dict[str, Node] = {} + for node_key, node_value in nodes_payload.items(): + if not isinstance(node_value, dict): + continue + node_id = str(node_value.get("id") or node_key).strip() + label = str(node_value.get("label") or node_id).strip() + if not node_id or not label: + continue + + zone_ids = _string_list(node_value.get("zones")) + if zones: + zone_ids = [zone_id for zone_id in zone_ids if zone_id in zones] + zone_ids = sort_zone_ids_by_hierarchy(zone_ids, zones) + zone = _strip_optional_str(node_value.get("zone")) + if zones and zone_ids: + zone = representative_zone_name(zone_ids, zones) or zone + + confidence_out[node_id] = _confidence(node_value.get("confidence")) + nodes[node_id] = Node( + id=node_id, + label=label, + zone=zone, + zones=zone_ids, + type=_strip_optional_str(node_value.get("type")), + data=_string_list(node_value.get("data")), + auth=_optional_bool(node_value.get("auth")), + notes=_strip_optional_str(node_value.get("notes")), + ) + return nodes + + +def _parse_edges( + edges_payload: Any, + nodes: Dict[str, Node], + confidence_out: Dict[str, str], +) -> List[Edge]: + if not isinstance(edges_payload, list): + return [] + edges: List[Edge] = [] + for edge_value in edges_payload: + if not isinstance(edge_value, dict): + continue + src = str(edge_value.get("src") or "").strip() + dst = str(edge_value.get("dst") or "").strip() + if not src or not dst or src not in nodes or dst not in nodes: + continue + edge = Edge( + src=src, + dst=dst, + label=_strip_optional_str(edge_value.get("label")), + protocol=_strip_optional_str(edge_value.get("protocol")), + data=_string_list(edge_value.get("data")), + id=_strip_optional_str(edge_value.get("id")), + ) + edges.append(edge) + confidence_out[_edge_key(edge)] = _confidence(edge_value.get("confidence")) + return edges + + +def _dict_payload(value: Any) -> Dict[str, Any]: + if isinstance(value, dict): + return value + if isinstance(value, list): + result: Dict[str, Any] = {} + for item in value: + if isinstance(item, dict): + item_id = str(item.get("id") or item.get("name") or "").strip() + if item_id: + result[item_id] = item + return result + return {} + + +def _string_list(value: Any) -> List[str]: + if value is None: + return [] + if not isinstance(value, list): + return [] + return [str(item).strip() for item in value if str(item).strip()] + + +def _strip_optional_str(value: Any) -> Optional[str]: + if value is None: + return None + text = str(value).strip() + return text or None + + +def _optional_bool(value: Any) -> Optional[bool]: + return value if isinstance(value, bool) else None + + +def _confidence(value: Any) -> str: + text = str(value or "").strip().lower() + if text in CONFIDENCE_VALUES: + return text + return "assumed" + + +def _edge_key(edge: Edge) -> str: + if edge.id: + return edge.id + label = f":{edge.label}" if edge.label else "" + return f"{edge.src}->{edge.dst}{label}" diff --git a/src/threat_thinker/constants.py b/src/threat_thinker/constants.py index 9f15325..5441cb1 100644 --- a/src/threat_thinker/constants.py +++ b/src/threat_thinker/constants.py @@ -51,6 +51,55 @@ "- Return ONLY the JSON object, no other text or formatting.\n" ) +# LLM-driven DFD generation from a system description +DFD_SYSTEM = ( + "You are a senior security architect who builds Data Flow Diagrams (DFDs) " + "from natural-language descriptions of systems. Produce a conservative " + "Threat Thinker native Graph IR that a threat modeler can reason against." +) + +DFD_INSTRUCTIONS = ( + "Return ONLY a valid JSON object (no markdown formatting, no code blocks, no ```json markers).\n\n" + "Required JSON structure:\n" + "{\n" + ' "summary": "One-paragraph restatement of the system.",\n' + ' "graph": {\n' + ' "nodes": {\n' + ' "": {\n' + ' "id": "",\n' + ' "label": "string",\n' + ' "zone": "optional zone name",\n' + ' "zones": ["outer_zone_id","inner_zone_id"],\n' + ' "type": "actor|service|pod|database|s3|elb|ingress|queue|cache|lambda|external|unknown",\n' + ' "data": ["PII","Credentials","Internal","Secrets"],\n' + ' "auth": true,\n' + ' "notes": "optional short rationale",\n' + ' "confidence": "stated|implied|assumed"\n' + " }\n" + " },\n" + ' "edges": [\n' + ' {"src":"","dst":"","label":"string","protocol":"HTTPS|HTTP|TCP|gRPC|AMQP|unknown","data":["PII"],"id":"optional","confidence":"stated|implied|assumed"}\n' + " ],\n" + ' "zones": {\n' + ' "": {"id":"","name":"string","parent_id":null,"confidence":"stated|implied|assumed"}\n' + " }\n" + " },\n" + ' "assumptions": ["..."],\n' + ' "clarifying_questions": ["..."]\n' + "}\n\n" + "Rules:\n" + "- Use Threat Thinker's native Graph IR names: nodes, edges, and zones. Do not use components, data_flows, or trust_boundaries keys.\n" + "- Do NOT invent components that are not reasonably implied by the description.\n" + '- Every node, edge, and zone must include confidence: "stated", "implied", or "assumed".\n' + "- Put ambiguity you resolved in assumptions. Put missing information the user should answer in clarifying_questions.\n" + "- Draw zones where the description implies a change in control: internet to internal, tenant to tenant, user device to server, or first-party to third-party SaaS.\n" + "- If the description is too thin to build a useful DFD, return an empty graph with clarifying_questions instead of guessing.\n" + "- Prefer a small correct DFD over a large speculative one.\n" + "- Node ids and zone ids must be stable ASCII identifiers using lowercase letters, digits, underscores, or hyphens.\n" + "- Edges must reference existing node ids with src and dst.\n" + "- Return ONLY the JSON object, no other text or formatting.\n" +) + # LLM-driven threat inference prompts LLM_SYSTEM = ( "You are Threat Thinker, an expert security analyst. " diff --git a/src/threat_thinker/llm/inference.py b/src/threat_thinker/llm/inference.py index a8e935b..c713573 100644 --- a/src/threat_thinker/llm/inference.py +++ b/src/threat_thinker/llm/inference.py @@ -7,6 +7,8 @@ from threat_thinker.models import Graph, Threat from threat_thinker.constants import ( + DFD_SYSTEM, + DFD_INSTRUCTIONS, HINT_SYSTEM, HINT_INSTRUCTIONS, LLM_SYSTEM, @@ -18,6 +20,7 @@ # Token budgets tuned for the JSON-heavy responses we expect from each flow. HINT_INFERENCE_MAX_TOKENS = 4096 +DFD_GENERATION_MAX_TOKENS = 8000 THREAT_INFERENCE_MAX_TOKENS = ( 10000 # Headroom for 10-12 verbose multilingual threats with evidence metadata ) @@ -32,6 +35,24 @@ "policies": {"type": "object"}, }, } +DFD_JSON_SCHEMA: Dict = { + "type": "object", + "properties": { + "summary": {"type": "string"}, + "graph": { + "type": "object", + "properties": { + "nodes": {"type": "object"}, + "edges": {"type": "array", "items": {"type": "object"}}, + "zones": {"type": "object"}, + }, + "required": ["nodes", "edges", "zones"], + }, + "assumptions": {"type": "array", "items": {"type": "string"}}, + "clarifying_questions": {"type": "array", "items": {"type": "string"}}, + }, + "required": ["summary", "graph", "assumptions", "clarifying_questions"], +} THREAT_JSON_SCHEMA: Dict = { "type": "object", "properties": { @@ -103,6 +124,62 @@ def _validate_hints_payload(payload: dict) -> None: raise ValueError("'policies' must be an object when present") +def _validate_dfd_payload(payload: dict) -> None: + if not isinstance(payload, dict): + raise ValueError("DFD payload must be a JSON object") + graph = payload.get("graph") + if not isinstance(graph, dict): + raise ValueError("DFD payload missing 'graph' object") + nodes = graph.get("nodes") + edges = graph.get("edges") + zones = graph.get("zones") + if not isinstance(nodes, dict): + raise ValueError("DFD graph.nodes must be an object keyed by node id") + if not isinstance(edges, list): + raise ValueError("DFD graph.edges must be an array") + if not isinstance(zones, dict): + raise ValueError("DFD graph.zones must be an object keyed by zone id") + node_ids = set() + for node_key, node in nodes.items(): + if not isinstance(node, dict): + raise ValueError(f"DFD node '{node_key}' must be an object") + node_id = str(node.get("id") or node_key).strip() + if not node_id: + raise ValueError("DFD node id cannot be empty") + if not str(node.get("label") or "").strip(): + raise ValueError(f"DFD node '{node_id}' must include label") + if node.get("confidence") not in {"stated", "implied", "assumed"}: + raise ValueError(f"DFD node '{node_id}' must include valid confidence") + node_ids.add(node_id) + for zone_key, zone in zones.items(): + if not isinstance(zone, dict): + raise ValueError(f"DFD zone '{zone_key}' must be an object") + zone_id = str(zone.get("id") or zone_key).strip() + if not zone_id: + raise ValueError("DFD zone id cannot be empty") + if not str(zone.get("name") or "").strip(): + raise ValueError(f"DFD zone '{zone_id}' must include name") + if zone.get("confidence") not in {"stated", "implied", "assumed"}: + raise ValueError(f"DFD zone '{zone_id}' must include valid confidence") + for index, edge in enumerate(edges): + if not isinstance(edge, dict): + raise ValueError(f"DFD edge at index {index} must be an object") + src = str(edge.get("src") or "").strip() + dst = str(edge.get("dst") or "").strip() + if not src or not dst: + raise ValueError(f"DFD edge at index {index} must include src and dst") + if src not in node_ids or dst not in node_ids: + raise ValueError( + f"DFD edge at index {index} references unknown nodes '{src}' -> '{dst}'" + ) + if edge.get("confidence") not in {"stated", "implied", "assumed"}: + raise ValueError(f"DFD edge at index {index} must include valid confidence") + if not isinstance(payload.get("assumptions"), list): + raise ValueError("DFD assumptions must be an array") + if not isinstance(payload.get("clarifying_questions"), list): + raise ValueError("DFD clarifying_questions must be an array") + + def _validate_threats_payload(payload: dict) -> None: if not isinstance(payload, dict): raise ValueError("Threat payload must be a JSON object") @@ -299,6 +376,73 @@ def llm_infer_hints( return data +def llm_generate_dfd_from_description( + description: str, + api: str, + model: str, + aws_profile: str = None, + aws_region: str = None, + ollama_host: str = None, + prompt_token_limit: Optional[int] = None, + lang: str = "en", +) -> dict: + """ + Use LLM to generate a native Graph IR DFD from a natural-language system description. + """ + description = (description or "").strip() + if not description: + raise ValueError("description is required") + + if lang == "en": + lang_instruction = "" + else: + lang_name = _get_language_name(lang) + lang_instruction = ( + f"Write human-readable DFD content in {lang_name}: graph node labels, " + "zone names, edge labels, summary, assumptions, clarifying questions, " + "and notes. Keep JSON field names, node ids, zone ids, protocol values, " + "and confidence enum values in English/ASCII exactly as specified.\n\n" + ) + + user_prompt = ( + "Here is the system description:\n\n" + "\n" + f"{description}\n" + "\n\n" + f"{lang_instruction}" + "Build a DFD per the rules. Mark each element's confidence honestly. " + "If the description is too vague, return clarifying questions rather than guessing.\n\n" + f"{DFD_INSTRUCTIONS}" + ) + + _validate_prompt_token_limit( + system_prompt=DFD_SYSTEM, + user_prompt=user_prompt, + api=api, + model=model, + prompt_token_limit=prompt_token_limit, + ) + + llm_client = LLMClient( + api=api, + model=model, + aws_profile=aws_profile, + aws_region=aws_region, + ollama_host=ollama_host, + ) + return _call_llm_json_with_retry( + lambda: llm_client.call_llm( + system_prompt=DFD_SYSTEM, + user_prompt=user_prompt, + response_format={"type": "json_object"}, + json_schema=DFD_JSON_SCHEMA, + temperature=0.1, + max_tokens=DFD_GENERATION_MAX_TOKENS, + ), + _validate_dfd_payload, + ) + + def llm_rerank_chunks( query: str, chunks: List[dict], diff --git a/src/threat_thinker/main.py b/src/threat_thinker/main.py index b4a3a95..dbe855f 100644 --- a/src/threat_thinker/main.py +++ b/src/threat_thinker/main.py @@ -58,8 +58,14 @@ detect_input_format, load_input, ) +from threat_thinker.business_context import ( + BusinessContextDfdResult, + dfd_result_from_payload, + dfd_result_to_sidecar_json, +) from threat_thinker.hint_processor import merge_llm_hints from threat_thinker.llm.inference import ( + llm_generate_dfd_from_description, llm_infer_hints, llm_infer_threats, llm_rerank_chunks, @@ -158,6 +164,10 @@ def _prepare_output_paths( return target_dir, json_path, md_path, html_path +def _prepare_dfd_sidecar_path(report_json_path: Path) -> Path: + return report_json_path.with_name(f"{report_json_path.stem}_dfd.json") + + def _prepare_diff_output_paths( after_report: str, out_dir: str ) -> tuple[Path, Path, Path]: @@ -170,7 +180,7 @@ def _prepare_diff_output_paths( return target_dir, json_path, md_path -def _select_think_input(args) -> tuple[str, str]: +def _select_think_input(args) -> tuple[str | None, str | None]: if args.diagram: diagram_file = args.diagram diagram_format = detect_input_format(diagram_file) @@ -198,11 +208,63 @@ def _select_think_input(args) -> tuple[str, str]: if args.ir: return args.ir, INPUT_FORMAT_IR - ui.error( - "No diagram file specified", - "Please specify a diagram file using --diagram, --mermaid, --drawio, --threat-dragon, --image, or --ir", + return None, None + + +def _combine_text_blocks(*blocks: str | None) -> str | None: + combined = "\n\n".join(block.strip() for block in blocks if block and block.strip()) + return combined or None + + +def _load_document_text(paths: list[str], model: str, *, label: str) -> str | None: + if not paths: + return None + try: + docs = load_context_documents(paths, model) + doc_count, token_count, sources = context_summary(docs) + ui.success( + f"Loaded {doc_count} {label} document(s), approximately {token_count} tokens" + ) + ui.info(f"{label.title()} documents: {', '.join(sources)}") + return format_context_documents(docs) + except ContextDocumentError as e: + ui.error(f"Failed to load {label} documents", str(e)) + sys.exit(2) + + +def _load_description_text(args) -> str | None: + inline = "\n\n".join( + text.strip() + for text in (getattr(args, "description", None) or []) + if text.strip() ) - sys.exit(2) + file_text = _load_document_text( + getattr(args, "description_file", None) or [], + args.llm_model, + label="description", + ) + return _combine_text_blocks(inline, file_text) + + +def _load_context_text(args) -> str | None: + context_paths = list(getattr(args, "context", None) or []) + context_paths.extend(getattr(args, "context_file", None) or []) + return _load_document_text(context_paths, args.llm_model, label="context") + + +def _show_dfd_notes(result: BusinessContextDfdResult) -> None: + if result.summary: + ui.info("Generated DFD summary", result.summary) + if result.assumptions: + ui.warning( + "Generated DFD assumptions", + "\n".join(f"- {item}" for item in result.assumptions), + ) + if result.clarifying_questions: + ui.warning( + "Generated DFD clarifying questions", + "\n".join(f"- {item}" for item in result.clarifying_questions), + ) def main(): @@ -248,6 +310,27 @@ def main(): default=[], help="Business context document path to inject into the threat prompt. Repeat for multiple PDF, Markdown, or text files.", ) + p_think.add_argument( + "--context-file", + type=str, + action="append", + default=[], + help="Alias for --context. Business context document path to inject into the threat prompt.", + ) + p_think.add_argument( + "--description", + type=str, + action="append", + default=[], + help="Natural-language system description. Used to generate a DFD when no diagram input is provided; also injected into the threat prompt.", + ) + p_think.add_argument( + "--description-file", + type=str, + action="append", + default=[], + help="System description document path. Repeat for multiple PDF, Markdown, or text files.", + ) p_think.add_argument( "--infer-hints", action="store_true", @@ -493,14 +576,26 @@ def main(): # Set verbose mode set_verbose(args.verbose) + has_context_files = bool(args.context or args.context_file) + has_description_input = bool(args.description or args.description_file) + # Set up progress tracking - total_steps = 5 + (1 if args.rag else 0) + (1 if args.context else 0) - ui.set_total_steps( - total_steps - ) # Parse, Infer hints, (Context), (Retrieve), Analyze threats, Denoise, Export + total_steps = ( + 5 + + (1 if args.rag else 0) + + (1 if has_context_files else 0) + + (1 if has_description_input else 0) + ) + ui.set_total_steps(total_steps) - # Determine diagram file and format + # Determine optional diagram file and format diagram_file, diagram_format = _select_think_input(args) + if not diagram_file and not has_description_input: + ui.error( + "No input specified", + "Provide a diagram with --diagram/--mermaid/--drawio/--threat-dragon/--image/--ir, or provide --description/--description-file to generate a DFD.", + ) + sys.exit(2) supported_apis = ["openai", "anthropic", "bedrock", "ollama"] if args.llm_api.lower() not in supported_apis: @@ -589,37 +684,90 @@ def main(): ui.error("--prompt-token-limit must be a positive integer.") sys.exit(2) - # 1) Parse diagram to skeleton graph (+ metrics) - ui.step("Parsing architecture diagram") - ui.info(f"Loading {diagram_format} diagram: {diagram_file}") + description_text = None + if has_description_input: + ui.step("Loading system description") + description_text = _load_description_text(args) + if not description_text: + ui.error( + "System description is empty", + "Provide text with --description or readable files with --description-file.", + ) + sys.exit(2) + + dfd_result = None + if diagram_file and diagram_format: + # 1) Parse diagram to skeleton graph (+ metrics) + ui.step("Parsing architecture diagram") + ui.info(f"Loading {diagram_format} diagram: {diagram_file}") - thinking = ui.create_thinking_indicator("Parsing diagram structure") - thinking.start() + thinking = ui.create_thinking_indicator("Parsing diagram structure") + thinking.start() - try: - g, metrics = load_input( - diagram_format, - diagram_file, - drawio_page=args.drawio_page, - api=args.llm_api, - model=args.llm_model, - aws_profile=args.aws_profile, - aws_region=args.aws_region, - ollama_host=ollama_host, - ) + try: + g, metrics = load_input( + diagram_format, + diagram_file, + drawio_page=args.drawio_page, + api=args.llm_api, + model=args.llm_model, + aws_profile=args.aws_profile, + aws_region=args.aws_region, + ollama_host=ollama_host, + ) - thinking.stop() - ui.success("Successfully parsed diagram") - ui.show_metrics_summary(metrics) - ui.debug("Parsed graph details", str(g)) + thinking.stop() + ui.success("Successfully parsed diagram") + ui.show_metrics_summary(metrics) + ui.debug("Parsed graph details", str(g)) - except Exception as e: - thinking.stop() - ui.error("Failed to parse diagram", str(e)) - sys.exit(2) + except Exception as e: + thinking.stop() + ui.error("Failed to parse diagram", str(e)) + sys.exit(2) + else: + ui.step("Generating DFD from system description") + thinking = ui.create_thinking_indicator( + "AI is reconstructing the architecture graph" + ) + thinking.start() + try: + payload = llm_generate_dfd_from_description( + description_text or "", + args.llm_api, + args.llm_model, + args.aws_profile, + args.aws_region, + ollama_host, + args.prompt_token_limit, + lang=args.lang, + ) + dfd_result = dfd_result_from_payload(payload) + dfd_result.metrics.total_lines = len( + (description_text or "").splitlines() + ) + g = dfd_result.graph + metrics = dfd_result.metrics + thinking.stop() + ui.success( + f"Generated DFD with {len(g.nodes)} nodes and {len(g.edges)} edges" + ) + ui.show_metrics_summary(metrics) + _show_dfd_notes(dfd_result) + ui.debug("Generated DFD graph details", str(g)) + if not g.nodes: + ui.error( + "System description is too vague to generate a useful DFD", + "Answer the clarifying questions and rerun with a more specific --description.", + ) + sys.exit(2) + except Exception as e: + thinking.stop() + ui.error("Failed to generate DFD from system description", str(e)) + sys.exit(2) # 2) (Optional) LLM-based attribute inference from skeleton - if args.infer_hints: + if args.infer_hints and diagram_file: ui.step("Inferring node and edge attributes") ui.thinking( "AI is analyzing diagram components to infer security-relevant attributes" @@ -660,24 +808,15 @@ def main(): thinking.stop() ui.error("Failed to infer hints", str(e)) sys.exit(2) - else: + elif diagram_file: ui.step("Skipping attribute inference") ui.info("Using basic component attributes from diagram") - business_context_text = None - if args.context: + context_text = None + if has_context_files: ui.step("Loading business context") - try: - context_docs = load_context_documents(args.context, args.llm_model) - doc_count, token_count, sources = context_summary(context_docs) - business_context_text = format_context_documents(context_docs) - ui.success( - f"Loaded {doc_count} business context document(s), approximately {token_count} tokens" - ) - ui.info(f"Context documents: {', '.join(sources)}") - except ContextDocumentError as e: - ui.error("Failed to load business context", str(e)) - sys.exit(2) + context_text = _load_context_text(args) + business_context_text = _combine_text_blocks(description_text, context_text) rag_context_text = None retrieval = None @@ -805,7 +944,7 @@ def _rerank_with_llm(q, candidates): # 6) Export ui.step("Generating reports") out_dir, out_json, out_md, out_html = _prepare_output_paths( - diagram_file, args.out_dir, args.out_name + diagram_file or "description", args.out_dir, args.out_name ) ui.info( f"Exporting reports to {out_dir} " @@ -829,6 +968,12 @@ def _rerank_with_llm(q, candidates): ui.success(f"JSON report saved to: {out_json}") ui.success(f"Markdown report saved to: {out_md}") ui.success(f"HTML report saved to: {out_html}") + if dfd_result: + dfd_path = _prepare_dfd_sidecar_path(out_json) + dfd_path.write_text( + dfd_result_to_sidecar_json(dfd_result), encoding="utf-8" + ) + ui.success(f"Generated DFD sidecar saved to: {dfd_path}") if args.verbose: print("\nJSON Output:") diff --git a/src/threat_thinker/webui.py b/src/threat_thinker/webui.py index d49e15b..0ab1f1f 100644 --- a/src/threat_thinker/webui.py +++ b/src/threat_thinker/webui.py @@ -15,6 +15,10 @@ import gradio as gr import threat_thinker.main as cli +from threat_thinker.business_context import ( + dfd_result_from_payload, + dfd_result_to_sidecar_json, +) from threat_thinker.constants import AI_OUTPUT_DISCLAIMER_MD from threat_thinker.input_loader import ( INPUT_FORMAT_DRAWIO, @@ -53,7 +57,10 @@ retrieve_context_for_graph, attach_rag_sources_to_threats, ) -from threat_thinker.llm.inference import llm_rerank_chunks +from threat_thinker.llm.inference import ( + llm_generate_dfd_from_description, + llm_rerank_chunks, +) from threat_thinker.rag.local import SUPPORTED_EXTENSIONS @@ -370,12 +377,13 @@ def _generate_diff_report( def _generate_report( + system_description: str, + context_files, input_method: str, diagram_text: str, diagram_format: str, drawio_page: str, image_file: str, - context_files, infer_hints: bool, llm_api: str, llm_model: str, @@ -394,22 +402,27 @@ def _generate_report( rag_candidates: int, rag_min_score: float, prompt_token_limit: int, -) -> Tuple[str, str, Optional[str], Optional[str], Optional[str], Optional[str]]: - # Validate input based on method - if input_method == "Text": - diagram_text = (diagram_text or "").strip() - if not diagram_text: - raise gr.Error("Diagram input is required.") - - diagram_format = _validate_text_input_format(diagram_format) - drawio_page = (drawio_page or "").strip() or None - else: # Image - if not image_file: - raise gr.Error("Image file is required when using image input method.") - - # Validate image file format - from pathlib import Path +) -> Tuple[ + str, + str, + Optional[str], + Optional[str], + Optional[str], + Optional[str], + Optional[str], +]: + system_description = (system_description or "").strip() + diagram_text = (diagram_text or "").strip() + diagram_format = _validate_text_input_format(diagram_format) + drawio_page = (drawio_page or "").strip() or None + has_text_diagram = input_method == "Text" and bool(diagram_text) + has_image_diagram = input_method == "Image" and bool(image_file) + has_diagram = has_text_diagram or has_image_diagram + if not has_diagram and not system_description: + raise gr.Error("System description is required when no diagram is provided.") + context_paths = _normalize_context_uploads(context_files) + if has_image_diagram: ext = Path(image_file).suffix.lower() supported_formats = {".jpg", ".jpeg", ".png", ".gif", ".bmp", ".webp"} if ext not in supported_formats: @@ -437,7 +450,7 @@ def _generate_report( or "http://localhost:11434" ) if llm_api == "ollama": - if input_method == "Image": + if has_image_diagram: raise gr.Error( "Image diagrams are not supported with the Ollama backend. " "Use OpenAI/Anthropic/Bedrock for image extraction or provide Mermaid/Draw.io/Threat Dragon/IR input." @@ -477,24 +490,22 @@ def _generate_report( # Prepare diagram file path diagram_path = None - if input_method == "Text": - # Determine file extension based on format + if has_text_diagram: diagram_path = _write_temp_file( diagram_text, suffix_for_text_input(diagram_format) ) - else: # Image - diagram_path = image_file # Use the uploaded file directly + elif has_image_diagram: + diagram_path = image_file status_lines = [] - context_paths = _normalize_context_uploads(context_files) - business_context_text = None + business_context_text = system_description or None + dfd_result = None rag_context_text = None retrieval = None rerank_fn = None try: - # Parse diagram based on input method and format - if input_method == "Text": + if has_text_diagram: graph, metrics = load_input( diagram_format, diagram_path, @@ -503,7 +514,7 @@ def _generate_report( status_lines.append( f"Parsed {diagram_format} diagram: {len(graph.nodes)} nodes, {len(graph.edges)} edges." ) - else: # Image + elif has_image_diagram: graph, metrics = load_input( "image", diagram_path, @@ -516,6 +527,38 @@ def _generate_report( status_lines.append( f"Parsed image diagram: {len(graph.nodes)} nodes, {len(graph.edges)} edges." ) + else: + payload = llm_generate_dfd_from_description( + system_description, + llm_api, + llm_model, + aws_profile, + aws_region, + ollama_host, + prompt_token_limit_val, + lang=lang, + ) + dfd_result = dfd_result_from_payload(payload) + dfd_result.metrics.total_lines = len(system_description.splitlines()) + graph = dfd_result.graph + metrics = dfd_result.metrics + status_lines.append( + f"Generated DFD from system description: {len(graph.nodes)} nodes, {len(graph.edges)} edges." + ) + if dfd_result.summary: + status_lines.append(f"DFD summary: {dfd_result.summary}") + if dfd_result.assumptions: + status_lines.append("Assumptions: " + "; ".join(dfd_result.assumptions)) + if dfd_result.clarifying_questions: + status_lines.append( + "Clarifying questions: " + + "; ".join(dfd_result.clarifying_questions) + ) + if not graph.nodes: + raise gr.Error( + "System description is too vague to generate a useful DFD. " + "Answer the clarifying questions and retry." + ) status_lines.append( f"Import success ~{metrics.import_success_rate * 100:.1f}% " @@ -523,7 +566,7 @@ def _generate_report( f"labels {metrics.node_labels_parsed}/{metrics.node_label_candidates})" ) - if infer_hints: + if has_diagram and infer_hints: skeleton = json.dumps( { "nodes": [ @@ -554,7 +597,12 @@ def _generate_report( try: context_docs = load_context_documents(context_paths, llm_model) doc_count, token_count, sources = context_summary(context_docs) - business_context_text = format_context_documents(context_docs) + context_text = format_context_documents(context_docs) + business_context_text = "\n\n".join( + text + for text in [business_context_text, context_text] + if text and text.strip() + ) status_lines.append( f"Loaded {doc_count} business context document(s), approximately {token_count} tokens: {', '.join(sources)}." ) @@ -663,12 +711,22 @@ def _rerank_with_llm(q, candidates): download_md_path = _write_temp_file(md_report, ".md") download_json_path = _write_temp_file(json_report, ".json") download_html_path = _write_temp_file(html_report, ".html") + dfd_download_path = None + if dfd_result: + dfd_download_path = _write_temp_file( + dfd_result_to_sidecar_json(dfd_result), ".dfd.json" + ) download_paths = {download_md_path, download_json_path, download_html_path} if td_download_path: download_paths.add(td_download_path) + if dfd_download_path: + download_paths.add(dfd_download_path) _DOWNLOAD_PATHS.update(download_paths) + status_lines.append("Report generated successfully.") + status_text = "\n".join(status_lines) report_text = ( + f"Status:\n{status_text}\n\n" f"JSON Report:\n{json_report}\n\n" f"Markdown Report:\n{md_report}\n\n" f"HTML Report:\n{html_report}" @@ -676,8 +734,6 @@ def _rerank_with_llm(q, candidates): if td_report: report_text += f"\n\nThreat Dragon Report:\n{td_report}" - status_lines.append("Report generated successfully.") - markdown_report = md_report return ( @@ -687,6 +743,7 @@ def _rerank_with_llm(q, candidates): download_json_path, download_html_path, td_download_path, + dfd_download_path, ) except gr.Error: raise @@ -696,7 +753,7 @@ def _rerank_with_llm(q, candidates): finally: # clean up intermediate files; keep the report download file around cleanup_paths = [] - if input_method == "Text" and diagram_path: + if has_text_diagram and diagram_path: cleanup_paths.append(diagram_path) for path in cleanup_paths: @@ -718,53 +775,62 @@ def _build_webui() -> gr.Blocks: with gr.Tabs(): with gr.Tab("Think - Threat Analysis"): - # Input method selection - input_method = gr.Radio( - label="Input Method - Choose whether to input diagram as text or upload an image file", - choices=["Text", "Image"], - value="Text", - ) - - # Text input (visible by default) - diagram_input = gr.TextArea( - label="Diagram Content", - placeholder="Paste your diagram content here (Mermaid, Draw.io XML, Threat Dragon JSON, or native IR JSON)...", - lines=20, + system_description_input = gr.TextArea( + label="System Description", + placeholder=( + "Describe the system, users, data, deployment, and external services. " + "Example: Customers use a web app to manage orders. The app runs on AWS behind a load balancer, stores PII in Postgres, and sends emails through a third-party provider." + ), + lines=10, autofocus=True, - visible=True, - ) - diagram_format_input = gr.Radio( - label="Diagram Format", - choices=[ - INPUT_FORMAT_MERMAID, - INPUT_FORMAT_DRAWIO, - INPUT_FORMAT_THREAT_DRAGON, - INPUT_FORMAT_IR, - ], - value=INPUT_FORMAT_MERMAID, - visible=True, - ) - drawio_page_input = gr.Textbox( - label="Draw.io Page (optional)", - placeholder="Page id, name, or 0-based index", - visible=False, - ) - - # Image input (hidden by default) - image_input = gr.File( - label="Upload Diagram Image (JPG, PNG, GIF, BMP, WebP)", - file_types=["image"], - type="filepath", - visible=False, ) context_files_input = gr.File( - label="Business Context (PDF, Markdown, Text)", + label="Business Context (supplemental PDF, Markdown, Text)", file_types=sorted(SUPPORTED_CONTEXT_EXTENSIONS), type="filepath", file_count="multiple", ) + with gr.Accordion( + "Advanced: provide an existing DFD or diagram", open=False + ): + input_method = gr.Radio( + label="Diagram Input Method", + choices=["Text", "Image"], + value="Text", + ) + + diagram_input = gr.TextArea( + label="Diagram Content", + placeholder="Paste Mermaid, Draw.io XML, Threat Dragon JSON, or native IR JSON...", + lines=16, + visible=True, + ) + diagram_format_input = gr.Radio( + label="Diagram Format", + choices=[ + INPUT_FORMAT_MERMAID, + INPUT_FORMAT_DRAWIO, + INPUT_FORMAT_THREAT_DRAGON, + INPUT_FORMAT_IR, + ], + value=INPUT_FORMAT_MERMAID, + visible=True, + ) + drawio_page_input = gr.Textbox( + label="Draw.io Page (optional)", + placeholder="Page id, name, or 0-based index", + visible=False, + ) + + image_input = gr.File( + label="Upload Diagram Image (JPG, PNG, GIF, BMP, WebP)", + file_types=["image"], + type="filepath", + visible=False, + ) + with gr.Row(): llm_api_input = gr.Dropdown( label="LLM API", @@ -912,16 +978,20 @@ def _build_webui() -> gr.Blocks: download_td_output = gr.File( label="Download Threat Dragon JSON (Threat Dragon inputs only)", ) + download_dfd_output = gr.File( + label="Download generated DFD JSON (description inputs only)", + ) generate_button.click( fn=_generate_report, inputs=[ + system_description_input, + context_files_input, input_method, diagram_input, diagram_format_input, drawio_page_input, image_input, - context_files_input, infer_hints_input, llm_api_input, llm_model_input, @@ -948,6 +1018,7 @@ def _build_webui() -> gr.Blocks: download_json_output, download_html_output, download_td_output, + download_dfd_output, ], api_name=False, ) diff --git a/tests/test_business_context.py b/tests/test_business_context.py new file mode 100644 index 0000000..bf85c0a --- /dev/null +++ b/tests/test_business_context.py @@ -0,0 +1,80 @@ +from threat_thinker.business_context import ( + dfd_result_from_payload, + dfd_result_to_sidecar_dict, +) + + +def test_dfd_result_from_payload_builds_graph_and_sidecar_confidence(): + payload = { + "summary": "Customers use a web app backed by an API and database.", + "graph": { + "zones": { + "internet": { + "id": "internet", + "name": "Internet", + "confidence": "implied", + }, + "private": { + "id": "private", + "name": "Private", + "parent_id": "internet", + "confidence": "assumed", + }, + }, + "nodes": { + "customer": { + "id": "customer", + "label": "Customer", + "type": "actor", + "zones": ["internet"], + "confidence": "stated", + }, + "api": { + "id": "api", + "label": "API", + "type": "service", + "zones": ["internet", "private"], + "data": ["PII"], + "auth": True, + "confidence": "implied", + }, + }, + "edges": [ + { + "src": "customer", + "dst": "api", + "label": "uses", + "protocol": "HTTPS", + "confidence": "implied", + } + ], + }, + "assumptions": ["API is server-side."], + "clarifying_questions": [], + } + + result = dfd_result_from_payload(payload) + sidecar = dfd_result_to_sidecar_dict(result) + + assert result.graph.source_format == "description" + assert result.graph.nodes["api"].zone == "Private" + assert result.graph.edges[0].src == "customer" + assert result.element_confidence["nodes"]["customer"] == "stated" + assert result.element_confidence["edges"]["customer->api:uses"] == "implied" + assert sidecar["graph"]["nodes"]["api"]["data"] == ["PII"] + assert sidecar["assumptions"] == ["API is server-side."] + + +def test_dfd_result_from_payload_allows_empty_graph_with_questions(): + payload = { + "summary": "Not enough information.", + "graph": {"nodes": {}, "edges": [], "zones": {}}, + "assumptions": [], + "clarifying_questions": ["Who uses the system?"], + } + + result = dfd_result_from_payload(payload) + + assert result.graph.nodes == {} + assert result.graph.edges == [] + assert result.clarifying_questions == ["Who uses the system?"] diff --git a/tests/test_cli_outputs.py b/tests/test_cli_outputs.py index 532c235..fcff709 100644 --- a/tests/test_cli_outputs.py +++ b/tests/test_cli_outputs.py @@ -9,6 +9,7 @@ import threat_thinker.main as cli from threat_thinker.main import ( _prepare_diff_output_paths, + _prepare_dfd_sidecar_path, _prepare_output_paths, _select_think_input, ) @@ -61,6 +62,12 @@ def test_prepare_diff_output_paths_use_after_stem(tmp_path: Path): assert md_path.name == "new-report_diff.md" +def test_prepare_dfd_sidecar_path_uses_report_stem(tmp_path: Path): + report_path = tmp_path / "description_report.json" + + assert _prepare_dfd_sidecar_path(report_path).name == "description_report_dfd.json" + + def test_version_command_prints_installed_version(monkeypatch, capsys): monkeypatch.setattr(cli, "get_threat_thinker_version", lambda: "9.8.7") monkeypatch.setattr(sys, "argv", ["threat-thinker", "version"]) @@ -134,3 +141,19 @@ def test_select_think_input_keeps_json_autodetect_as_threat_dragon(): assert diagram_file == str(fixture_path) assert diagram_format == INPUT_FORMAT_THREAT_DRAGON + + +def test_select_think_input_allows_description_without_diagram(): + args = SimpleNamespace( + diagram=None, + mermaid=None, + drawio=None, + threat_dragon=None, + image=None, + ir=None, + ) + + diagram_file, diagram_format = _select_think_input(args) + + assert diagram_file is None + assert diagram_format is None diff --git a/tests/test_llm_context_prompt.py b/tests/test_llm_context_prompt.py index 2d072d0..2ae31de 100644 --- a/tests/test_llm_context_prompt.py +++ b/tests/test_llm_context_prompt.py @@ -89,3 +89,89 @@ def call_llm(self, *, system_prompt, user_prompt, **kwargs): ) assert "include `rag_sources`" not in captured["user_prompt"] + + +def test_llm_generate_dfd_from_description_requests_native_graph_ir(monkeypatch): + captured = {} + + class _Client: + def __init__(self, *args, **kwargs): + pass + + def call_llm(self, *, system_prompt, user_prompt, **kwargs): + captured["system_prompt"] = system_prompt + captured["user_prompt"] = user_prompt + captured["json_schema"] = kwargs.get("json_schema") + return """ + { + "summary": "A web app stores customer data.", + "graph": { + "nodes": { + "user": {"id": "user", "label": "User", "confidence": "stated"}, + "web": {"id": "web", "label": "Web App", "confidence": "stated"} + }, + "edges": [ + {"src": "user", "dst": "web", "label": "uses", "confidence": "implied"} + ], + "zones": {} + }, + "assumptions": [], + "clarifying_questions": [] + } + """ + + monkeypatch.setattr(inference, "LLMClient", _Client) + + payload = inference.llm_generate_dfd_from_description( + "Users sign in to a web app.", + "openai", + "gpt-4.1", + prompt_token_limit=60000, + lang="ja", + ) + + assert payload["graph"]["nodes"]["web"]["label"] == "Web App" + assert "nodes, edges, and zones" in captured["user_prompt"] + assert "components, data_flows, or trust_boundaries" in captured["user_prompt"] + assert "Write human-readable DFD content in Japanese" in captured["user_prompt"] + assert "Keep JSON field names, node ids, zone ids" in captured["user_prompt"] + assert captured["json_schema"]["required"] == [ + "summary", + "graph", + "assumptions", + "clarifying_questions", + ] + + +def test_llm_generate_dfd_from_description_rejects_unknown_edge(monkeypatch): + class _Client: + def __init__(self, *args, **kwargs): + pass + + def call_llm(self, *, system_prompt, user_prompt, **kwargs): + return """ + { + "summary": "Bad graph.", + "graph": { + "nodes": { + "user": {"id": "user", "label": "User", "confidence": "stated"} + }, + "edges": [ + {"src": "user", "dst": "missing", "confidence": "implied"} + ], + "zones": {} + }, + "assumptions": [], + "clarifying_questions": [] + } + """ + + monkeypatch.setattr(inference, "LLMClient", _Client) + + with pytest.raises(RuntimeError, match="references unknown nodes"): + inference.llm_generate_dfd_from_description( + "Users sign in.", + "openai", + "gpt-4.1", + prompt_token_limit=60000, + ) diff --git a/tests/test_webui_helpers.py b/tests/test_webui_helpers.py index 0289c70..71842b1 100644 --- a/tests/test_webui_helpers.py +++ b/tests/test_webui_helpers.py @@ -86,3 +86,17 @@ def test_build_webui_smoke(): assert isinstance(demo, gr.Blocks) assert demo.title == "Threat Thinker WebUI" assert len(demo.blocks) > 0 + + +def test_build_webui_has_system_description_entrypoint(): + demo = webui._build_webui() + labels = { + getattr(block, "label", None) + for block in demo.blocks.values() + if getattr(block, "label", None) + } + + assert "System Description" in labels + assert "Business Context (supplemental PDF, Markdown, Text)" in labels + assert "Diagram Content" in labels + assert "Download generated DFD JSON (description inputs only)" in labels From b104fe5fcaa1114055be6ba5ceec0d8661385dbc Mon Sep 17 00:00:00 2001 From: melonattacker <41631269+melonattacker@users.noreply.github.com> Date: Tue, 28 Apr 2026 22:08:07 +0900 Subject: [PATCH 2/3] feat: Increase DFD generation token limit to accommodate larger responses --- src/threat_thinker/llm/inference.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/threat_thinker/llm/inference.py b/src/threat_thinker/llm/inference.py index c713573..8f91b1f 100644 --- a/src/threat_thinker/llm/inference.py +++ b/src/threat_thinker/llm/inference.py @@ -20,7 +20,7 @@ # Token budgets tuned for the JSON-heavy responses we expect from each flow. HINT_INFERENCE_MAX_TOKENS = 4096 -DFD_GENERATION_MAX_TOKENS = 8000 +DFD_GENERATION_MAX_TOKENS = 16000 THREAT_INFERENCE_MAX_TOKENS = ( 10000 # Headroom for 10-12 verbose multilingual threats with evidence metadata ) From 9e6dde7b96e14db77ecb1691cffc435f0b86d7cc Mon Sep 17 00:00:00 2001 From: melonattacker <41631269+melonattacker@users.noreply.github.com> Date: Tue, 28 Apr 2026 22:40:01 +0900 Subject: [PATCH 3/3] feat: Implement default report base name logic and enhance DFD report generation for empty cases --- src/threat_thinker/main.py | 16 ++++++++++- src/threat_thinker/webui.py | 42 ++++++++++++++++++++++++++-- tests/test_cli_outputs.py | 9 ++++++ tests/test_webui_helpers.py | 56 +++++++++++++++++++++++++++++++++++++ 4 files changed, 119 insertions(+), 4 deletions(-) diff --git a/src/threat_thinker/main.py b/src/threat_thinker/main.py index dbe855f..fac856c 100644 --- a/src/threat_thinker/main.py +++ b/src/threat_thinker/main.py @@ -164,6 +164,16 @@ def _prepare_output_paths( return target_dir, json_path, md_path, html_path +def _default_report_base_name( + diagram_file: str | None, description_files: list[str] | None = None +) -> str: + if diagram_file: + return Path(diagram_file).stem or "threat" + if description_files: + return Path(description_files[0]).stem or "description" + return "description" + + def _prepare_dfd_sidecar_path(report_json_path: Path) -> Path: return report_json_path.with_name(f"{report_json_path.stem}_dfd.json") @@ -943,8 +953,12 @@ def _rerank_with_llm(q, candidates): # 6) Export ui.step("Generating reports") + base_name = args.out_name or _default_report_base_name( + diagram_file, + getattr(args, "description_file", None) or [], + ) out_dir, out_json, out_md, out_html = _prepare_output_paths( - diagram_file or "description", args.out_dir, args.out_name + diagram_file or base_name, args.out_dir, base_name ) ui.info( f"Exporting reports to {out_dir} " diff --git a/src/threat_thinker/webui.py b/src/threat_thinker/webui.py index 0ab1f1f..1de14df 100644 --- a/src/threat_thinker/webui.py +++ b/src/threat_thinker/webui.py @@ -117,6 +117,23 @@ def _write_temp_file(content: str, suffix: str) -> str: return tmp.name +def _build_incomplete_dfd_markdown(result) -> str: + lines = [ + "## System Description Needs More Detail", + "", + "Threat inference did not run because the generated DFD was empty.", + ] + if result.summary: + lines.extend(["", f"Summary: {result.summary}"]) + if result.assumptions: + lines.extend(["", "### Assumptions"]) + lines.extend(f"- {item}" for item in result.assumptions) + if result.clarifying_questions: + lines.extend(["", "### Clarifying Questions"]) + lines.extend(f"- {item}" for item in result.clarifying_questions) + return "\n".join(lines) + + def _validate_text_input_format(diagram_format: str) -> str: value = (diagram_format or INPUT_FORMAT_MERMAID).strip().lower() if value not in TEXT_INPUT_FORMATS: @@ -555,9 +572,28 @@ def _generate_report( + "; ".join(dfd_result.clarifying_questions) ) if not graph.nodes: - raise gr.Error( - "System description is too vague to generate a useful DFD. " - "Answer the clarifying questions and retry." + status_lines.append( + "Threat inference skipped because the generated DFD is empty. " + "Expand the system description and retry." + ) + _cleanup_downloads() + dfd_download_path = _write_temp_file( + dfd_result_to_sidecar_json(dfd_result), ".dfd.json" + ) + _DOWNLOAD_PATHS.add(dfd_download_path) + status_text = "\n".join(status_lines) + report_text = ( + f"Status:\n{status_text}\n\n" + f"Generated DFD JSON:\n{dfd_result_to_sidecar_json(dfd_result)}" + ) + return ( + _build_incomplete_dfd_markdown(dfd_result), + report_text, + None, + None, + None, + None, + dfd_download_path, ) status_lines.append( diff --git a/tests/test_cli_outputs.py b/tests/test_cli_outputs.py index fcff709..64024a8 100644 --- a/tests/test_cli_outputs.py +++ b/tests/test_cli_outputs.py @@ -8,6 +8,7 @@ import threat_thinker.main as cli from threat_thinker.main import ( + _default_report_base_name, _prepare_diff_output_paths, _prepare_dfd_sidecar_path, _prepare_output_paths, @@ -51,6 +52,14 @@ def test_prepare_output_paths_with_override(tmp_path: Path): assert html_path.name == "custom-base_report.html" +def test_default_report_base_name_prefers_description_file_stem(): + assert _default_report_base_name(None, ["docs/drone-system.txt"]) == "drone-system" + + +def test_default_report_base_name_falls_back_to_description(): + assert _default_report_base_name(None, []) == "description" + + def test_prepare_diff_output_paths_use_after_stem(tmp_path: Path): out_dir, json_path, md_path = _prepare_diff_output_paths( "results/new-report.json", tmp_path / "diffs" diff --git a/tests/test_webui_helpers.py b/tests/test_webui_helpers.py index 71842b1..2e5bcda 100644 --- a/tests/test_webui_helpers.py +++ b/tests/test_webui_helpers.py @@ -100,3 +100,59 @@ def test_build_webui_has_system_description_entrypoint(): assert "Business Context (supplemental PDF, Markdown, Text)" in labels assert "Diagram Content" in labels assert "Download generated DFD JSON (description inputs only)" in labels + + +def test_generate_report_returns_clarifying_questions_for_empty_dfd(monkeypatch): + monkeypatch.setattr( + webui, + "llm_generate_dfd_from_description", + lambda *args, **kwargs: { + "summary": "A drone delivery service.", + "graph": {"nodes": {}, "edges": [], "zones": {}}, + "assumptions": ["The service has a backend API."], + "clarifying_questions": [ + "Who places orders?", + "How are payments processed?", + ], + }, + ) + + markdown_report, report_text, md_path, json_path, html_path, td_path, dfd_path = ( + webui._generate_report( + system_description="Drone food delivery.", + context_files=[], + input_method="Text", + diagram_text="", + diagram_format="mermaid", + drawio_page="", + image_file="", + infer_hints=False, + llm_api="openai", + llm_model="gpt-4.1", + aws_profile="", + aws_region="", + ollama_host="", + topn=10, + min_confidence=0.5, + require_asvs=False, + lang="en", + use_rag=False, + kb_names=[], + rag_topk=5, + rag_strategy=webui.DEFAULT_RAG_STRATEGY, + rag_reranker=webui.DEFAULT_RAG_RERANKER, + rag_candidates=webui.DEFAULT_RAG_CANDIDATES, + rag_min_score=webui.DEFAULT_RAG_MIN_SCORE, + prompt_token_limit=1000, + ) + ) + + assert "System Description Needs More Detail" in markdown_report + assert "Clarifying Questions" in markdown_report + assert "Who places orders?" in report_text + assert "Threat inference skipped because the generated DFD is empty." in report_text + assert md_path is None + assert json_path is None + assert html_path is None + assert td_path is None + assert dfd_path is not None