Skip to content

An intelligent threat model evaluation tool that uses Large Language Models (LLMs) to analyse Threat Dragon models and generate comprehensive security assessment reports.

License

Notifications You must be signed in to change notification settings

InfosecOTB/td-ai-evaluator

Repository files navigation

AI-Powered Threat Model Evaluator

td-ai

An intelligent threat model evaluation tool that uses Large Language Models (LLMs) to analyze Threat Dragon models and generate comprehensive security assessment reports.

Features

  • AI-Powered Analysis: Uses state-of-the-art LLMs to analyze Threat Dragon models and evaluate security posture
  • Comprehensive Reports: Generates detailed markdown reports with threat analysis, risk assessment, and recommendations
  • Multi-LLM Support: Tested on OpenAI, Anthropic, Google, Novita, and xAI. Works with any LiteLLM-compatible provider
  • Threat Dragon Integration: Works seamlessly with Threat Dragon JSON models
  • Multiple Output Formats: Generates reports in Markdown, DOCX, and PDF formats
  • Batch Processing: Automatically processes multiple threat models and generates summary reports
  • Structured Output: Clean, professional reports ready for documentation and review

Quick Start

Prerequisites

  • Python 3.8+
  • API key for your chosen LLM provider

Installation

  1. Clone the repository

    git clone <repository-url>
    cd td-ai-evaluator
  2. Install dependencies

    pip install -r requirements.txt
  3. Configure environment

    cp env.example .env

    Edit .env with your configuration:

    LLM_MODEL_NAME=openai/gpt-5
    OPENAI_API_KEY=your_openai_api_key_here
    
    INPUT_THREAT_SCHEMA_JSON=owasp.threat-dragon.schema.V2.json
  4. Prepare input files

    • The Threat Dragon schema file is already in ./input/schema/
    • Place your threat model JSON files in ./input/
  5. Run the application

    python src/main.py
  6. Check results

    • Markdown reports will be in ./output/markdown/
    • DOCX reports will be in ./output/docx/
    • PDF reports will be in ./output/

Configuration

Environment Variables

Variable Description Example
LLM_MODEL_NAME LLM model identifier openai/gpt-5
INPUT_THREAT_SCHEMA_JSON Threat Dragon schema filename owasp.threat-dragon.schema.V2.json

Advanced Configuration

You can adjust LLM parameters in src/ai_client.py:

  • temperature: Controls randomness (0.0=deterministic, 1.0=creative). Default: 0.1 (both individual and combined reports)
  • max_tokens: Maximum tokens in response. Default: 48000 (individual reports), 30000 (combined reports)
  • timeout: Request timeout in seconds. Default: 14400 (4 hours)

Project Structure

td-ai-evaluator/
├── src/
│   ├── main.py              # Main application entry point
│   ├── ai_client.py         # LLM integration and report generation
│   ├── utils.py             # File operations and format conversions
│   └── config.py            # Configuration and path constants
├── input/                   # Input files directory
│   ├── schema/
│   │   └── owasp.threat-dragon.schema.V2.json
│   └── *.json              # Threat model files
├── output/                  # Generated output directory
│   ├── markdown/           # Markdown reports
│   ├── docx/               # DOCX reports
│   └── logs/               # Application logs
├── prompt-report.txt        # AI threat model evaluation prompt
├── prompt-combined.txt      # AI combined report generation prompt
├── env.example              # Environment configuration template
├── requirements.txt         # Python dependencies
└── README.md               # This file

How It Works

  1. Input Processing: Loads Threat Dragon schema and discovers all threat model JSON files
  2. Automatic Mode Detection:
    • Single File: Generates individual threat model report
    • Multiple Files: Generates comprehensive combined report comparing all models
  3. AI Analysis: Uses LLM to analyze threat model(s) and evaluate security posture
  4. Report Generation: Creates comprehensive markdown reports with threat analysis and recommendations
  5. Format Conversion: Converts markdown reports to DOCX and PDF formats
  6. Output Organization: Saves all reports in organized output directories

Report Features

The tool generates comprehensive security assessment reports with:

Report Contents

  • Threat Analysis: Detailed evaluation of identified threats and vulnerabilities
  • Risk Assessment: Risk levels and impact analysis for each threat
  • Recommendations: Actionable security recommendations and best practices
  • Model Overview: Summary of system components and architecture
  • Security Posture: Overall security assessment and maturity evaluation

Output Formats

  • Markdown: Clean, readable format for documentation and version control
  • DOCX: Professional Word documents with proper formatting
  • PDF: Final reports ready for distribution and review

Combined Report (Multiple Models)

  • Comparative Analysis: Side-by-side comparison of all threat models
  • Threats & Mitigations Ranking: Ranks models by their security maturity
  • Overall Model Maturity: Evaluates shared DFD architecture across models
  • Individual Model Evaluations: Detailed analysis of each model's threats and mitigations
  • Executive Summary: High-level overview suitable for stakeholders

Troubleshooting

Common Issues

LLM Response Errors

  • Timeout Issues: Increase timeout value in ai_client.py for large models
  • Token Limits: Adjust max_tokens based on model capabilities
  • Empty Responses: Check LLM API key and model availability

Configuration Issues

  • API Key Errors: Ensure correct environment variables are set in .env
  • Model Not Found: Verify model name format matches provider requirements
  • File Not Found: Ensure threat model JSON files are in ./input/ directory

Report Generation Issues

  • Missing Reports: Check ./output/logs/ for error messages
  • Format Issues: Ensure Microsoft Word is installed for PDF conversion
  • Empty Reports: Verify threat model files are valid JSON

Performance Tips

  • Use faster models for quick assessments
  • Process multiple models in batch for efficiency
  • Monitor token usage to control costs

Development

Running the Application

# Install dependencies
pip install -r requirements.txt

# Run the application
python src/main.py

Code Structure

  • main.py: Orchestrates report generation and batch processing
  • ai_client.py: Handles LLM communication and report generation
  • utils.py: File operations and format conversions (Markdown, DOCX, PDF)
  • config.py: Centralized configuration and path management

Customization

Modifying the AI Prompts

Edit prompt-report.txt and prompt-combined.txt to customize:

  • Report structure and format
  • Analysis depth and focus areas
  • Security frameworks and methodologies

Adding New LLM Providers

  1. Add provider configuration to env.example
  2. Update provider table in README
  3. Test with sample threat model

License

This project is licensed under the Apache 2.0 License - see the LICENSE file for details.

Acknowledgments

Additional Resources

For more information about cybersecurity and AI projects, visit my blog at https://infosecotb.com.


Built for security professionals and threat modeling evaluators

About

An intelligent threat model evaluation tool that uses Large Language Models (LLMs) to analyse Threat Dragon models and generate comprehensive security assessment reports.

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages