Skip to content

castlebbs/llm-threat-modeling

Folders and files

NameName
Last commit message
Last commit date

Latest commit

Β 

History

2 Commits
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 

Repository files navigation

OWASP September 2025 Dublin Talk - David Robert - AI/LLM Security Threat Modeling

A collection of Jupyter notebooks for automated security threat modeling and analysis, with a focus on AI/LLM systems. This toolkit leverages Neo4j graph databases and Claude Sonnet Foundation Model to automated threat identification in software architectures.

Project Goals

This repository provides automated tools to:

  1. Analyze Data Flow Diagrams - Extract architectural components and relationships from diagrams images
  2. Model System Architecture - Model application architectures in a Neo4j graph database
  3. Generate Threat Libraries - Create Threat libraries represented by cypher queries for automated threat detection
  4. Execute Threat Detection - Run threat queries against system models to identify threats
  5. Identify Assets - Extract and classify data assets according to protection requirements
  6. Generate Abuse Cases - Derive security abuse cases from functional user stories
  7. Document Architecture - Generate architecture documentation artefacts from codebases

⚠️ Disclaimer:

This codebase was created for the OWASP September 2025 Dublin Talk as a demonstration and prototyping toolkit.
It is not production-ready and should not be used in real production environments.
Use at your own risk.

Repository Structure

owasp-jupyter/
β”œβ”€β”€ README.md                          # This file
β”œβ”€β”€ requirements.txt                   # Python dependencies
β”œβ”€β”€ .gitignore                         # Git ignore patterns
β”‚
β”œβ”€β”€ notebooks/                         # Jupyter notebooks (main interface)
β”‚   β”œβ”€β”€ diagram_analysis.ipynb        # Analyze architecture diagrams
β”‚   β”œβ”€β”€ threat_analysis.ipynb         # Execute threat detection queries
β”‚   β”œβ”€β”€ threats_library_generator.ipynb   # Generate threat detection queries
β”‚   β”œβ”€β”€ abuse_cases_and_assets.ipynb  # Asset identification & abuse cases
β”‚   └── architecture_document_generator.ipynb  # Generate architecture docs
β”‚
β”œβ”€β”€ utils/                             # Shared utility modules
β”‚   β”œβ”€β”€ common_utils.py               # Common helper functions
β”‚   β”œβ”€β”€ threat_modeling_utils.py      # Neo4j & diagram analysis utilities
β”‚   β”œβ”€β”€ cypher_query_utils.py         # Cypher query loading & execution
β”‚   β”œβ”€β”€ llm_analysis_utils.py         # LLM-powered threat analysis
β”‚   β”œβ”€β”€ abuse_case_utils.py           # Abuse case & asset identification
β”‚   └── code_analysis.py              # Code analysis utilities
β”‚
β”œβ”€β”€ schema/                            # Neo4j schema definitions
β”‚   └── mapping.txt                   # Node/relationship type mappings
β”‚
β”œβ”€β”€ threat_library/                   # Generated Cypher query libraries
β”‚   └── *.cypher                      # Threat detection query files
β”‚
└── data_classification.md            # Data classification example

Architecture Overview

Core Components

1. Diagram Analysis Pipeline

  • Input: Data flow diagrams (images)
  • Process: LLM vision model extracts components and relationships
  • Output: Cypher statements to populate Neo4j graph
  • Use Case: Rapidly model existing system architectures

2. Threat Library Generator

  • Input: Deep research agents output. For the demo we used OWASP LLM vulnerability documentation.
  • Process: LLM generates detection queries from threat descriptions
  • Output: Categorized Cypher query libraries
  • Use Case: Keep threat detection current with emerging vulnerabilities

3. Threat Detection Engine

  • Input: System architecture graph in Neo4j
  • Process: Execute threat libraries represented by Cypher queries
  • Output: Identified threats with severity ratings
  • Use Case: Automated threat identification

4. Asset & Abuse Case Analysis

  • Input: Application descriptions and user stories
  • Process: LLM extracts assets and generates abuse cases
  • Output: Classified assets and security abuse cases
  • Use Case: Can be used for further analysis for business logic threats for example

5. Architecture Documents generators

  • Input: Source repo
  • Process: Coding Agent following instructions to identify architecture details
  • Output: In this example, list of APIs, list of components
  • Use Case: Generate architecture documents from source code

Technology Stack

  • Python 3.8+ - Core programming language
  • Jupyter Notebooks - Interactive analysis environment
  • Neo4j - Graph database for architecture modeling
  • Claude Sonnet 4 - AI for diagram analysis and threat generation
  • Pandas - Data manipulation and visualization
  • OWASP Top 10 for LLMs - Threat classification framework used as example (instead of deep research agents because of demo constraints)

Quick Install

  1. Clone and Setup

    git clone https://github.com/castlebbs/ai-threat-modeling.git
    cd ai-threat-modeling
    python3 -m venv .venv
    source .venv/bin/activate  # On Windows: .venv\Scripts\activate
    pip install -r requirements.txt
  2. Configure Environment

    export NEO4J_URI="bolt://localhost:7687"
    export NEO4J_USERNAME="neo4j"
    export NEO4J_PASSWORD="your-password"
    export ANTHROPIC_API_KEY="your-anthropic-api-key"
  3. Start Jupyter

    jupyter notebook

For detailed setup instructions including Neo4j installation and troubleshooting, see SETUP.md.

