Skip to main content

AI-powered threat modeling that turns architecture diagrams into actionable risks

Project description

Threat Thinker

AI-powered threat modeling that turns architecture diagrams and business context into actionable threats.

English | 日本語

Public Demo: https://threat-thinker.melonattacker.com

[!IMPORTANT] This is a public demo environment. Please do not upload sensitive or confidential architecture diagrams. For sensitive use cases, use the local CLI or Web UI.

[!IMPORTANT] AI can make mistakes. Do not trust Threat Thinker's output as-is; review the results and judge their correctness before using them.

threat-thinker-logo

What is Threat Thinker?

Threat Thinker is an open-source tool that turns system descriptions, architecture diagrams, and business context into threat models automatically. Provide a natural-language system description or a DFD/architecture diagram as the system shape, add Business Context for scope and assumptions, and optionally use RAG to bring in supporting standards or internal guidance.

Key Features:

  • Description-to-DFD: Generates an intermediate Graph IR DFD from a natural-language system description when no diagram is available.
  • Diagram coverage: Ingests Mermaid, draw.io, Threat Dragon JSON, native Graph IR JSON, and images.
  • Business Context: Injects scope, actors, assets, assumptions, and constraints from PDF, Markdown, or text files.
  • Attribute inference: Uses LLMs to enrich components, data flows, and trust boundaries.
  • RAG boost: Strengthens threat reasoning with retrieved local docs/KB snippets (e.g., OWASP/MITRE/internal).
  • Threat Dragon: Imports Threat Dragon diagrams and can export findings back in Threat Dragon format.
  • Reports: Exports Markdown, JSON, and HTML for reviews and automation.

Key Features

Diagram-to-threat reasoning

  • Provide --description when you do not have a diagram, or drop in a diagram via CLI (--diagram or format-specific flags) or Web UI.
  • Supports Mermaid, draw.io, Threat Dragon JSON, native Graph IR JSON, and image-based diagrams.
  • Deterministic parsing plus LLM reasoning fills missing labels, trust boundaries, and protocols.
  • Outputs prioritized threats with short rationales and OWASP ASVS/CWE references for quick review.

diagram-to-threats
Input diagram and get prioritized threats automatically

Business Context as first-class input

  • Use --description for the system description that can generate a DFD when no diagram is provided.
  • Use --context to add required business context that is not visible in the DFD or architecture diagram.
  • Include scope, actors, sensitive assets, workflows, regulatory assumptions, availability needs, and audit expectations.
  • Threat Thinker injects the full extracted text from PDF, Markdown, or text files into the threat prompt.
  • Combine Business Context with RAG when you also want supporting references retrieved from larger KBs.

Local RAG to boost accuracy

  • Build on-disk knowledge bases from PDFs/Markdown/HTML with threat-thinker kb build under ~/.threat-thinker/kb/<name>.
  • Enable --rag in CLI or the “Use Knowledge Base” toggle in Web UI to retrieve relevant chunks from security guidelines and your org's guidance.
  • Retrieval stays local; only the final prompts go to your chosen LLM provider.
  • Tune top-k per run and swap KBs per project to balance depth, speed, and relevance.

knowledge-base
Build local knowledge bases and use them to strengthen threat reasoning

Threat Dragon round-trip

  • Import Threat Dragon v2 JSON with --threat-dragon, preserving layout and cell metadata.
  • Export a Threat Dragon-compatible JSON that embeds detected threats without regenerating positions.
  • Re-open the exported JSON in Threat Dragon to review or adjust cells with the added findings.
  • Markdown/JSON/HTML reports stay available alongside Threat Dragon output for broader sharing.

threat-dragon-output
Import and export Threat Dragon diagrams with embedded threat findings

Getting Started

Set Up API Keys

Threat Thinker uses LLM for extracting diagrams from images, extracting components, data flows, and trust boundaries from architecture diagrams, and for inferring threats. Threat Thinker supports OpenAI, Anthropic Claude, AWS Bedrock (Claude v3+ models), and local Ollama APIs (for text-only flows).

You must set at least one of the following environment variables before use:

# For OpenAI API (e.g., gpt-4.1)
export OPENAI_API_KEY=...

# For Claude API (e.g., claude-sonnet-4-5)
export ANTHROPIC_API_KEY=...

# For Bedrock API (e.g., anthropic.claude-sonnet-4-5-20250929-v1:0)
# Option 1: Use AWS Profile (recommended)
aws configure --profile my-profile
# Then use --aws-profile my-profile in the command

# Option 2: Use environment variables
export AWS_ACCESS_KEY_ID=...
export AWS_SECRET_ACCESS_KEY=...
export AWS_SESSION_TOKEN=...

