Skip to content

Back to the catalog

AI-Powered Threat Modeling
github.comrepo

AI-Powered Threat Modeling

AITM is an open-source Python CLI (MIT licensed) that automates STRIDE threat modeling of system architectures using LLMs via Bedrock, OpenAI, or local Ollama. Given a text description or architecture diagram image, it identifies assets and data flows, generates STRIDE-categorized threats with CWE mapping (e.g. CWE-89 for SQLi), scores risk on a 5-level matrix, and outputs actionable mitigations in Markdown or JSON. The engine is stateless with a Pydantic data contract, designed to later support a REST API and frontend; an enterprise edition with additional features is offered commercially.

TL;DR Open-source LLM-powered CLI that auto-generates STRIDE threat models with CWE mappings and mitigations from architecture descriptions or diagrams.
Open original ↗
Document

What it is

AITM is an MIT-licensed Python CLI that sends a system description, or a PNG/JPEG architecture diagram, to an LLM and writes a STRIDE threat model. The summary covers what it produces. This note covers how to run it and what the page leaves unclear.

What to do

Install from source. The page shows no PyPI package.

git clone https://github.com/aisecuritytools-ai/ai-powered-threat-modeling.git
cd ai-powered-threat-modeling
pip install -e .

For a local, no-account trial, the page uses Ollama with llama3.1:

brew install ollama && ollama pull llama3.1
aitm analyze --provider ollama --description-file examples/sample-description.txt

Diagram input is shown with Bedrock:

aitm analyze --provider bedrock --image architecture.png -d "Our payment platform"

Options worth knowing:

  • -p selects bedrock, openai or ollama. The default is bedrock, so a bare aitm analyze needs AWS access.
  • -m sets the model ID. It is auto-detected if omitted.
  • --format takes markdown (default) or json. -o sets the output path.
  • -r sets reasoning depth from 0 (off) to 3 (deep). The page does not say how this maps to each provider or what it costs.
  • -d and -f take the description as text or from a file.

Output

The sample run used Bedrock with claude-sonnet-4. It reported 18 threats, 11 assets and 14 data flows. Each threat lists its STRIDE category, target component, impact, likelihood, CWE ID and a list of mitigations. The sample SQL injection threat (CWE-89, Order Service) is conditional: "if the Order Service constructs SQL queries using string concatenation". The model infers this from the description and does not inspect any code. Example reports are in examples/output/.

Caveats

  • Output quality depends on the model and on how detailed the description is. The page gives no accuracy figures or comparison between providers.
  • The roadmap lists PDF export, an agentic loop with gap analysis, CVE enrichment via web search, a FastAPI REST API, a React frontend and Terraform. The fetched text does not show which of these exist. Treat the CLI with Markdown/JSON export as the shipped part.
  • The page claims "no security expertise required". Review of the generated threats and mitigations by a person is still sensible.
  • The enterprise edition is advertised as having "16+ advanced features". The fetched text does not name them. Contact is by email, with on-premise and volume licensing mentioned.
  • Cloud providers receive your architecture description or diagram. Ollama keeps it local.

Organisations named here

  • CWE-89

Technologies named here

  • AITM
  • Order Service
  • PostgreSQL
  • Ollama
  • nginx
  • AWS Bedrock
  • OpenAI
  • Kubernetes
  • Redis
  • Python

Techniques named here

  • STRIDE methodology
  • SQL injection

Related in AI Security