Quick Start

  1. Start Jupyter

    jupyter notebook
  2. Choose Your Workflow

    Option A: Analyze an Existing Architecture Diagram

    • Open notebooks/diagram_analysis.ipynb
    • Upload or provide URL to your architecture diagram
    • Run cells to extract components and populate Neo4j
    • Review the generated graph structure

    Option B: Generate a Threat Library

    • Open notebooks/threats_library_generator.ipynb
    • Specify OWASP LLM vulnerability URL
    • Generate Cypher queries for threat detection
    • Export to threat_library/ directory

    Option C: Run Threat Detection

    • Open notebooks/threat_analysis.ipynb
    • Select a threat query library
    • Execute queries against your system model
    • Review detected threats and severity levels

πŸ“š Notebooks Guide

All notebooks are located in the notebooks/ directory.

1. diagram_analysis.ipynb

Purpose: Convert architecture diagrams to Neo4j graph models

Workflow:

  1. Provide diagram (URL or local file)
  2. Claude analyzes diagram and generates Cypher
  3. Creates nodes (components) and relationships (data flows)
  4. Populates Neo4j database
  5. Verifies data integrity

Key Features:

  • Automatic component normalization
  • Semantic understanding of diagram elements
  • Support for various diagram formats

2. threats_library_generator.ipynb

Purpose: Generate threat detection queries from OWASP documentation

Workflow:

  1. Fetch OWASP LLM vulnerability documentation
  2. Load Neo4j schema mapping
  3. Claude generates STRIDE-categorized queries
  4. Exports to .cypher files
  5. Ready for threat detection execution

Key Features:

  • STRIDE threat categorization
  • Both specific and general detection patterns
  • Schema-compliant query generation

3. threat_analysis.ipynb

Purpose: Execute threat detection against system models

Workflow:

  1. Connect to Neo4j database
  2. Load threat query library
  3. Execute queries by category
  4. Display findings with severity
  5. Generate threat analysis report

Key Features:

  • Organized by STRIDE categories
  • Severity-based highlighting
  • Detailed threat descriptions
  • Markdown report generation

4. abuse_cases_and_assets.ipynb

Purpose: Identify assets and generate abuse cases

Workflow:

  1. Provide application description or user stories
  2. Claude extracts data assets with classification
  3. Claude generates security abuse cases
  4. Display categorized results
  5. Export for security requirements

Key Features:

  • Data classification framework integration
  • Focus on information assets (not infrastructure)
  • Realistic abuse case scenarios

5. architecture_document_generator.ipynb

Purpose: Generate architecture documentation from code

Workflow:

  1. Analyze repository structure
  2. Extract components and patterns
  3. Generate architecture diagrams
  4. Document key design decisions
  5. Export comprehensive documentation

Key Features:

  • Automated component discovery
  • Technology stack identification
  • Relationship mapping

πŸ” Example Workflows

Complete Threat Modeling Workflow

  1. Model Your System

    # In notebooks/diagram_analysis.ipynb
    image_path = "architecture_diagram.png"
    # Run analysis to populate Neo4j
  2. Generate Threat Library

    # In notebooks/threats_library_generator.ipynb
    threat_url = "https://raw.githubusercontent.com/OWASP/www-project-top-10-for-large-language-model-applications/main/2_0_vulns/LLM01_PromptInjection.md"
    # Generate queries
  3. Detect Threats

    # In notebooks/threat_analysis.ipynb
    query_file = "threat_library/llm_threat_queries_20250917_180146.cypher"
    # Execute detection
  4. Review Results

    • Examine threats by STRIDE category
    • Prioritize by severity
    • Generate remediation plan

πŸ› οΈ Utilities Reference

common_utils.py

  • format_results_table() - Display results with formatting
  • display_summary_statistics() - Show detection summaries
  • validate_environment_variables() - Check required config
  • print_section_header(), print_success(), etc. - Formatted output

threat_modeling_utils.py

  • Neo4jConnection - Database connection management
  • ImageProcessor - Image handling for Claude
  • DiagramAnalyzer - Diagram-to-Cypher conversion
  • initialize_connections() - Setup connections
  • execute_dynamic_cypher() - Execute Cypher with validation

cypher_query_utils.py

  • CypherQueryLoader - Load and parse query files
  • ThreatQueryExecutor - Execute threat queries
  • ThreatReportGenerator - Generate analysis reports
  • load_and_execute_queries() - Complete workflow
  • run_complete_threat_detection() - End-to-end detection

llm_analysis_utils.py

  • ThreatDocumentationRetriever - Fetch threat docs
  • SchemaManager - Load Neo4j schema
  • ThreatAnalyzer - Claude AI threat analysis
  • QueryParser - Parse generated queries
  • LLMThreatAnalyzer - Complete workflow orchestration

abuse_case_utils.py

  • AbuseCaseGenerator - Generate abuse cases from user stories
  • AssetIdentifier - Identify and classify data assets

code_analysis.py

  • do_code_analysis() - Execute code analysis with Claude CLI

ecurity Considerations

  • API Keys: Never commit API keys. Use environment variables
  • Neo4j: Use strong passwords and restrict network access
  • Data: Be cautious with sensitive architectural information
  • Generated Queries: Review before executing in production environments

Contributing

Contributions are welcome! Areas for improvement:

  • Additional threat patterns and queries
  • Support for more diagram formats
  • Integration with other security tools
  • Performance optimizations
  • Documentation improvements

+ License

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

πŸ“ž Support

For questions, issues, or feature requests, please open an issue in the GitHub repository.


Note: This toolkit is designed for security research and threat modeling purposes. Always validate findings and follow responsible disclosure practices.

About

OWASP September 2025 Dublin Talk - David Robert - AI/LLM Security Threat Modeling

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

 
 
 

Contributors