Local Ollama (no API key)

  • Start Ollama locally (default host http://localhost:11434) and pull a model (e.g., ollama pull llama3.1).
  • Run Threat Thinker with --llm-api ollama --llm-model <model> [--ollama-host http://localhost:11434] for Mermaid/Draw.io/Threat Dragon inputs.
  • Image extraction is not supported with the Ollama backend; use text-based diagram inputs instead.

Installation

Choose one of the following methods:

Using pipx

pipx install threat-thinker

Using uv

uv tool install threat-thinker

Using pip

pip install threat-thinker

# Or install from a GitHub Release wheel
pip install https://github.com/melonattacker/threat-thinker/releases/download/v0.8.0/threat_thinker-0.8.0-py3-none-any.whl

# Or install the latest from main
pip install git+https://github.com/melonattacker/threat-thinker.git

Note: If you see externally-managed-environment error, use pipx or uv instead, or create a virtual environment first.

For development

git clone https://github.com/melonattacker/threat-thinker.git
cd threat-thinker
uv sync --extra dev --frozen

# Fallback if uv is unavailable
python3 -m venv .venv
source .venv/bin/activate
pip install -e .[dev]

Verify installation

threat-thinker version
threat-thinker -v
threat-thinker --help

CLI Usage

Here is an example of command using CLI mode.

# Think: Generate a DFD from a system description, then analyze threats
threat-thinker think \
    --description "Customers use a web app to manage orders. The app stores customer PII in Postgres and sends email through a third-party provider." \
    --topn 5 \
    --llm-api openai \
    --llm-model gpt-4.1 \
    --out-dir reports/

# Think: Analyze a diagram
threat-thinker think \
    --diagram examples/diagrams/web/system.mmd \
    --context examples/diagrams/web/business-context.md \
    --infer-hints \
    --topn 5 \
    --llm-api openai \
    --llm-model gpt-4.1 \
    --out-dir reports/

# Diff: Compare two threat reports and analyze changes
threat-thinker diff \
    --after reports/new-report.json \
    --before reports/old-report.json \
    --llm-api openai \
    --llm-model gpt-4.1 \
    --out-dir reports/ \
    --lang en

# Run threat analysis with local Ollama (text-only diagrams)
threat-thinker think \
    --mermaid examples/diagrams/web/system.mmd \
    --llm-api ollama \
    --llm-model llama3.1 \
    --ollama-host http://localhost:11434 \
    --out-dir reports/

# Serve: Launch API server
threat-thinker serve --config examples/demo-app/serve.example.yaml

# Worker: Start background processor for queued jobs
threat-thinker worker --config examples/demo-app/serve.example.yaml

Web UI

# Launch Web UI
threat-thinker webui

Then visit http://localhost:7860 to use Threat Thinker interactively.

Documentation

  • docs/tutorials.md — Guided runs for web, AWS, and diff scenarios.
  • docs/cli.md — Flag reference and examples for think/diff/kb commands.
  • docs/design.md — Architecture and processing flow across the five layers.
  • docs/rag.md — Building and using local knowledge bases to strengthen threat reasoning.
  • docs/reports.md - Report formats and contents for Markdown, JSON, HTML, Threat Dragon and diff outputs.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

threat_thinker-0.8.0.tar.gz (149.2 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

threat_thinker-0.8.0-py3-none-any.whl (135.5 kB view details)

Uploaded Python 3

File details

Details for the file threat_thinker-0.8.0.tar.gz.

File metadata

  • Download URL: threat_thinker-0.8.0.tar.gz
  • Upload date:
  • Size: 149.2 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.7

File hashes

Hashes for threat_thinker-0.8.0.tar.gz
Algorithm Hash digest
SHA256 008eeb26f1b5eb4f72d1ecc58221e6048d640c3b788d42b2c2bdd81d435a2d08
MD5 cde352397765daaa775b55985ec36e6a
BLAKE2b-256 d9b21226f68b6afd03b152c7e0b6328c09b6da7ed6ebc42bca3c464702e21d8a

See more details on using hashes here.

Provenance

The following attestation bundles were made for threat_thinker-0.8.0.tar.gz:

Publisher: release.yml on melonattacker/threat-thinker

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file threat_thinker-0.8.0-py3-none-any.whl.

File metadata

  • Download URL: threat_thinker-0.8.0-py3-none-any.whl
  • Upload date:
  • Size: 135.5 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.7

File hashes

Hashes for threat_thinker-0.8.0-py3-none-any.whl
Algorithm Hash digest
SHA256 d7d5da25b4ad780085150e7f3410b35d638971ccb62c6c830133a8bfff0b3c19
MD5 80cc52e5f721a66d4a5502ec2ceb29f9
BLAKE2b-256 2e801d9e3ad245f4cf14ff17db91706bcdf0066418e303c17070f0f61ffd77c9

See more details on using hashes here.

Provenance

The following attestation bundles were made for threat_thinker-0.8.0-py3-none-any.whl:

Publisher: release.yml on melonattacker/threat-thinker

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page