AI-driven Threat Modeling-as-a-Code (TaaC-AI)

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.

Design

How It Works

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

Model supported

GPT-3.5
GPT-4
Claude 3 Haiku
‎Gemini (Planned)

Set API Key

Set the openai api key as an environment variable OPENAI_KEY and ANTHROPIC_KEY in your operating system

OpenAI API Anthropic API

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

Install dependencies

pip3 install -r requirements.txt

How to Use ❓

Create a valid service description using these guidelines
Execute the script (GPT-3.5 is used by default)

python3 TaaC-AI.py <path_to_yaml_file>

Use GPT-4 or Claude 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. Cross-validation will be done by 'claude' for gpt-4, by gpt-4 for claude and gpt-3.5-turbo.

python3 TaaC-AI.py --model claude --cross-validation <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

Usage Example 🏁

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]