diff --git a/README.md b/README.md index 5c46bab..8ef903a 100644 --- a/README.md +++ b/README.md @@ -37,6 +37,7 @@ If you find STRIDE GPT useful, please consider supporting the project: - **MITRE ATT&CK & ATLAS mapping**: Threats are annotated with standardized adversary technique IDs (MITRE ATT&CK Enterprise for traditional infrastructure attacks, ATLAS for ML/LLM-specific attacks) — surfaced as columns in markdown, linked pills in HTML, and `mitre_attack` properties in SARIF - **Architectural pattern detection**: Automatically detects RAG pipelines, multi-agent systems, code execution environments, tool ecosystems, and more from application descriptions (inspired by [CSA MAESTRO](https://cloudsecurityalliance.org/blog/2025/02/06/agentic-ai-threat-modeling-framework-maestro)) - Multi-modal: Use architecture diagrams, flowcharts, etc. as inputs for threat modelling across all supported vision-capable models +- **Data Flow Diagrams**: Generate DFDs from your application description (or parse an uploaded DFD image), edit the Mermaid source live, and feed the confirmed diagram back into the Threat Model and Attack Tree prompts as the authoritative system model. CLI `/analyze` also emits a system-level DFD alongside its findings - Generates attack trees to enumerate possible attack paths - Suggests possible mitigations for identified threats - Supports DREAD risk scoring for identified threats @@ -70,6 +71,10 @@ This video is an excellent resource for anyone interested in understanding how S ## Changelog +### Unreleased + +- **Data Flow Diagram support** (closes #56): New "Data Flow Diagram" tab in the Streamlit UI generates DFDs from the application description, parses uploaded DFD images via vision-capable models, and renders an editable Mermaid source pane with a live preview. Ticking *"Use this DFD for the threat model"* stores the confirmed diagram in session state and splices it into subsequent Threat Model and Attack Tree prompts as the authoritative system model — closing the description → DFD → review → refined threat model loop the issue asks for. The CLI's `/analyze` agent also produces a system-level DFD after synthesis, rendered as a Mermaid block in the markdown report, carried in the JSON `data_flow_diagram` field, and shown in the HTML view via a CDN-loaded Mermaid runtime. DFD generation is wrapped so a bad diagram never fails a good report. + ### Version 0.17 (latest) - **MITRE ATT&CK and ATLAS integration** (closes #36): Threats emitted by `/analyze` and `/quick` can now carry standardized adversary technique IDs from MITRE ATT&CK Enterprise (v17.1) and MITRE ATLAS (2026.05). Two new progressive-disclosure reference cards (`mitre_enterprise`, `mitre_atlas`) instruct the agent to use only the catalogued IDs and names, eliminating ID hallucination. Cards are regenerated from upstream STIX/YAML via `scripts/refresh_mitre_cards.py` rather than written from model recall. Techniques render as a `MITRE ATT&CK` column in markdown tables, clickable sky-tinted pills in HTML, and `properties.mitre_attack` arrays in SARIF. diff --git a/apps/web/attack_tree.py b/apps/web/attack_tree.py index 53604d8..37ec1ff 100644 --- a/apps/web/attack_tree.py +++ b/apps/web/attack_tree.py @@ -8,7 +8,12 @@ # Function to create a prompt to generate an attack tree def create_attack_tree_prompt( - app_type, authentication, internet_facing, sensitive_data, app_input + app_type, + authentication, + internet_facing, + sensitive_data, + app_input, + confirmed_dfd: str | None = None, ): prompt = f"""APPLICATION TYPE: {app_type} AUTHENTICATION METHODS: {authentication} @@ -17,6 +22,10 @@ def create_attack_tree_prompt( APPLICATION DESCRIPTION: {app_input} """ + if confirmed_dfd: + from stride_gpt.core.prompts import dfd_to_prompt_section + prompt += dfd_to_prompt_section(confirmed_dfd) + # Add GenAI attack vectors for both Generative AI and Agentic AI applications if app_type in ["Generative AI application", "Agentic AI application"]: prompt += """ diff --git a/apps/web/dfd.py b/apps/web/dfd.py new file mode 100644 index 0000000..a07239c --- /dev/null +++ b/apps/web/dfd.py @@ -0,0 +1,111 @@ +"""Streamlit-facing provider wrappers for Data Flow Diagram generation. + +Mirrors `apps/web/attack_tree.py`: thin per-provider wrappers around the +core `generate_dfd` / `parse_dfd_from_image` functions, with session-state +plumbing for thinking/reasoning displays. +""" + +from __future__ import annotations + +import streamlit as st + +from stride_gpt.core.dfd import generate_dfd, parse_dfd_from_image +from stride_gpt.core.schemas import LLMConfig + + +__all__ = [ + "get_dfd_anthropic", + "get_dfd_from_image_anthropic", + "get_dfd_from_image_google", + "get_dfd_from_image_openai", + "get_dfd_google", + "get_dfd_groq", + "get_dfd_lm_studio", + "get_dfd_mistral", + "get_dfd_openai", +] + + +# --------------------------------------------------------------------------- +# Generation from a textual description +# --------------------------------------------------------------------------- + + +def get_dfd_openai(api_key: str, model_name: str, prompt: str) -> str: + config = LLMConfig(provider="OpenAI API", model_name=model_name, api_key=api_key) + mermaid, _ = generate_dfd(config, prompt) + return mermaid + + +def get_dfd_anthropic(api_key: str, model_name: str, prompt: str) -> str: + config = LLMConfig( + provider="Anthropic API", + model_name=model_name, + api_key=api_key, + use_thinking=st.session_state.get("use_thinking", False), + ) + mermaid, response = generate_dfd(config, prompt) + if response.thinking: + st.session_state["last_thinking_content"] = response.thinking + return mermaid + + +def get_dfd_google(api_key: str, model_name: str, prompt: str) -> str: + config = LLMConfig(provider="Google AI API", model_name=model_name, api_key=api_key) + mermaid, response = generate_dfd(config, prompt) + if response.thinking: + st.session_state["last_thinking_content"] = response.thinking + return mermaid + + +def get_dfd_groq(api_key: str, model_name: str, prompt: str) -> str: + config = LLMConfig(provider="Groq API", model_name=model_name, api_key=api_key) + mermaid, response = generate_dfd(config, prompt) + if response.reasoning: + with st.expander("View model's reasoning process", expanded=False): + st.write(response.reasoning) + return mermaid + + +def get_dfd_lm_studio( + lm_studio_endpoint: str, model_name: str, prompt: str, api_key: str = "not-needed" +) -> str: + config = LLMConfig( + provider="LM Studio Server", + model_name=model_name, + api_key=api_key, + api_base=lm_studio_endpoint, + ) + mermaid, _ = generate_dfd(config, prompt) + return mermaid + + +def get_dfd_mistral(api_key: str, model_name: str, prompt: str) -> str: + config = LLMConfig(provider="Mistral API", model_name=model_name, api_key=api_key) + mermaid, _ = generate_dfd(config, prompt) + return mermaid + + +# --------------------------------------------------------------------------- +# Parsing a user-uploaded DFD image +# --------------------------------------------------------------------------- + + +def get_dfd_from_image_openai(api_key: str, model_name: str, base64_image: str) -> str: + config = LLMConfig(provider="OpenAI API", model_name=model_name, api_key=api_key) + mermaid, _ = parse_dfd_from_image(config, base64_image) + return mermaid + + +def get_dfd_from_image_anthropic( + api_key: str, model_name: str, base64_image: str, media_type: str = "image/png" +) -> str: + config = LLMConfig(provider="Anthropic API", model_name=model_name, api_key=api_key) + mermaid, _ = parse_dfd_from_image(config, base64_image, media_type=media_type) + return mermaid + + +def get_dfd_from_image_google(api_key: str, model_name: str, base64_image: str) -> str: + config = LLMConfig(provider="Google AI API", model_name=model_name, api_key=api_key) + mermaid, _ = parse_dfd_from_image(config, base64_image) + return mermaid diff --git a/apps/web/main.py b/apps/web/main.py index 6a700d0..fbd5f49 100644 --- a/apps/web/main.py +++ b/apps/web/main.py @@ -33,6 +33,18 @@ get_attack_tree_lm_studio, get_attack_tree_mistral, ) +from dfd import ( + get_dfd_anthropic, + get_dfd_from_image_anthropic, + get_dfd_from_image_google, + get_dfd_from_image_openai, + get_dfd_google, + get_dfd_groq, + get_dfd_lm_studio, + get_dfd_mistral, + get_dfd_openai, +) +from stride_gpt.core.prompts import create_dfd_prompt from dread import ( create_dread_assessment_prompt, dread_json_to_markdown, @@ -1003,10 +1015,26 @@ def on_model_selection_change(): ) +# ------------------ Provider API keys (module scope) ------------------ # +# +# `load_env_variables()` reads .env into FUNCTION-LOCAL variables, and the +# sidebar's password fields only write to st.session_state — neither path +# defines `openai_api_key` / `anthropic_api_key` / etc. at module scope. +# Hoist them here, after the sidebar has run, so the bare names referenced +# throughout the tab bodies (image analysis, threat model, attack tree, +# mitigations, DREAD, test cases) resolve to the live values. Streamlit's +# hot-reload was masking this by leaking stale locals across reruns. +openai_api_key = st.session_state.get("openai_api_key", "") +anthropic_api_key = st.session_state.get("anthropic_api_key", "") +google_api_key = st.session_state.get("google_api_key", "") +mistral_api_key = st.session_state.get("mistral_api_key", "") +groq_api_key = st.session_state.get("groq_api_key", "") + + # ------------------ Main App UI ------------------ # -tab1, tab2, tab3, tab4, tab5 = st.tabs( - ["Threat Model", "Attack Tree", "Mitigations", "DREAD", "Test Cases"] +tab1, tab2, dfd_tab, tab3, tab4, tab5 = st.tabs( + ["Threat Model", "Attack Tree", "Data Flow Diagram", "Mitigations", "DREAD", "Test Cases"] ) with tab1: @@ -1196,13 +1224,17 @@ def encode_image(uploaded_file): if threat_model_submit_button and st.session_state.get("app_input"): app_input = st.session_state["app_input"] # Retrieve from session state - # Generate the prompt using the create_prompt function + # Generate the prompt using the create_prompt function. If the user + # has confirmed a DFD on the Data Flow Diagram tab, splice it in as + # an authoritative system model so the threats align with the + # diagram's components and trust boundaries. threat_model_prompt = create_threat_model_prompt( app_type, authentication, internet_facing, sensitive_data, app_input, + confirmed_dfd=st.session_state.get("confirmed_dfd"), ) # Clear thinking content when switching models or starting a new operation @@ -1341,13 +1373,16 @@ def encode_image(uploaded_file): if attack_tree_submit_button and st.session_state.get("app_input"): app_input = st.session_state.get("app_input") - # Generate the prompt using the create_attack_tree_prompt function + # Generate the prompt using the create_attack_tree_prompt function. + # When a DFD has been confirmed on the Data Flow Diagram tab it + # is appended so the attack paths reference the same components. attack_tree_prompt = create_attack_tree_prompt( app_type, authentication, internet_facing, sensitive_data, app_input, + confirmed_dfd=st.session_state.get("confirmed_dfd"), ) # Clear thinking content when switching models or starting a new operation @@ -1447,6 +1482,222 @@ def encode_image(uploaded_file): st.error(f"Error generating attack tree: {e}") +# ------------------ Data Flow Diagram Generation ------------------ # + +with dfd_tab: + st.markdown( + """ +A Data Flow Diagram (DFD) makes the system under analysis legible: it names every external entity, +process, and data store, and shows the data flows between them. Crucially, it lets you mark +**trust boundaries** — the lines a threat must cross to attack the system. DFDs are foundational to +threat modelling; reviewing one is the cleanest way to confirm STRIDE-GPT has understood your system +the same way you have. + +Three ways to use this tab: + +1. **Generate a DFD** from your application description, then edit it to match your understanding. +2. **Upload an existing DFD image** to convert it into editable Mermaid form. +3. Once the diagram looks right, tick **Use this DFD for the threat model** to feed it into the next + Threat Model and Attack Tree runs. +""" + ) + st.markdown("""---""") + + # --- Section 1: Inputs (generate from description OR parse an image) --- # + + dfd_input_col1, dfd_input_col2 = st.columns([1, 1]) + + with dfd_input_col1: + st.markdown("**Generate from description**") + dfd_generate_button = st.button( + label="Generate DFD", + help="Uses your application description and sidebar settings to draft a DFD.", + ) + + with dfd_input_col2: + st.markdown("**Or parse an existing DFD image**") + # Vision support set is the same as for architecture diagrams. + dfd_vision_supported = ( + (model_provider == "OpenAI API" and ( + selected_model.startswith("gpt-4") or selected_model.startswith("gpt-5") + )) + or model_provider == "Google AI API" + or (model_provider == "Anthropic API" and selected_model.startswith("claude")) + ) + if dfd_vision_supported: + uploaded_dfd_image = st.file_uploader( + "Upload DFD image", + type=["jpg", "jpeg", "png"], + key="dfd_image_uploader", + help="PNG or JPEG. The diagram is parsed into editable Mermaid below.", + ) + else: + uploaded_dfd_image = None + st.caption( + "Image parsing requires a vision-capable model " + "(OpenAI GPT-4/5, Anthropic Claude, or Google Gemini)." + ) + + # --- Section 2: Handle Generation --- # + + # API keys live in st.session_state once the sidebar has run (loaded from + # .env in load_env_variables and/or typed into the password field). The + # bare names openai_api_key / anthropic_api_key etc. are only ever local + # to load_env_variables, so always go through session_state here. + _dfd_api_key = st.session_state.get( + { + "OpenAI API": "openai_api_key", + "Anthropic API": "anthropic_api_key", + "Google AI API": "google_api_key", + "Mistral API": "mistral_api_key", + "Groq API": "groq_api_key", + }.get(model_provider, ""), + "", + ) + + if dfd_generate_button and st.session_state.get("app_input"): + dfd_prompt = create_dfd_prompt( + app_type, + authentication, + internet_facing, + sensitive_data, + st.session_state["app_input"], + ) + if model_provider != "LM Studio Server" and not _dfd_api_key: + st.error( + f"Please enter your {model_provider} key in the sidebar (or set it in " + "your .env file) before generating a DFD." + ) + else: + with st.spinner("Generating Data Flow Diagram..."): + try: + if model_provider == "OpenAI API": + dfd_mermaid = get_dfd_openai(_dfd_api_key, selected_model, dfd_prompt) + elif model_provider == "Anthropic API": + dfd_mermaid = get_dfd_anthropic(_dfd_api_key, selected_model, dfd_prompt) + elif model_provider == "Google AI API": + dfd_mermaid = get_dfd_google(_dfd_api_key, selected_model, dfd_prompt) + elif model_provider == "Mistral API": + dfd_mermaid = get_dfd_mistral(_dfd_api_key, selected_model, dfd_prompt) + elif model_provider == "Groq API": + dfd_mermaid = get_dfd_groq(_dfd_api_key, selected_model, dfd_prompt) + elif model_provider == "LM Studio Server": + dfd_mermaid = get_dfd_lm_studio( + st.session_state["lm_studio_endpoint"], + selected_model, + dfd_prompt, + st.session_state.get("lm_studio_api_key", ""), + ) + else: + dfd_mermaid = "" + # Generation invalidates any previous confirmation — the user + # needs to re-tick "Use this DFD" to opt back in. + st.session_state["dfd_mermaid"] = dfd_mermaid + st.session_state.pop("confirmed_dfd", None) + except Exception as e: + st.error(f"Error generating DFD: {e}") + elif dfd_generate_button and not st.session_state.get("app_input"): + st.error("Please enter your application details on the Threat Model tab first.") + + # --- Section 3: Handle Image Parsing --- # + + if uploaded_dfd_image is not None: + # Re-parse on every change of the upload — the file_uploader returns + # the same object across reruns, so we key on the file name + size. + upload_key = f"{uploaded_dfd_image.name}:{uploaded_dfd_image.size}" + if st.session_state.get("dfd_last_upload_key") != upload_key: + if not _dfd_api_key: + st.error( + f"Please enter your {model_provider} key in the sidebar before " + "uploading a DFD image." + ) + else: + with st.spinner("Parsing DFD image..."): + try: + image_bytes = uploaded_dfd_image.read() + base64_image = base64.b64encode(image_bytes).decode("utf-8") + media_type = ( + "image/png" if uploaded_dfd_image.type == "image/png" else "image/jpeg" + ) + if model_provider == "OpenAI API": + dfd_mermaid = get_dfd_from_image_openai( + _dfd_api_key, selected_model, base64_image + ) + elif model_provider == "Anthropic API": + dfd_mermaid = get_dfd_from_image_anthropic( + _dfd_api_key, selected_model, base64_image, media_type + ) + elif model_provider == "Google AI API": + dfd_mermaid = get_dfd_from_image_google( + _dfd_api_key, selected_model, base64_image + ) + else: + dfd_mermaid = "" + st.session_state["dfd_mermaid"] = dfd_mermaid + st.session_state["dfd_last_upload_key"] = upload_key + st.session_state.pop("confirmed_dfd", None) + except Exception as e: + st.error(f"Error parsing DFD image: {e}") + + # --- Section 4: Editor + Preview --- # + + st.markdown("---") + st.markdown("**Mermaid source** (edit freely — the preview updates on rerun):") + dfd_editor_value = st.text_area( + "DFD Mermaid code", + value=st.session_state.get("dfd_mermaid", ""), + height=260, + key="dfd_editor", + label_visibility="collapsed", + placeholder="No DFD yet. Use 'Generate DFD' or upload an image to populate this editor.", + ) + # Keep dfd_mermaid in sync with the editor so it survives reruns. + st.session_state["dfd_mermaid"] = dfd_editor_value + + if dfd_editor_value.strip(): + st.markdown("**Diagram preview:**") + try: + mermaid(dfd_editor_value) + except Exception as e: + st.warning(f"Could not render preview: {e}") + + # --- Section 5: Confirm + Download --- # + + confirmed = st.checkbox( + "Use this DFD for the threat model", + value=bool(st.session_state.get("confirmed_dfd")), + help=( + "When ticked, this DFD is included in the Threat Model and Attack Tree " + "prompts so the analysis aligns with the diagram." + ), + key="dfd_confirm_checkbox", + ) + # Keep confirmed_dfd in lockstep with the checkbox AND the editor. + if confirmed: + st.session_state["confirmed_dfd"] = dfd_editor_value + st.success( + "DFD confirmed — it will be used as the system model in the next " + "Threat Model and Attack Tree runs." + ) + else: + st.session_state.pop("confirmed_dfd", None) + st.caption( + "Not in use. Tick the checkbox above to feed this DFD into the next analysis." + ) + + dfd_dl_col1, dfd_dl_col2, _ = st.columns([1, 1, 3]) + with dfd_dl_col1: + st.download_button( + label="Download Diagram Code", + data=dfd_editor_value, + file_name="data_flow_diagram.md", + mime="text/plain", + help="Download the Mermaid source for this DFD.", + ) + with dfd_dl_col2: + st.link_button("Open Mermaid Live", "https://mermaid.live") + + # ------------------ Mitigations Generation ------------------ # with tab3: diff --git a/stride_gpt/agent/html_report.py b/stride_gpt/agent/html_report.py index 3657c21..ca980b6 100644 --- a/stride_gpt/agent/html_report.py +++ b/stride_gpt/agent/html_report.py @@ -61,6 +61,7 @@ def render_html(report: AnalysisReport) -> str: "generated_at": datetime.now(timezone.utc).isoformat(), "target": report.plan.target_path, "overview": report.plan.overall_description, + "data_flow_diagram": report.data_flow_diagram, "subsystems": [ { "name": f.subsystem, @@ -86,6 +87,7 @@ def render_html_from_json(data: dict[str, Any]) -> str: target_name = Path(target).name or target or "(unknown target)" generated_at = _format_generated_at(data.get("generated_at", "")) overview = (data.get("overview") or "").strip() + dfd_mermaid = (data.get("data_flow_diagram") or "").strip() subsystems = data.get("subsystems") or [] cross_cutting = data.get("cross_cutting_threats") or [] metadata = data.get("metadata") or {} @@ -106,6 +108,9 @@ def render_html_from_json(data: dict[str, Any]) -> str: if overview: parts.append(_render_overview(overview)) + if dfd_mermaid: + parts.append(_render_dfd(dfd_mermaid)) + for sub in subsystems: parts.append(_render_subsystem(sub)) @@ -137,6 +142,10 @@ def _scaffold(*, target_name: str, body: str) -> str:
{body}
+ """ @@ -195,6 +204,21 @@ def _render_overview(overview: str) -> str: """ +def _render_dfd(dfd_mermaid: str) -> str: + """Render the system-level Data Flow Diagram as a Mermaid block. + + Escapes the diagram source — LLM-generated, so treat as untrusted markup + even though Mermaid would re-escape internally. The CDN-loaded Mermaid + runtime in the scaffold picks up `.mermaid` blocks on startup. + """ + return f"""
+

Data Flow Diagram

+
+ 
{html.escape(dfd_mermaid)}
+
+
""" + + def _render_footer(metadata: dict[str, Any]) -> str: bits: list[str] = [] llm_calls = metadata.get("llm_calls") diff --git a/stride_gpt/agent/loop.py b/stride_gpt/agent/loop.py index 6c13607..7a55e53 100644 --- a/stride_gpt/agent/loop.py +++ b/stride_gpt/agent/loop.py @@ -202,10 +202,24 @@ def run_analysis( llm_calls += 1 progress.synthesis_done(len(cross_cutting)) + # --- Phase 4: System-level DFD --- + # Optional pass — gives the report a visual map of components and trust + # boundaries. Wrapped in try/except so a bad DFD never nukes a good + # report. Skipped when no findings exist or the call budget is spent. + data_flow_diagram: str | None = None + if findings and (not max_llm_calls or llm_calls < max_llm_calls): + progress.status("Generating system-level Data Flow Diagram...") + try: + data_flow_diagram = _generate_system_dfd(models, plan, findings) + llm_calls += 1 + except Exception: + data_flow_diagram = None + report = AnalysisReport( plan=plan, findings=findings, cross_cutting_threats=cross_cutting, + data_flow_diagram=data_flow_diagram, metadata=_build_metadata( models, plan, llm_calls=llm_calls, tool_calls=tool_calls, subsystems_analyzed=len(findings), @@ -461,6 +475,49 @@ def _synthesize(models: ModelPair, findings: list[SubsystemFinding]) -> list[dic return data.get("cross_cutting_threats", []) +def _generate_system_dfd( + models: ModelPair, plan: AnalysisPlan, findings: list[SubsystemFinding] +) -> str | None: + """Produce a system-level DFD in Mermaid form, or None on any failure. + + Single architect-tier LLM call. The architect already saw the + cross-cutting synthesis; this gives it the same finding-level summary + plus the plan, and asks for a JSON DFD that we render to Mermaid via + the canonical converter. + + DFD generation is auxiliary — never let it fail the whole report. + Callers wrap in try/except too as belt-and-braces. + """ + from stride_gpt.core.dfd import generate_dfd + from stride_gpt.core.prompts import create_dfd_prompt + + # Build a description that combines the plan overview with the per- + # subsystem context the agent actually surfaced. Files-analyzed lists + # are noisy and not useful here. + description_parts = [plan.overall_description, "", "Subsystems and files identified:"] + for sub in plan.subsystems: + description_parts.append(f"- {sub.name}: {sub.description}") + + description_parts.append("") + description_parts.append("Subsystem threat findings (for context only):") + for finding in findings: + if not finding.threats: + continue + threat_types = sorted({t.get("Threat Type", "") for t in finding.threats if t.get("Threat Type")}) + description_parts.append(f"- {finding.subsystem}: {', '.join(threat_types)}") + + prompt = create_dfd_prompt( + app_type=plan.detected_app_type, + authentication="See subsystem details", + internet_facing="Inferred from findings", + sensitive_data="Inferred from findings", + app_input="\n".join(description_parts), + ) + + mermaid, _ = generate_dfd(models.for_architect(), prompt) + return mermaid or None + + def _truncate_findings_to_fit( architect: LLMConfig, findings: list[SubsystemFinding], diff --git a/stride_gpt/agent/report.py b/stride_gpt/agent/report.py index efb0a71..4a9bd53 100644 --- a/stride_gpt/agent/report.py +++ b/stride_gpt/agent/report.py @@ -32,6 +32,15 @@ def render_markdown(report: AnalysisReport) -> str: lines.append(report.plan.overall_description) lines.append("") + # Data Flow Diagram (optional — generated during synthesis) + if report.data_flow_diagram: + lines.append("## Data Flow Diagram") + lines.append("") + lines.append("```mermaid") + lines.append(report.data_flow_diagram) + lines.append("```") + lines.append("") + # Analysis plan summary lines.append("## Analysis Plan") lines.append("") @@ -127,6 +136,7 @@ def render_json(report: AnalysisReport) -> dict[str, Any]: "generated_at": datetime.now(timezone.utc).isoformat(), "target": report.plan.target_path, "overview": report.plan.overall_description, + "data_flow_diagram": report.data_flow_diagram, "subsystems": [ { "name": f.subsystem, @@ -511,6 +521,15 @@ def render_markdown_from_json(data: dict[str, Any]) -> str: lines.append(data.get("overview", "")) lines.append("") + dfd = data.get("data_flow_diagram") + if dfd: + lines.append("## Data Flow Diagram") + lines.append("") + lines.append("```mermaid") + lines.append(dfd) + lines.append("```") + lines.append("") + cross_cutting = data.get("cross_cutting_threats", []) all_threats: list[dict[str, Any]] = [] for sub in data.get("subsystems", []): diff --git a/stride_gpt/core/attack_tree.py b/stride_gpt/core/attack_tree.py index dbffe76..6ab73d3 100644 --- a/stride_gpt/core/attack_tree.py +++ b/stride_gpt/core/attack_tree.py @@ -3,9 +3,13 @@ from __future__ import annotations import json -import re from stride_gpt.core.llm import call_llm +from stride_gpt.core.mermaid_utils import ( + clean_json_response, + clean_mermaid_syntax, + extract_mermaid_code, +) from stride_gpt.core.prompts import ( create_json_structure_prompt, create_reasoning_system_prompt, @@ -14,6 +18,14 @@ from stride_gpt.models import model_uses_completion_tokens +__all__ = [ + "clean_json_response", + "clean_mermaid_syntax", + "convert_tree_to_mermaid", + "extract_mermaid_code", + "generate_attack_tree", +] + def generate_attack_tree(config: LLMConfig, prompt: str) -> tuple[str, LLMResponse]: """Generate attack tree as Mermaid code. Returns (mermaid_string, response).""" @@ -117,106 +129,6 @@ def process_node(node, parent_id=None): return "\n".join(mermaid_lines) -def clean_json_response(response_text): - """ - Clean JSON response by removing any markdown code block markers and finding the JSON content. - - Args: - response_text (str): The raw response text that might contain JSON - - Returns: - str: Cleaned JSON string - """ - # Remove markdown JSON code block if present - json_pattern = r"```json\s*(.*?)\s*```" - match = re.search(json_pattern, response_text, re.DOTALL) - if match: - return match.group(1).strip() - - # If no JSON code block, try to find content between any code blocks - code_pattern = r"```\s*(.*?)\s*```" - match = re.search(code_pattern, response_text, re.DOTALL) - if match: - return match.group(1).strip() - - # If no code blocks, return the original text - return response_text.strip() - - -def extract_mermaid_code(text): - """ - Extract the Mermaid diagram code from text that may contain additional content. - Looks for code between ```mermaid, ``` or just ``` tags, and extracts the graph content. - Also cleans and validates the Mermaid syntax. - - Args: - text (str): The text containing the Mermaid code - - Returns: - str: The cleaned Mermaid code, or the original text if no code block is found - """ - # Try to find code block with explicit mermaid tag - mermaid_pattern = r"```mermaid\s*(graph[\s\S]*?)```" - match = re.search(mermaid_pattern, text, re.MULTILINE) - - if not match: - # Try to find any code block containing graph definition - code_pattern = r"```\s*(graph[\s\S]*?)```" - match = re.search(code_pattern, text, re.MULTILINE) - - code = match.group(1).strip() if match else text.strip() - - # Only proceed if we have a graph definition - if not code.startswith("graph "): - if "graph " in code: - # Find the start of the graph definition - code = code[code.find("graph "):] - else: - return text - - # Clean up common issues in Mermaid syntax - return clean_mermaid_syntax(code) - - -def clean_mermaid_syntax(code): - """ - Clean up common issues in Mermaid syntax. - - Args: - code (str): The Mermaid code to clean - - Returns: - str: The cleaned Mermaid code - """ - # Ensure proper spacing around arrows - code = re.sub(r"(\w+|\]|\)|\})(-->|==>|-.->)(\w+|\[|\(|\{)", r"\1 \2 \3", code) - - # Fix missing brackets around node labels - def fix_node_brackets(match): - node_id = match.group(1) - if not any(c in node_id for c in "[](){}"): - return f"{node_id}[{node_id}]" - return node_id - - code = re.sub(r"(?:^|\s)(\w+)(?:\s|$)", fix_node_brackets, code) - - # Ensure node IDs with spaces are properly quoted - def quote_node_labels(match): - label = match.group(1) - if " " in label and not label.startswith('"'): - return f'["{label}"]' - return f"[{label}]" - - code = re.sub(r"\[(.*?)\]", quote_node_labels, code) - - # Fix parentheses in node labels - def fix_parentheses(match): - label = match.group(1) - if "(" in label or ")" in label: - return f'["{label}"]' - return f"[{label}]" - - code = re.sub(r"\[(.*?)\]", fix_parentheses, code) - - # Ensure proper line endings - return code.replace("\r\n", "\n").strip() +# clean_json_response, extract_mermaid_code, and clean_mermaid_syntax now +# live in stride_gpt.core.mermaid_utils — re-exported above for callers that +# still import them from this module. diff --git a/stride_gpt/core/dfd.py b/stride_gpt/core/dfd.py new file mode 100644 index 0000000..b3521a3 --- /dev/null +++ b/stride_gpt/core/dfd.py @@ -0,0 +1,256 @@ +"""Core Data Flow Diagram (DFD) generation logic. Zero Streamlit imports. + +Mirrors the shape of `core/attack_tree.py`: build a prompt, call the LLM +asking for structured JSON, then convert the JSON into Mermaid `flowchart` +syntax. Trust boundaries render as Mermaid `subgraph` blocks with dashed +styling — the threat-modelling convention. + +Two entry points: + +- `generate_dfd(config, prompt)` — produce a DFD from a textual application + description. +- `parse_dfd_from_image(config, base64_image, media_type)` — parse a + user-uploaded DFD image (PNG/JPEG) into the same canonical JSON form, + then render to Mermaid. Used by the iterative workflow so users can + upload an existing DFD as input. +""" + +from __future__ import annotations + +import json +import re + +from stride_gpt.core.llm import call_llm, call_llm_with_image +from stride_gpt.core.mermaid_utils import clean_json_response, extract_mermaid_code +from stride_gpt.core.prompts.builder import ( + create_dfd_image_analysis_prompt, + create_reasoning_system_prompt, +) +from stride_gpt.core.schemas import LLMConfig, LLMResponse +from stride_gpt.models import model_uses_completion_tokens + + +VALID_NODE_TYPES = ("external_entity", "process", "data_store") + + +DFD_SCHEMA: dict = { + "type": "object", + "properties": { + "nodes": { + "type": "array", + "items": { + "type": "object", + "properties": { + "id": {"type": "string"}, + "label": {"type": "string"}, + "type": {"enum": list(VALID_NODE_TYPES)}, + }, + "required": ["id", "label", "type"], + }, + }, + "edges": { + "type": "array", + "items": { + "type": "object", + "properties": { + "from": {"type": "string"}, + "to": {"type": "string"}, + "label": {"type": "string"}, + }, + "required": ["from", "to", "label"], + }, + }, + "trust_boundaries": { + "type": "array", + "items": { + "type": "object", + "properties": { + "name": {"type": "string"}, + "node_ids": {"type": "array", "items": {"type": "string"}}, + }, + "required": ["name", "node_ids"], + }, + }, + }, + "required": ["nodes", "edges"], +} + + +def generate_dfd(config: LLMConfig, prompt: str) -> tuple[str, LLMResponse]: + """Generate a DFD as Mermaid code. Returns (mermaid_string, response).""" + system_prompt = _get_system_prompt(config) + json_config = config.model_copy(update={"response_format": "json"}) + messages = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": prompt}, + ] + response = call_llm(json_config, messages) + mermaid = _parse_dfd_response(response.content) + return mermaid, response + + +def parse_dfd_from_image( + config: LLMConfig, base64_image: str, media_type: str = "image/png" +) -> tuple[str, LLMResponse]: + """Vision-parse an uploaded DFD image into the canonical Mermaid form. + + Returns (mermaid_string, response). The vision provider is asked for + JSON conforming to `DFD_SCHEMA`; the same fallback parser as + `generate_dfd` covers cases where the model returns Mermaid directly. + """ + prompt = create_dfd_image_analysis_prompt() + # call_llm_with_image takes the prompt as a positional arg and routes + # provider-specific multimodal kwargs internally. JSON response_format + # is set on the config so json-mode is respected where supported. + json_config = config.model_copy(update={"response_format": "json"}) + response = call_llm_with_image(json_config, prompt, base64_image, media_type) + mermaid = _parse_dfd_response(response.content) + return mermaid, response + + +def _get_system_prompt(config: LLMConfig) -> str: + """Reasoning-model variant gets explicit step-by-step framing.""" + if model_uses_completion_tokens(config.model_name): + return create_reasoning_system_prompt( + task_description=( + "Produce a Data Flow Diagram (DFD) for threat modelling as " + "structured JSON." + ), + approach_description=_DFD_RULES_BLOCK, + ) + return _DFD_BASE_PROMPT + + +_DFD_RULES_BLOCK = """Analyse the application and emit a JSON object describing the system as a Data Flow Diagram. + +Rules: +- Use simple alphanumeric IDs (user, web, api, db, etc.) — no spaces or hyphens. +- Every node has a `type`: "external_entity" (actors / 3rd-party services), "process" (application components doing work), or "data_store" (databases, caches, queues, blob stores). +- Every edge has a descriptive `label` naming the data crossing it (e.g. "Login credentials", "Encrypted session token"). +- Group nodes inside the same trust zone with `trust_boundaries` (e.g. "Internal VPC", "Browser/Untrusted"). +- Cover the whole system — every external entity, process, data store, and the flows between them. +- Output JSON only — no commentary. + +Example shape: +{ + "nodes": [ + {"id": "user", "label": "End User", "type": "external_entity"}, + {"id": "web", "label": "Web Frontend", "type": "process"}, + {"id": "api", "label": "API Server", "type": "process"}, + {"id": "db", "label": "User Database", "type": "data_store"} + ], + "edges": [ + {"from": "user", "to": "web", "label": "Login credentials"}, + {"from": "web", "to": "api", "label": "Auth request"}, + {"from": "api", "to": "db", "label": "User lookup"}, + {"from": "db", "to": "api", "label": "Hashed password"} + ], + "trust_boundaries": [ + {"name": "Internal Network", "node_ids": ["api", "db"]} + ] +} + +ONLY RESPOND WITH THE JSON STRUCTURE, NO ADDITIONAL TEXT.""" + + +_DFD_BASE_PROMPT = "Your task is to model the application as a Data Flow Diagram for threat modelling.\n\n" + _DFD_RULES_BLOCK + + +def _parse_dfd_response(content: str) -> str: + """JSON-first, Mermaid-fallback parser. + + The configured response_format=json should give us a JSON object, but + a few providers (LM Studio, Groq with DeepSeek) sometimes emit raw + Mermaid inside a code fence. Try the structured path first and fall + back to extracting Mermaid directly. + """ + try: + cleaned = clean_json_response(content) + dfd_data = json.loads(cleaned) + return convert_dfd_to_mermaid(dfd_data) + except (json.JSONDecodeError, KeyError, TypeError): + return extract_mermaid_code(content, start_keywords=("flowchart", "graph")) + + +def convert_dfd_to_mermaid(dfd_data: dict) -> str: + """Render the canonical DFD JSON form into Mermaid `flowchart TD` syntax. + + Node shapes follow DFD-on-Mermaid convention: + - external_entity -> `[Label]` (rectangle) + - process -> `((Label))` (circle) + - data_store -> `[(Label)]` (cylinder) + + Trust boundaries become Mermaid `subgraph` blocks with dashed styling + applied via a `classDef` on each boundary. + """ + nodes = dfd_data.get("nodes", []) + edges = dfd_data.get("edges", []) + boundaries = dfd_data.get("trust_boundaries", []) or [] + + # Index nodes for shape lookup + which boundary they live in. + node_lookup: dict[str, dict] = {n["id"]: n for n in nodes if "id" in n} + boundary_for: dict[str, int] = {} + for idx, b in enumerate(boundaries): + for nid in b.get("node_ids", []) or []: + boundary_for[nid] = idx + + lines: list[str] = ["flowchart TD"] + + # Emit nodes inside their boundary subgraph, then loose nodes outside. + for idx, boundary in enumerate(boundaries): + boundary_name = boundary.get("name", f"Trust Boundary {idx + 1}") + safe_name = _sanitize_label(boundary_name) + lines.append(f' subgraph tb{idx}["{safe_name}"]') + for nid in boundary.get("node_ids", []) or []: + node = node_lookup.get(nid) + if node is None: + continue + lines.append(" " + _render_node(node)) + lines.append(" end") + + for node in nodes: + if node.get("id") in boundary_for: + continue + lines.append(" " + _render_node(node)) + + for edge in edges: + src = edge.get("from") + dst = edge.get("to") + if not src or not dst: + continue + label = edge.get("label", "") + if label: + lines.append(f' {src} -->|{_sanitize_label(label)}| {dst}') + else: + lines.append(f" {src} --> {dst}") + + # Dashed trust-boundary styling so they read as boundaries, not groups. + if boundaries: + lines.append(" classDef trustBoundary stroke-dasharray: 5 5,fill:transparent;") + for idx in range(len(boundaries)): + lines.append(f" class tb{idx} trustBoundary;") + + return "\n".join(lines) + + +def _render_node(node: dict) -> str: + node_id = node["id"] + label = _sanitize_label(node.get("label", node_id)) + node_type = node.get("type", "process") + if node_type == "external_entity": + return f'{node_id}["{label}"]' + if node_type == "data_store": + return f'{node_id}[("{label}")]' + # process (default) + return f'{node_id}(("{label}"))' + + +def _sanitize_label(label: str) -> str: + """Strip characters that confuse Mermaid label parsing. + + Mermaid treats `"`, `|`, and unescaped brackets as syntax. We replace + them with safe equivalents rather than failing the render. + """ + cleaned = re.sub(r'["|`]', "'", str(label)) + cleaned = cleaned.replace("\n", " ").strip() + return cleaned diff --git a/stride_gpt/core/mermaid_utils.py b/stride_gpt/core/mermaid_utils.py new file mode 100644 index 0000000..df4643c --- /dev/null +++ b/stride_gpt/core/mermaid_utils.py @@ -0,0 +1,86 @@ +"""Shared helpers for LLM → JSON → Mermaid conversion. + +Used by both `core/attack_tree.py` and `core/dfd.py`. Kept provider-agnostic +and free of Streamlit imports. +""" + +from __future__ import annotations + +import re + + +def clean_json_response(response_text: str) -> str: + """Strip markdown code fences from an LLM response, returning the JSON body.""" + json_pattern = r"```json\s*(.*?)\s*```" + match = re.search(json_pattern, response_text, re.DOTALL) + if match: + return match.group(1).strip() + + code_pattern = r"```\s*(.*?)\s*```" + match = re.search(code_pattern, response_text, re.DOTALL) + if match: + return match.group(1).strip() + + return response_text.strip() + + +def extract_mermaid_code(text: str, start_keywords: tuple[str, ...] = ("graph", "flowchart")) -> str: + """Pull a Mermaid diagram out of free-form text. + + Looks for fenced ```mermaid ... ``` or plain ``` ... ``` blocks containing + one of the recognised diagram keywords. Falls through to the original + text if nothing usable is found. + """ + keyword_alternation = "|".join(re.escape(k) for k in start_keywords) + mermaid_pattern = rf"```mermaid\s*((?:{keyword_alternation})[\s\S]*?)```" + match = re.search(mermaid_pattern, text, re.MULTILINE) + + if not match: + code_pattern = rf"```\s*((?:{keyword_alternation})[\s\S]*?)```" + match = re.search(code_pattern, text, re.MULTILINE) + + code = match.group(1).strip() if match else text.strip() + + starts_with_keyword = any(code.startswith(f"{k} ") for k in start_keywords) + if not starts_with_keyword: + for keyword in start_keywords: + marker = f"{keyword} " + if marker in code: + code = code[code.find(marker):] + starts_with_keyword = True + break + if not starts_with_keyword: + return text + + return clean_mermaid_syntax(code) + + +def clean_mermaid_syntax(code: str) -> str: + """Best-effort cleanup of arrow spacing and bracketing in LLM-generated Mermaid.""" + code = re.sub(r"(\w+|\]|\)|\})(-->|==>|-.->)(\w+|\[|\(|\{)", r"\1 \2 \3", code) + + def fix_node_brackets(match: re.Match[str]) -> str: + node_id = match.group(1) + if not any(c in node_id for c in "[](){}"): + return f"{node_id}[{node_id}]" + return node_id + + code = re.sub(r"(?:^|\s)(\w+)(?:\s|$)", fix_node_brackets, code) + + def quote_node_labels(match: re.Match[str]) -> str: + label = match.group(1) + if " " in label and not label.startswith('"'): + return f'["{label}"]' + return f"[{label}]" + + code = re.sub(r"\[(.*?)\]", quote_node_labels, code) + + def fix_parentheses(match: re.Match[str]) -> str: + label = match.group(1) + if "(" in label or ")" in label: + return f'["{label}"]' + return f"[{label}]" + + code = re.sub(r"\[(.*?)\]", fix_parentheses, code) + + return code.replace("\r\n", "\n").strip() diff --git a/stride_gpt/core/prompts/__init__.py b/stride_gpt/core/prompts/__init__.py index c8e567c..7ba7fa1 100644 --- a/stride_gpt/core/prompts/__init__.py +++ b/stride_gpt/core/prompts/__init__.py @@ -8,12 +8,15 @@ create_agentic_stride_prompt_section, create_attack_tree_schema, create_attack_tree_schema_lm_studio, + create_dfd_image_analysis_prompt, + create_dfd_prompt, create_image_analysis_prompt, create_insider_threat_prompt_section, create_json_structure_prompt, create_llm_stride_prompt_section, create_reasoning_system_prompt, create_threat_model_prompt, + dfd_to_prompt_section, ) from stride_gpt.core.prompts.variants import ( AppType, @@ -30,12 +33,15 @@ "create_agentic_stride_prompt_section", "create_attack_tree_schema", "create_attack_tree_schema_lm_studio", + "create_dfd_image_analysis_prompt", + "create_dfd_prompt", "create_image_analysis_prompt", "create_insider_threat_prompt_section", "create_json_structure_prompt", "create_llm_stride_prompt_section", "create_reasoning_system_prompt", "create_threat_model_prompt", + "dfd_to_prompt_section", "load_reference", "quick_base_prompt", ] diff --git a/stride_gpt/core/prompts/builder.py b/stride_gpt/core/prompts/builder.py index a1ca600..59bc721 100644 --- a/stride_gpt/core/prompts/builder.py +++ b/stride_gpt/core/prompts/builder.py @@ -155,14 +155,108 @@ def create_insider_threat_prompt_section() -> str: return "\n" + load_reference("insider_threat") +def dfd_to_prompt_section(dfd_mermaid: str) -> str: + """Render a confirmed DFD into a prompt section. + + When the user has reviewed (and possibly edited) a DFD in the web UI + and confirmed it for use, the Mermaid source is spliced into the + threat-model and attack-tree prompts. The DFD is the user's canonical + statement of the system under analysis, so we surface it explicitly + rather than just appending to the free-form description. + """ + return f""" +CONFIRMED DATA FLOW DIAGRAM (Mermaid): +The user has reviewed the following Data Flow Diagram and confirmed it as +an accurate representation of the system under analysis. Use it as the +authoritative model of components, data flows, and trust boundaries when +identifying threats — pay particular attention to flows that cross trust +boundaries. + +```mermaid +{dfd_mermaid.strip()} +``` +""" + + +def create_dfd_prompt( + app_type: str, + authentication: str, + internet_facing: str, + sensitive_data: str, + app_input: str, +) -> str: + """Build the DFD-generation prompt for a given application profile. + + The output of this prompt is fed to `core.dfd.generate_dfd`, which + expects JSON conforming to `DFD_SCHEMA`. The detailed JSON shape and + rules live in the system prompt set by `core.dfd._get_system_prompt`; + this user prompt provides the application context. + """ + return f"""APPLICATION TYPE: {app_type} +AUTHENTICATION METHODS: {authentication} +INTERNET FACING: {internet_facing} +SENSITIVE DATA: {sensitive_data} +APPLICATION DESCRIPTION: {app_input} + +Produce a Data Flow Diagram (DFD) for this application as JSON conforming +to the schema described in your system instructions. Cover every external +entity, process, and data store you can identify from the description, and +group internal components inside trust boundaries (e.g. "Internal Network", +"Cloud VPC") so the diagram is useful for STRIDE threat modelling. Edge +labels should name the data crossing the flow.""" + + +def create_dfd_image_analysis_prompt() -> str: + """Instruction for vision models parsing a user-supplied DFD image. + + Asks the model to extract structured DFD JSON (nodes/edges/trust + boundaries) so we can re-render the user's existing diagram in our + canonical Mermaid form. Falls back gracefully if the model emits raw + Mermaid in a code fence — `_parse_dfd_response` handles that path. + """ + return """You are reviewing a Data Flow Diagram (DFD) that a security architect has provided. + +Extract the diagram's structure as a JSON object with this exact shape: + +{ + "nodes": [ + {"id": "", "label": "", "type": "external_entity" | "process" | "data_store"} + ], + "edges": [ + {"from": "", "to": "", "label": ""} + ], + "trust_boundaries": [ + {"name": "", "node_ids": ["", "..."]} + ] +} + +Classification rules: +- external_entity: actors outside the system (users, third-party APIs, browsers) +- process: components that transform data (services, functions, workers) +- data_store: persistent stores (databases, queues, caches, file systems) + +If the diagram shows trust boundaries (dashed lines, labelled zones, network +perimeters), capture them in `trust_boundaries`. If no boundaries are shown, +return an empty list. + +ONLY RESPOND WITH THE JSON OBJECT, NO ADDITIONAL TEXT. Do not infer +components that aren't visible in the diagram.""" + + def create_threat_model_prompt( app_type: str, authentication: str, internet_facing: str, sensitive_data: str, app_input: str, + confirmed_dfd: str | None = None, ) -> str: - """Build the threat-model prompt for a given application profile.""" + """Build the threat-model prompt for a given application profile. + + If `confirmed_dfd` is supplied (Mermaid source the user has confirmed + in the DFD tab), it's spliced in as an authoritative system model so + the LLM threats reference the same components and trust boundaries. + """ is_genai = app_type == "Generative AI application" is_agentic = app_type == "Agentic AI application" include_llm_risks = is_genai or is_agentic @@ -244,6 +338,9 @@ def create_threat_model_prompt( prompt += create_agentic_stride_prompt_section() prompt += create_insider_threat_prompt_section() + if confirmed_dfd: + prompt += dfd_to_prompt_section(confirmed_dfd) + prompt += f""" CODE SUMMARY, README CONTENT, AND APPLICATION DESCRIPTION: {app_input} diff --git a/stride_gpt/core/schemas.py b/stride_gpt/core/schemas.py index e35864b..39d7e4d 100644 --- a/stride_gpt/core/schemas.py +++ b/stride_gpt/core/schemas.py @@ -92,6 +92,10 @@ class AnalysisReport(BaseModel): plan: AnalysisPlan findings: list[SubsystemFinding] cross_cutting_threats: list[dict[str, Any]] = [] + # System-level Data Flow Diagram in Mermaid `flowchart` form. Generated + # during synthesis from the full set of subsystem findings. None when + # generation was skipped or failed — DFD is auxiliary, not load-bearing. + data_flow_diagram: str | None = None metadata: dict[str, Any] = {} diff --git a/tests/test_dfd.py b/tests/test_dfd.py new file mode 100644 index 0000000..6fdbe5f --- /dev/null +++ b/tests/test_dfd.py @@ -0,0 +1,275 @@ +"""Tests for Data Flow Diagram generation, parsing, and rendering.""" + +from __future__ import annotations + +import json +from unittest.mock import patch + +from stride_gpt.core.dfd import ( + convert_dfd_to_mermaid, + generate_dfd, + parse_dfd_from_image, + _parse_dfd_response, +) +from stride_gpt.core.prompts.builder import ( + create_dfd_image_analysis_prompt, + create_dfd_prompt, + create_threat_model_prompt, + dfd_to_prompt_section, +) +from stride_gpt.core.schemas import ( + AnalysisReport, + LLMResponse, +) + + +# --------------------------------------------------------------------------- +# convert_dfd_to_mermaid +# --------------------------------------------------------------------------- + + +class TestConvertDfdToMermaid: + def test_basic_three_nodes_two_edges(self): + dfd = { + "nodes": [ + {"id": "user", "label": "End User", "type": "external_entity"}, + {"id": "api", "label": "API Server", "type": "process"}, + {"id": "db", "label": "User Database", "type": "data_store"}, + ], + "edges": [ + {"from": "user", "to": "api", "label": "Login request"}, + {"from": "api", "to": "db", "label": "User lookup"}, + ], + } + result = convert_dfd_to_mermaid(dfd) + + assert result.startswith("flowchart TD") + # External entity → rectangle + assert 'user["End User"]' in result + # Process → circle + assert 'api(("API Server"))' in result + # Data store → cylinder + assert 'db[("User Database")]' in result + # Edges keep their labels + assert "user -->|Login request| api" in result + assert "api -->|User lookup| db" in result + + def test_trust_boundary_emits_subgraph(self): + dfd = { + "nodes": [ + {"id": "user", "label": "User", "type": "external_entity"}, + {"id": "api", "label": "API", "type": "process"}, + {"id": "db", "label": "DB", "type": "data_store"}, + ], + "edges": [ + {"from": "user", "to": "api", "label": "Request"}, + {"from": "api", "to": "db", "label": "Query"}, + ], + "trust_boundaries": [ + {"name": "Internal VPC", "node_ids": ["api", "db"]}, + ], + } + result = convert_dfd_to_mermaid(dfd) + + assert 'subgraph tb0["Internal VPC"]' in result + assert "end" in result + # Boundary styling + assert "stroke-dasharray" in result + assert "class tb0 trustBoundary" in result + # User is outside the boundary + user_idx = result.find("user[") + subgraph_idx = result.find("subgraph tb0") + end_idx = result.find("\n end") + assert user_idx > end_idx or user_idx < subgraph_idx + + def test_handles_label_with_special_chars(self): + """LLM-supplied labels with pipes or quotes shouldn't break Mermaid syntax. + + The outer `"..."` are Mermaid's label delimiters — what matters is that + the *content* between them is free of pipes and stray inner quotes + (we replace `"` with `'` so the label text doesn't terminate early). + """ + dfd = { + "nodes": [{"id": "api", "label": 'Has "quotes" | and pipe', "type": "process"}], + "edges": [], + } + result = convert_dfd_to_mermaid(dfd) + # Extract the label content between Mermaid's `(("..."))` delimiters. + between = result.split('api(("')[1].split('"))')[0] + assert '"' not in between, f"Inner quotes survived sanitization: {between!r}" + assert "|" not in between, f"Pipe survived sanitization: {between!r}" + + +# --------------------------------------------------------------------------- +# _parse_dfd_response +# --------------------------------------------------------------------------- + + +class TestParseDfdResponse: + def test_json_path(self): + payload = json.dumps({ + "nodes": [{"id": "a", "label": "A", "type": "process"}], + "edges": [], + }) + result = _parse_dfd_response(payload) + assert result.startswith("flowchart TD") + assert 'a(("A"))' in result + + def test_json_fenced(self): + content = '```json\n{"nodes": [{"id":"x","label":"X","type":"data_store"}], "edges": []}\n```' + result = _parse_dfd_response(content) + assert 'x[("X")]' in result + + def test_mermaid_fallback_when_json_fails(self): + # Model emitted raw Mermaid instead of JSON — fallback should extract + # the diagram code (`clean_mermaid_syntax` is best-effort and adds + # default-label brackets, matching the attack-tree path's behaviour). + content = "```mermaid\nflowchart TD\n a-->b\n```" + result = _parse_dfd_response(content) + # Fallback returned the body of the fence, not the original wrapped text. + assert "```" not in result + assert "a" in result and "b" in result + assert "flowchart" in result + + +# --------------------------------------------------------------------------- +# generate_dfd / parse_dfd_from_image +# --------------------------------------------------------------------------- + + +class TestGenerateDfd: + def test_uses_json_response_format(self, llm_config): + """generate_dfd must force JSON response_format so providers honour it.""" + captured: dict = {} + + def fake_call(config, messages): + captured["config"] = config + return LLMResponse( + content=json.dumps({ + "nodes": [{"id": "a", "label": "A", "type": "process"}], + "edges": [], + }) + ) + + with patch("stride_gpt.core.dfd.call_llm", side_effect=fake_call): + mermaid, _ = generate_dfd(llm_config, "Describe the system.") + + assert captured["config"].response_format == "json" + # Original config must NOT be mutated. + assert llm_config.response_format == "text" + assert "flowchart TD" in mermaid + + def test_parse_dfd_from_image_uses_call_llm_with_image(self, llm_config): + captured: dict = {} + + def fake_image_call(config, prompt, base64_image, media_type): + captured["config"] = config + captured["media_type"] = media_type + return LLMResponse( + content=json.dumps({ + "nodes": [{"id": "u", "label": "User", "type": "external_entity"}], + "edges": [], + }) + ) + + with patch( + "stride_gpt.core.dfd.call_llm_with_image", side_effect=fake_image_call + ): + mermaid, _ = parse_dfd_from_image(llm_config, "Zm9v", media_type="image/png") + + assert captured["media_type"] == "image/png" + assert captured["config"].response_format == "json" + assert 'u["User"]' in mermaid + + +# --------------------------------------------------------------------------- +# Prompt builders — confirmed_dfd wiring (the iteration loop) +# --------------------------------------------------------------------------- + + +class TestPromptWiring: + def test_dfd_section_appears_when_confirmed_dfd_set(self): + mermaid = "flowchart TD\n user-->api" + prompt = create_threat_model_prompt( + app_type="Web application", + authentication="OAuth2", + internet_facing="Yes", + sensitive_data="PII", + app_input="A web app.", + confirmed_dfd=mermaid, + ) + assert "CONFIRMED DATA FLOW DIAGRAM" in prompt + assert "user-->api" in prompt + + def test_dfd_section_absent_when_no_confirmed_dfd(self): + prompt = create_threat_model_prompt( + app_type="Web application", + authentication="OAuth2", + internet_facing="Yes", + sensitive_data="PII", + app_input="A web app.", + ) + assert "CONFIRMED DATA FLOW DIAGRAM" not in prompt + + def test_dfd_to_prompt_section_contains_mermaid_fence(self): + section = dfd_to_prompt_section("flowchart TD\n a-->b") + assert "```mermaid" in section + assert "flowchart TD" in section + + def test_create_dfd_prompt_carries_application_context(self): + prompt = create_dfd_prompt( + "Web application", "OAuth2", "Yes", "PII", "My description." + ) + assert "APPLICATION TYPE: Web application" in prompt + assert "My description." in prompt + + def test_image_analysis_prompt_requests_canonical_json(self): + prompt = create_dfd_image_analysis_prompt() + # Names the three node-type values so the model can't invent its own. + assert "external_entity" in prompt + assert "process" in prompt + assert "data_store" in prompt + + +# --------------------------------------------------------------------------- +# Agent report rendering — DFD section +# --------------------------------------------------------------------------- + + +class TestAgentReportRendersDfd: + def test_renders_mermaid_fence_when_dfd_present(self, sample_plan, sample_finding): + from stride_gpt.agent.report import render_markdown + + report = AnalysisReport( + plan=sample_plan, + findings=[sample_finding], + data_flow_diagram="flowchart TD\n user --> api", + metadata={}, + ) + md = render_markdown(report) + assert "## Data Flow Diagram" in md + assert "```mermaid" in md + assert "flowchart TD" in md + + def test_no_section_when_dfd_absent(self, sample_report): + from stride_gpt.agent.report import render_markdown + + # sample_report is constructed without data_flow_diagram → defaults to None + assert sample_report.data_flow_diagram is None + md = render_markdown(sample_report) + assert "## Data Flow Diagram" not in md + + def test_render_json_carries_dfd(self, sample_plan, sample_finding): + from stride_gpt.agent.report import render_markdown_from_json, render_json + + report = AnalysisReport( + plan=sample_plan, + findings=[sample_finding], + data_flow_diagram="flowchart TD\n a --> b", + metadata={}, + ) + data = render_json(report) + assert data["data_flow_diagram"] == "flowchart TD\n a --> b" + # And the from-JSON markdown renderer surfaces it too (replay path). + md = render_markdown_from_json(data) + assert "## Data Flow Diagram" in md