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.
This repository provides automated tools to:
- Analyze Data Flow Diagrams - Extract architectural components and relationships from diagrams images
- Model System Architecture - Model application architectures in a Neo4j graph database
- Generate Threat Libraries - Create Threat libraries represented by cypher queries for automated threat detection
- Execute Threat Detection - Run threat queries against system models to identify threats
- Identify Assets - Extract and classify data assets according to protection requirements
- Generate Abuse Cases - Derive security abuse cases from functional user stories
- Document Architecture - Generate architecture documentation artefacts from codebases
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.
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
- 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
- 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
- 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
- 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
- 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
- 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)
-
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
-
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"
-
Start Jupyter
jupyter notebook
For detailed setup instructions including Neo4j installation and troubleshooting, see SETUP.md.
-
Start Jupyter
jupyter notebook
-
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
- Open
All notebooks are located in the notebooks/ directory.
Purpose: Convert architecture diagrams to Neo4j graph models
Workflow:
- Provide diagram (URL or local file)
- Claude analyzes diagram and generates Cypher
- Creates nodes (components) and relationships (data flows)
- Populates Neo4j database
- Verifies data integrity
Key Features:
- Automatic component normalization
- Semantic understanding of diagram elements
- Support for various diagram formats
Purpose: Generate threat detection queries from OWASP documentation
Workflow:
- Fetch OWASP LLM vulnerability documentation
- Load Neo4j schema mapping
- Claude generates STRIDE-categorized queries
- Exports to
.cypherfiles - Ready for threat detection execution
Key Features:
- STRIDE threat categorization
- Both specific and general detection patterns
- Schema-compliant query generation
Purpose: Execute threat detection against system models
Workflow:
- Connect to Neo4j database
- Load threat query library
- Execute queries by category
- Display findings with severity
- Generate threat analysis report
Key Features:
- Organized by STRIDE categories
- Severity-based highlighting
- Detailed threat descriptions
- Markdown report generation
Purpose: Identify assets and generate abuse cases
Workflow:
- Provide application description or user stories
- Claude extracts data assets with classification
- Claude generates security abuse cases
- Display categorized results
- Export for security requirements
Key Features:
- Data classification framework integration
- Focus on information assets (not infrastructure)
- Realistic abuse case scenarios
Purpose: Generate architecture documentation from code
Workflow:
- Analyze repository structure
- Extract components and patterns
- Generate architecture diagrams
- Document key design decisions
- Export comprehensive documentation
Key Features:
- Automated component discovery
- Technology stack identification
- Relationship mapping
-
Model Your System
# In notebooks/diagram_analysis.ipynb image_path = "architecture_diagram.png" # Run analysis to populate Neo4j
-
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
-
Detect Threats
# In notebooks/threat_analysis.ipynb query_file = "threat_library/llm_threat_queries_20250917_180146.cypher" # Execute detection
-
Review Results
- Examine threats by STRIDE category
- Prioritize by severity
- Generate remediation plan
format_results_table()- Display results with formattingdisplay_summary_statistics()- Show detection summariesvalidate_environment_variables()- Check required configprint_section_header(),print_success(), etc. - Formatted output
Neo4jConnection- Database connection managementImageProcessor- Image handling for ClaudeDiagramAnalyzer- Diagram-to-Cypher conversioninitialize_connections()- Setup connectionsexecute_dynamic_cypher()- Execute Cypher with validation
CypherQueryLoader- Load and parse query filesThreatQueryExecutor- Execute threat queriesThreatReportGenerator- Generate analysis reportsload_and_execute_queries()- Complete workflowrun_complete_threat_detection()- End-to-end detection
ThreatDocumentationRetriever- Fetch threat docsSchemaManager- Load Neo4j schemaThreatAnalyzer- Claude AI threat analysisQueryParser- Parse generated queriesLLMThreatAnalyzer- Complete workflow orchestration
AbuseCaseGenerator- Generate abuse cases from user storiesAssetIdentifier- Identify and classify data assets
do_code_analysis()- Execute code analysis with Claude CLI
- 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
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
This project is licensed under the MIT License - see the LICENSE file for details.
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.