Deterministic threat modeling for modern AI systems.
The AI Threat Modeling Assistant helps security teams identify exploitable attack paths, weak controls, compliance gaps, and architecture-specific risks across modern AI applications including:
- Large Language Models (LLMs)
- Retrieval-Augmented Generation (RAG)
- Agentic AI systems
- MCP-integrated architectures
- Multimodal AI
- Classical Machine Learning systems
Unlike LLM-generated threat modeling tools, this engine uses deterministic rules, attack-graph correlation, framework mappings, and architecture-aware logic to produce explainable and repeatable findings.
🌐 Live Demo: https://aimlthreats.com
Most AI security assessments today are:
- Generic checklists
- Spreadsheet-driven reviews
- Non-contextual questionnaires
- LLM-generated threat models that are difficult to audit or reproduce
Modern AI systems expose fundamentally different attack surfaces:
- Prompt injection
- Retrieval poisoning
- Excessive agency
- Model manipulation
- Tool abuse
- Insecure MCP integrations
- Unsafe downstream execution
- Hallucination-driven impact
- Training-data poisoning
- Model extraction
This project was built to provide:
- Deterministic reasoning
- Explainable findings
- Architecture-aware assessments
- Attack-chain visibility
- Actionable remediation guidance
Findings are generated using:
- Rule-based analysis
- Control correlation
- Architecture-aware logic
- Graph-based chain discovery
- Framework mappings
No AI-generated hallucinated findings.
Outputs are:
- Reproducible
- Explainable
- Audit-friendly
- Stable across assessments
The questionnaire dynamically changes based on selected AI architecture.
Different AI systems expose different attack surfaces.
The engine automatically enables:
- LLM-specific controls
- RAG retrieval security
- Agentic autonomy checks
- MCP trust boundaries
- Multimodal risks
- Classical ML attack vectors
This reduces irrelevant questions while improving precision.
The engine correlates weak controls into deterministic multi-step attack paths.
Instead of isolated findings, the engine identifies:
- Attacker progression
- Trust-boundary traversal
- Downstream exploitation paths
- Exploit dependencies
- Chained control failures
Each chain explains:
- Why the chain fired
- Required weak controls
- Impacted assets
- Attack progression
- Mitigation steps required to break the chain
External attacker
→ Submit crafted prompt
→ Prompt injection succeeds
→ Harmful output generated
→ Payload reaches downstream system
→ System manipulation achieved
The engine includes graph-based attack-path analysis.
It dynamically traverses attack edges where:
- Trust boundaries exist
- Controls are weak or missing
- Attacker reachability is possible
This enables:
- Chained exploitation visibility
- Downstream impact discovery
- Lateral movement analysis
- Trust-boundary modeling
Findings are ranked using:
- Likelihood
- Impact
- Exposure
- Architecture context
- Attack-chain amplification
- Control gaps
The goal is to surface:
The risks a senior security lead would escalate first.
Every finding includes:
- Reason for trigger
- Attack path
- Abuse case
- Mitigation guidance
- Quick-win remediation
- Confidence scoring
- NIST impact mapping
- Testing guidance
| Architecture | Supported |
|---|---|
| Large Language Models (LLMs) | ✅ |
| RAG / Vector DB Systems | ✅ |
| Agentic AI | ✅ |
| MCP Integrations | ✅ |
| Multimodal AI | ✅ |
| Generative AI | ✅ |
| Traditional ML | ✅ |
| NLP (non-LLM) | ✅ |
| Computer Vision | ✅ |
| Recommendation Systems | ✅ |
| Fraud / Anomaly Detection | ✅ |
- OWASP LLM Top 10
- OWASP ML Top 10
- OWASP Agentic AI
- OWASP MCP Top 10
- MITRE ATLAS technique mappings
- NIST AI RMF
- NIST CIA impact categories
- GDPR
- HIPAA
- PCI DSS
- SOC 2
- ISO 27001
- EU AI Act
The engine can identify findings such as:
- Prompt Injection
- Unsafe Tool Usage
- Excessive Agency
- System Prompt Leakage
- RAG Poisoning
- Retrieval Access Weaknesses
- Harmful Content Generation
- Output Injection
- Unsafe Downstream Execution
- Model Extraction
- Membership Inference
- Data Poisoning
- Drift Exploitation
- MCP Trust Boundary Violations
- Multi-Agent Escalation Risks
External attacker
→ Crafted prompt submission
→ Prompt sanitization bypass
→ Harmful content generated
→ Unsafe downstream rendering
→ System manipulation
Attacker-controlled document
→ Retrieval poisoning
→ Malicious context insertion
→ LLM manipulation
→ Unsafe response generation
Prompt manipulation
→ Tool invocation abuse
→ Credential misuse
→ External API interaction
→ Sensitive data exposure
- Application Security teams
- AI security engineers
- Security architects
- AI governance teams
- Enterprise risk teams
- Red teams
- AI platform teams
- Security researchers
Many AI threat-modeling tools rely entirely on LLMs to generate findings.
This engine intentionally avoids that approach.
Instead, findings are produced using:
- Deterministic rules
- Attack graphs
- Architecture-aware conditions
- Control correlation
- Framework mappings
Benefits:
- Reproducible outputs
- Explainable findings
- Reduced hallucination risk
- Audit-friendly assessments
- Stable severity scoring
- Easier governance adoption
- Python
- Streamlit
- Rule-based analysis engine
- Deterministic scoring engine
- Graph-based chain traversal
- MITRE / OWASP mapping logic
https://github.com/purplesectools/ai-threat-model-assistant.git
cd ai-threat-model-assistantpython3 -m venv venv
source venv/bin/activatepip install -r requirements.txtstreamlit run main.pyai-threat-model-assistant/
│
├── docs/
│ └── images/
├── main.py
├── rules_engine.py
├── requirements.txt
├── README.md
├── Procfile
└── threat_model_coverage.png
Planned improvements include:
- Expanded attack-chain library
- More graph traversal logic
- AI risk scoring improvements
- Exportable reports (PDF/JSON)
- Threat-model versioning
- Saved assessments
- Additional compliance mappings
- More MCP attack scenarios
- Expanded agentic AI coverage
- API support
- CI/CD integrations
This tool is intended to assist threat modeling and security reviews.
It does not replace:
- Penetration testing
- Architecture reviews
- Secure SDLC practices
- Formal compliance assessments
- Human security expertise
Contributions, ideas, rule improvements, attack-chain suggestions, and framework mappings are welcome. Please open:
- Issues
- Feature requests
- Pull requests
MIT License
Developed by Purple Security
GitHub: https://github.com/purplesectools/ai-threat-model-assistant