It leverages AI to analyze service descriptions, identify security threats, visualize data flow and suggest remediations based on the STRIDE framework and OWASP guidelines. The script generates in-depth HTML report and includes feature for manual risk assessment.
- YAML File Processing: Loads and validates a YAML file containing the service details
- AI Threat Analysis: If an OpenAI API key is provided, the script uses AI to generate a comprehensive threat modeling analysis
- AI Threat Cross validation: Various LLMs performs a validation results
- Data Flow Generation: Automatically generates a visual representation of the data flow within the service
- Manual Risk Management: Users can manually add, modify, or cross out risks in the generated report
- Report Generation: Produces a detailed HTML report, including both AI-generated and manually added risks
- GPT-3.5
- GPT-4
- Claude 3 Haiku
- Mistral 7b (through ollama)
- Gemini (Planned)
Set the openai api key as an environment variable OPENAI_KEY and ANTHROPIC_KEY in your operating system
Linux/Mac
export OPENAI_KEY=sk-ApiKeyExample
export ANTHROPIC_KEY=sk-ant-api03-ApiKeyExample
Windows
set OPENAI_KEY=sk-ApiKeyExample
set ANTHROPIC_KEY=sk-ant-api03-ApiKeyExample
pip3 install -r requirements.txt
- Download and install Ollama ollama.ai
- Install Mistral
ollama pull mistral
- Start ollama service (make sure the Ollama desktop app is closed)
ollama serve
Now you can use mistral as the main model:
python3 TaaC-AI.py --model mistral <path_to_yaml_file>
Or for cross-validation
python3 TaaC-AI.py --model claude --cross-validation mistral <path_to_yaml_file>
- Create a valid service description using these guidelines or use
taac_yaml_generator.py
that will guide you through the process of generating one
- Execute the script (GPT-3.5 is used by default)
python3 TaaC-AI.py <path_to_yaml_file>
Use gpt-4
, claude
or mistral
as a model to identify threats by specifying the --model
option
python3 TaaC-AI.py --model gpt-4 <path_to_yaml_file>
To perform Threats result validation by another LLM use --cross-validation
option.
python3 TaaC-AI.py --model claude --cross-validation claude <path_to_yaml_file>
- Open generate .html report
- Review/Edit AI-driven Threat Modeling Analysis table, and for false positives or resolved issues, mark the 'Status' checkbox
- Add manually identified threats to the table (optional)
- Download the report via the Download Report button
- Valid service description example
Version: '1.0'
Date: 14.11.2023
# Authentication Service Description
Description:
Name: AuthService
Type: Service
Criticality: Tier1
# Service Functionality
Functionality: Handles user authentication, including login and token generation.
# Data Processing Details
DataProcessed:
Type: Confidential
DataCategory: Auth
EncryptionAtRest: Yes
# Components Used by the Service
Components:
Internal:
Exist: Yes
Source: Private
Note: Scoped Package Access
External:
Exist: Yes
PackageManager: NPM
# Pipeline Configuration
Pipeline:
Type: GithubActions
CODEOWNERS: Yes
BranchProtection: Yes
SignCommits: Yes
PinActions: Yes
# Network Information
Network:
Access: Private
# Authentication Service Data Flow
dataFlow: # Removed the dash here
- name: UserAuthenticationFlow
description: Authenticates users and issues tokens.
source: UserLoginInterface
EncryptionTransit: Yes
Authentication:
Exist: Yes
Type: JWT
Authorization: read-write
Protocol: HTTPS
Communication:
Type: RESTful API
interactions:
- from: UserLoginInterface
to: AuthService
method: RESTful API
protocol: HTTPS
- from: AuthService
to: UserDatabase
method: Query
protocol: JDBC
servicesInvolved: [UserLoginInterface, AuthService, UserDatabase]
- Script execution
- Download and Review the generated HTML report
- Add Threats manually
- Download the latest report
Template DesignBasic FunctionalityGPT-3 IntegrationReport generationManually adding identified threatsGPT-4 IntegrationClaude IntegrationLLM Cross ValidationMistral Integration via Ollama- Gemini Integration
- Accuracy Comparison
All suggestions write to [email protected]