AI-Powered Threat Modelling Extension

[jgadsden]

It would be good to build in a generative AI engine into Threat Dragon that suggests threats for a given data flow diagram model.
This needs a lot of discussion on how to build it in and the consequences for the security posture of the threat models themselves.

Piotr Kowalczyk has built a Python tool that uses AI to analyses and modify Threat Dragon model - it works outside of the application and directly updates the .json file by generating and adding threats to DFD diagrams.
The article describing the tool and its application to the Threat Dragon files is here:

AI-Powered Threat Modeling with OWASP Threat Dragon – Part 2 (along with Part 1)

The TD AI Modeler tool is open source and Piotr has provided example threats generated for the threat models

We are often asked for help in generating threats, and we have never managed to create an engine that would satisfy this demand.
The AI generated threats may be the solution to this gap in Threat Dragon.

There may be logistics to work through, and it would be great if it was a 'magic wand' provided within Threat Dragon itself, so a lot of discussion involved.

[InfosecOTB]

I would like to suggest a potential solution for the “magic wand”. I’ve just adjusted the code in a way that (I believe) may make it easier to integrate with Threat Dragon. You can find it here:

td-ai-modeler-plugin

I think it might make creating the “magic wand” inside Threat Dragon a bit easier.

How I imagine this might work:

On the top menu, there could be an additional option called AI Tools (or something similar).
Inside AI Tools, there would be two options (and maybe more in the future):

• Generate Threats & Mitigations
• Settings

The Generate Threats & Mitigations option would trigger the Python application with the arguments configured in Settings.

Example command:

python src/main.py --llm-model gemini/gemini-2.5-pro --model-file "C:\temp\plugin\payments-processing-platform.json" --temperature 0.3 --response-format

The API key would be passed as an environment variable (so not visible in the command line).
This will generate threats for the JSON file specified in the --model-file argument.

While running, there should be a processing window showing progress and results.

Right now, I’ve been generating threats while Threat Dragon was open, but I had to close it and reload the JSON file to see the results.

"Settings" would have the following options:

API_KEY – masked and saved securely within Threat Dragon. Passed as an environment variable and loaded by the Python app with:

api_key = os.getenv('API_KEY')

The rest of the arguments would be set in Settings and passed using the Argument Parser function. Two are mandatory: (--llm-model, --model-file) and rest optional (depending of LLM model and provider).

I believe an instruction for Settings would need to be provided inside the Threat Dragon, as Arguments depend of selected AI.

Available Arguments

Argument	Required	Default	Description
--llm-model	Yes	-	LLM model identifier (e.g., "openai/gpt-5", "anthropic/claude-sonnet-4-5-20250929")
--model-file	Yes	-	Input threat model JSON file path (full path including filename)
--temperature	No	0.1	LLM temperature for randomness (range: 0-2)
--response-format	No	False	Enable structured JSON response format
--api-base	No	None	Custom API base URL
--log-level	No	INFO	Logging level (INFO or DEBUG)

One interesting thing about this new version of the application is that you can run it multiple times on the same Threat Model. You can switch (or not) between LLM providers and models - the AI will then review the existing threats and mitigations, decide if they should be adjusted, and add new ones if needed.

Maybe this functionality could also be added directly into AI Tools, so it can be used without touching the Settings each time.
What do you think?

[jgadsden]

Hello @InfosecOTB - certainly is a possible. Is this to be tried out for the desktop only, and not yet for the web application version?
It may be that AI tools should be restricted to the desktop version, so that the location of the LLM is certain, for example.

Is this something that you can take forward @InfosecOTB ?
I have a lot that I am working on within Threat Dragon (bug reports, dependency hell) and also preparing for a student engagement in December, so realistically my involvement can only be cheering on from the sideline :)

[InfosecOTB]

Hi @jgadsden - I believe that, following your suggestion, we should initially integrate this with the desktop version, as it seems easier in many aspects and we need to start somewhere. However, once it’s working well, I think we should try to address all concerns and eventually do the same for the web version.
I’m happy to take this forward with the community’s help.

[jgadsden]

Thanks for taking this on @InfosecOTB , we can certainly support you to get this feature into Threat Dragon
It would be the most significant feature added to Threat Dragon since @lreading migrated Threat Dragon to Vue + antV/x6 with all the good stuff that came from that