diff --git a/README.md b/README.md index cbdb9f5..4303f3f 100644 --- a/README.md +++ b/README.md @@ -37,6 +37,7 @@ If you find STRIDE GPT useful, please consider supporting the project: - **Generative AI support**: Threat modeling for GenAI applications with OWASP LLM Top 10 integration - **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)) +- **Embedded draw.io diagram editor**: Create and edit architecture diagrams directly in STRIDE-GPT using the integrated [diagrams.net](https://www.diagrams.net/) editor — no external tool needed. Diagrams are parsed as XML to extract components, connections, and trust boundaries, providing significantly richer context for threat model generation than image analysis alone. The existing image upload workflow is unchanged - 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 diff --git a/apps/web/components/__init__.py b/apps/web/components/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/apps/web/components/drawio_editor/__init__.py b/apps/web/components/drawio_editor/__init__.py new file mode 100644 index 0000000..9bae02f --- /dev/null +++ b/apps/web/components/drawio_editor/__init__.py @@ -0,0 +1,26 @@ +"""Streamlit custom component that embeds the diagrams.net editor. + +The component returns a dict {action: "save"|"close", xml: str} +when the user saves or cancels, and None while the editor is idle. +""" +from __future__ import annotations +import os +import streamlit.components.v1 as components + +_FRONTEND = os.path.join(os.path.dirname(__file__), "frontend") +_component_func = components.declare_component("drawio_editor", path=_FRONTEND) + + +def drawio_editor_component(xml: str = "", key: str | None = None) -> dict | None: + """Render the embedded draw.io editor. + + Args: + xml: Optional draw.io XML to pre-load. Pass "" for a blank canvas. + key: Streamlit widget key for state isolation. + + Returns: + None while the editor is open and the user hasn't acted. + {"action": "save", "xml": ""} on save. + {"action": "close", "xml": ""} on cancel. + """ + return _component_func(xml=xml, key=key, default=None) diff --git a/apps/web/components/drawio_editor/frontend/index.html b/apps/web/components/drawio_editor/frontend/index.html new file mode 100644 index 0000000..0f4c3a9 --- /dev/null +++ b/apps/web/components/drawio_editor/frontend/index.html @@ -0,0 +1,174 @@ + + + + + + draw.io Editor + + + +
+ + + Loading editor… +
+ + + + + diff --git a/apps/web/main.py b/apps/web/main.py index fbd5f49..70c3678 100644 --- a/apps/web/main.py +++ b/apps/web/main.py @@ -87,6 +87,8 @@ get_threat_model_mistral, json_to_markdown, ) +from components.drawio_editor import drawio_editor_component +from stride_gpt.core.drawio_parser import drawio_to_prompt_section # ------------------ Helper Functions ------------------ # @@ -1046,6 +1048,36 @@ def on_model_selection_change(): ) st.markdown("""---""") + # Editor renders full-width here (above the form columns) so it is + # immediately visible after the button in col1 is clicked. + if st.session_state.get("show_drawio_editor"): + _hdr_l, _hdr_r = st.columns([1, 3]) + with _hdr_l: + if st.button("Close Diagram Editor", key="close_drawio", use_container_width=True): + st.session_state["show_drawio_editor"] = False + st.rerun() + with _hdr_r: + if st.session_state.get("drawio_xml"): + st.success("Diagram saved – will be used for threat model generation.") + _drawio_result = drawio_editor_component( + xml=st.session_state.get("drawio_xml", ""), + key="drawio_editor_widget", + ) + if _drawio_result: + if _drawio_result.get("action") == "save" and _drawio_result.get("xml"): + _xml = _drawio_result["xml"] + st.session_state["drawio_xml"] = _xml + _summary = drawio_to_prompt_section(_xml) + if _summary: + st.session_state["app_input"] = _summary + st.session_state["_sync_app_desc"] = True + st.session_state["show_drawio_editor"] = False + st.rerun() + elif _drawio_result.get("action") == "close": + st.session_state["show_drawio_editor"] = False + st.rerun() + st.markdown("---") + # Two column layout for the main app content col1, col2 = st.columns([1, 1]) @@ -1138,6 +1170,16 @@ def encode_image(uploaded_file): except Exception as e: st.error(f"An error occurred while analyzing the image: {e!s}") + st.markdown("---") + if st.button("Open Diagram Editor", key="toggle_drawio", use_container_width=True): + st.session_state["show_drawio_editor"] = True + st.rerun() + if st.session_state.get("drawio_xml"): + st.success("Diagram saved – will be used for threat model generation.") + if st.button("Clear diagram", key="clear_drawio"): + st.session_state.pop("drawio_xml", None) + st.rerun() + # Use the get_input() function to get the application description and GitHub URL app_input = get_input() # Update session state only if the text area content has changed @@ -1224,6 +1266,15 @@ 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 + # If the user created a draw.io diagram, prepend extracted architecture + # context (components, data flows, trust boundaries) to the app description + # so the LLM has richer structural input for threat model generation. + _effective_app_input = app_input + if st.session_state.get("drawio_xml"): + _drawio_ctx = drawio_to_prompt_section(st.session_state["drawio_xml"]) + if _drawio_ctx: + _effective_app_input = _drawio_ctx + "\n\n" + app_input + # 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 @@ -1233,7 +1284,7 @@ def encode_image(uploaded_file): authentication, internet_facing, sensitive_data, - app_input, + _effective_app_input, confirmed_dfd=st.session_state.get("confirmed_dfd"), ) diff --git a/stride_gpt/core/drawio_parser.py b/stride_gpt/core/drawio_parser.py new file mode 100644 index 0000000..3951879 --- /dev/null +++ b/stride_gpt/core/drawio_parser.py @@ -0,0 +1,182 @@ +"""Parse draw.io XML (mxGraphModel format) into structured text for LLM prompts. + +Only handles uncompressed XML — which is what the diagrams.net embed API exports. +Compressed .drawio files (base64+deflate) are not supported; the embed component +always returns raw XML so this is fine in practice. +""" +from __future__ import annotations +import xml.etree.ElementTree as ET + + +def parse_drawio_xml(xml: str) -> dict: + """Extract components, connections, and trust boundaries from draw.io XML. + + Returns: + { + "components": [{"id": str, "label": str, "group": str|None}], + "connections": [{"label": str, "from": str, "to": str}], + "trust_boundaries": [{"name": str, "members": [str]}], + } + All lists are empty on parse error. + """ + try: + root = ET.fromstring(xml.strip()) + except ET.ParseError: + return {"components": [], "connections": [], "trust_boundaries": []} + + # Support both bare and ... wrappers. + graph_model = root if root.tag == "mxGraphModel" else root.find(".//mxGraphModel") + if graph_model is None: + return {"components": [], "connections": [], "trust_boundaries": []} + + cells = graph_model.findall(".//mxCell") + + # Build id → (label, parent, style) lookup + id_label: dict[str, str] = {} + id_parent: dict[str, str] = {} + id_style: dict[str, str] = {} + for c in cells: + cid = c.get("id", "") + id_label[cid] = (c.get("value") or "").strip() + id_parent[cid] = c.get("parent", "") + id_style[cid] = (c.get("style") or "").lower() + + # Identify group cells (swimlane or any vertex that acts as parent). + # pontytail: O(n) scan; fine for any diagram a human would draw. + parent_count: dict[str, int] = {} + for c in cells: + p = c.get("parent", "") + if p not in ("0", "1"): + parent_count[p] = parent_count.get(p, 0) + 1 + + group_ids: set[str] = { + cid + for cid, style in id_style.items() + if "swimlane" in style or cid in parent_count + } + + # Components: labeled vertices that are not root/default-layer cells. + components = [] + for c in cells: + cid = c.get("id", "") + if cid in ("0", "1") or c.get("vertex") != "1": + continue + label = id_label.get(cid, "") + if not label: + continue + parent = id_parent.get(cid, "") + group_label = id_label.get(parent, "") if parent not in ("0", "1") else "" + components.append({ + "id": cid, + "label": label, + "group": group_label or None, + }) + + # Connections: edge cells with source+target. + connections = [] + for c in cells: + if c.get("edge") != "1": + continue + src_id = c.get("source", "") + tgt_id = c.get("target", "") + src = id_label.get(src_id, src_id) or src_id + tgt = id_label.get(tgt_id, tgt_id) or tgt_id + if src and tgt: + connections.append({ + "label": (c.get("value") or "").strip(), + "from": src, + "to": tgt, + }) + + # Trust boundaries: labeled group cells with their direct vertex children. + trust_boundaries = [] + for gid in group_ids: + name = id_label.get(gid, "").strip() + if not name: + continue + members = [ + id_label[c.get("id", "")] + for c in cells + if c.get("parent") == gid + and c.get("vertex") == "1" + and id_label.get(c.get("id", ""), "").strip() + ] + trust_boundaries.append({"name": name, "members": members}) + + return { + "components": components, + "connections": connections, + "trust_boundaries": trust_boundaries, + } + + +def drawio_to_prompt_section(xml: str) -> str: + """Convert draw.io XML to a structured prompt section for threat modelling. + + Returns an empty string if the XML has no useful content (parse error, + empty diagram, etc.) so callers can safely check truthiness. + """ + parsed = parse_drawio_xml(xml) + if not parsed["components"] and not parsed["connections"]: + return "" + + lines = [ + "ARCHITECTURE DIAGRAM (from draw.io):", + "The user created the following architecture diagram. Use it as the", + "authoritative model of components and data flows when identifying threats.", + "", + ] + + if parsed["components"]: + lines.append("Components:") + for comp in parsed["components"]: + entry = f" - {comp['label']}" + if comp["group"]: + entry += f" [zone: {comp['group']}]" + lines.append(entry) + lines.append("") + + if parsed["connections"]: + lines.append("Data Flows:") + for conn in parsed["connections"]: + arrow = f" {conn['from']} -> {conn['to']}" + if conn["label"]: + arrow += f" ({conn['label']})" + lines.append(arrow) + lines.append("") + + if parsed["trust_boundaries"]: + lines.append("Trust Boundaries / Security Zones:") + for tb in parsed["trust_boundaries"]: + members = ", ".join(tb["members"]) if tb["members"] else "—" + lines.append(f" - {tb['name']}: {members}") + lines.append("") + + return "\n".join(lines) + + +if __name__ == "__main__": + _SAMPLE = """ + + + + + + + """ + + _r = parse_drawio_xml(_SAMPLE) + assert len(_r["components"]) == 3, f"Expected 3 components, got {len(_r['components'])}" + assert len(_r["connections"]) == 2, f"Expected 2 connections, got {len(_r['connections'])}" + + _section = drawio_to_prompt_section(_SAMPLE) + assert "User" in _section + assert "HTTPS" in _section + assert "->" in _section + + # Empty/invalid XML returns empty string without crashing. + assert drawio_to_prompt_section("") == "" + assert drawio_to_prompt_section("