mrwadams · Mavi9412 · Jun 30, 2026 · Jun 30, 2026
diff --git 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
diff --git a/apps/web/components/drawio_editor/__init__.py 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": "<mxGraphModel...>"}  on save.
+        {"action": "close", "xml": "<last known 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
@@ -0,0 +1,174 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+  <meta charset="utf-8"/>
+  <meta name="viewport" content="width=device-width, initial-scale=1"/>
+  <title>draw.io Editor</title>
+  <style>
+    * { box-sizing: border-box; margin: 0; padding: 0; }
+    body { font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif; background: #fff; }
+
+    #toolbar {
+      display: flex;
+      align-items: center;
+      gap: 8px;
+      padding: 8px 12px;
+      background: #f0f2f6;
+      border-bottom: 1px solid #d0d3da;
+    }
+
+    .btn {
+      padding: 5px 16px;
+      border-radius: 4px;
+      border: 1px solid #ccc;
+      cursor: pointer;
+      font-size: 13px;
+      font-family: inherit;
+      background: #fff;
+      transition: background 0.1s;
+    }
+    .btn:hover { background: #e8eaf0; }
+
+    .btn-primary {
+      background: #ff4b4b;
+      color: #fff;
+      border-color: #c83030;
+    }
+    .btn-primary:hover { background: #d93030; }
+
+    #status {
+      margin-left: auto;
+      font-size: 12px;
+      color: #6c757d;
+      white-space: nowrap;
+    }
+
+    #drawio-frame {
+      width: 100%;
+      height: 590px;
+      border: none;
+      display: block;
+    }
+  </style>
+</head>
+<body>
+  <div id="toolbar">
+    <button class="btn btn-primary" id="btn-save">Save &amp; Close</button>
+    <button class="btn" id="btn-cancel">Cancel</button>
+    <span id="status">Loading editor…</span>
+  </div>
+  <iframe
+    id="drawio-frame"
+    src="https://embed.diagrams.net/?embed=1&proto=json&spin=1&ui=min&dark=0&noSaveBtn=0&saveAndExit=0"
+    allowfullscreen
+  ></iframe>
+
+  <script>
+    /* ── Streamlit component IPC (raw postMessage, no npm build needed) ──
+       All outgoing messages MUST include isStreamlitMessage: true or
+       Streamlit's frontend will silently ignore them.               ── */
+    var FRAME_HEIGHT = 650;
+
+    function stPost(type, extra) {
+      var msg = Object.assign({ isStreamlitMessage: true, type: "streamlit:" + type }, extra || {});
+      window.parent.postMessage(msg, "*");
+    }
+
+    stPost("componentReady", { apiVersion: 1 });
+    stPost("setFrameHeight", { height: FRAME_HEIGHT });
+
+    /* ── draw.io editor state ── */
+    var drawioFrame = document.getElementById("drawio-frame");
+    var statusEl    = document.getElementById("status");
+
+    var editorReady = false;   // draw.io fired "init"
+    var pendingXml  = null;    // XML to load once the editor is ready
+    var savedXml    = "";      // last XML we know about (auto-saved or explicitly exported)
+    var pendingSave = false;   // we fired an export request and are awaiting the result
+
+    var BLANK_XML = [
+      '<mxGraphModel>',
+      '  <root>',
+      '    <mxCell id="0"/>',
+      '    <mxCell id="1" parent="0"/>',
+      '  </root>',
+      '</mxGraphModel>'
+    ].join("\n");
+
+    function postToDrawio(msg) {
+      drawioFrame.contentWindow.postMessage(JSON.stringify(msg), "*");
+    }
+
+    /* ── Toolbar buttons ── */
+    document.getElementById("btn-save").addEventListener("click", function () {
+      pendingSave = true;
+      statusEl.textContent = "Exporting…";
+      postToDrawio({ action: "export", format: "xml" });
+    });
+
+    document.getElementById("btn-cancel").addEventListener("click", function () {
+      stPost("setComponentValue", { value: { action: "close", xml: savedXml }, dataType: "json" });
+    });
+
+    /* ── Message hub ── */
+    window.addEventListener("message", function (evt) {
+      var data = evt.data;
+
+      /* --- Streamlit → component: render with new props --- */
+      if (data && typeof data === "object" && data.isStreamlitMessage && data.type === "streamlit:render") {
+        var xml = (data.args && data.args.xml) || "";
+        stPost("setFrameHeight", { height: FRAME_HEIGHT });
+
+        if (editorReady) {
+          /* Only reload if we don't have local unsaved edits. */
+          if (!savedXml) {
+            postToDrawio({ action: "load", xml: xml || BLANK_XML });
+          }
+        } else {
+          /* Buffer until the editor fires "init". */
+          if (!pendingXml) pendingXml = xml || null;
+        }
+        return;
+      }
+
+      /* --- draw.io → component: JSON string events --- */
+      if (typeof data !== "string") return;
+      var msg;
+      try { msg = JSON.parse(data); } catch (_) { return; }
+
+      switch (msg.event) {
+        case "init":
+          editorReady = true;
+          statusEl.textContent = "Ready";
+          postToDrawio({ action: "load", xml: pendingXml || BLANK_XML });
+          pendingXml = null;
+          break;
+
+        case "load":
+          statusEl.textContent = "Diagram loaded";
+          break;
+
+        case "save":
+          /* Native draw.io Ctrl+S / save button. */
+          savedXml = msg.xml || savedXml;
+          statusEl.textContent = "Saved";
+          break;
+
+        case "export":
+          if (pendingSave) {
+            pendingSave = false;
+            /* draw.io returns XML in msg.xml for format=xml; msg.data is a fallback. */
+            var xml = msg.xml || msg.data || savedXml;
+            savedXml = xml;
+            stPost("setComponentValue", { value: { action: "save", xml: xml }, dataType: "json" });
+          }
+          break;
+
+        case "close":
+          stPost("setComponentValue", { value: { action: "close", xml: savedXml }, dataType: "json" });
+          break;
+      }
+    });
+  </script>
+</body>
+</html>
diff --git 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"),
         )