Threat Modeling with pytm

Overview

pytm is a Python library for programmatic threat modeling based on the STRIDE methodology. It enables security engineers to define system architecture as code, automatically generate data flow diagrams (DFDs), identify security threats across trust boundaries, and produce comprehensive threat reports. This approach integrates threat modeling into CI/CD pipelines, enabling shift-left security and continuous threat analysis.

Quick Start

Create a basic threat model:

#!/usr/bin/env python3
from pytm import TM, Server, Dataflow, Boundary, Actor

# Initialize threat model
tm = TM("Web Application Threat Model")
tm.description = "E-commerce web application"

# Define trust boundaries
internet = Boundary("Internet")
dmz = Boundary("DMZ")
internal = Boundary("Internal Network")

# Define actors and components
user = Actor("Customer")
user.inBoundary = internet

web = Server("Web Server")
web.inBoundary = dmz

db = Server("Database")
db.inBoundary = internal

# Define data flows
user_to_web = Dataflow(user, web, "HTTPS Request")
user_to_web.protocol = "HTTPS"
user_to_web.data = "credentials, payment info"
user_to_web.isEncrypted = True

web_to_db = Dataflow(web, db, "Database Query")
web_to_db.protocol = "SQL/TLS"
web_to_db.data = "user data, transactions"

# Generate threat report and diagram
tm.process()

Install pytm:

pip install pytm
# Also requires graphviz for diagram generation
brew install graphviz  # macOS
# or: apt-get install graphviz  # Linux

Core Workflows

Workflow 1: Create New Threat Model

Progress: [ ] 1. Define system scope and trust boundaries [ ] 2. Identify all actors (users, administrators, external systems) [ ] 3. Map system components (servers, databases, APIs, services) [ ] 4. Define data flows between components with security attributes [ ] 5. Run tm.process() to generate threats and DFD [ ] 6. Review STRIDE threats and add mitigations [ ] 7. Generate threat report with scripts/generate_report.py

Work through each step systematically. Check off completed items.

Workflow 2: STRIDE Threat Analysis

pytm automatically identifies threats based on STRIDE categories:

• Spoofing: Identity impersonation attacks
• Tampering: Unauthorized modification of data
• Repudiation: Denial of actions without traceability
• Information Disclosure: Unauthorized access to sensitive data
• Denial of Service: Availability attacks
• Elevation of Privilege: Unauthorized access escalation

For each identified threat:

1. Review threat description and affected component
2. Assess likelihood and impact (use references/risk_matrix.md)
3. Determine if existing controls mitigate the threat
4. Add mitigation using threat.mitigation = "description"
5. Document residual risk and acceptance criteria

Workflow 3: Architecture as Code

Define system architecture programmatically:

from pytm import TM, Server, Datastore, Dataflow, Boundary, Actor, Lambda

tm = TM("Microservices Architecture")

# Cloud boundaries
internet = Boundary("Internet")
cloud_vpc = Boundary("Cloud VPC")

# API Gateway
api_gateway = Server("API Gateway")
api_gateway.inBoundary = cloud_vpc
api_gateway.implementsAuthentication = True
api_gateway.implementsAuthorization = True

# Microservices
auth_service = Lambda("Auth Service")
auth_service.inBoundary = cloud_vpc

order_service = Lambda("Order Service")
order_service.inBoundary = cloud_vpc

# Data stores
user_db = Datastore("User Database")
user_db.inBoundary = cloud_vpc
user_db.isEncryptedAtRest = True

# Data flows with security properties
client_to_api = Dataflow(Actor("Client"), api_gateway, "API Request")
client_to_api.protocol = "HTTPS"
client_to_api.isEncrypted = True
client_to_api.data = "user credentials, orders"

api_to_auth = Dataflow(api_gateway, auth_service, "Auth Check")
api_to_auth.protocol = "gRPC/TLS"

auth_to_db = Dataflow(auth_service, user_db, "User Lookup")
auth_to_db.protocol = "TLS"

tm.process()

Workflow 4: CI/CD Integration

Automate threat modeling in continuous integration:

# .github/workflows/threat-model.yml
name: Threat Model Analysis
on: [push, pull_request]

jobs:
  threat-model:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v3

      - name: Set up Python
        uses: actions/setup-python@v4
        with:
          python-version: '3.10'

      - name: Install dependencies
        run: |
          pip install pytm
          sudo apt-get install -y graphviz

      - name: Generate threat model
        run: python threat_model.py

      - name: Upload DFD diagram
        uses: actions/upload-artifact@v3
        with:
          name: threat-model-dfd
          path: '*.png'

      - name: Check for unmitigated threats
        run: python scripts/check_mitigations.py threat_model.py

Workflow 5: Threat Report Generation

Generate comprehensive threat documentation:

# Run threat model with report generation
python threat_model.py

# Generate markdown report
./scripts/generate_report.py --model threat_model.py --output threat_report.md

# Generate JSON for tool integration
./scripts/generate_report.py --model threat_model.py --format json --output threats.json

Report includes:

• System architecture overview
• Trust boundary analysis
• Complete STRIDE threat enumeration
• Existing and recommended mitigations
• Risk prioritization matrix

Security Considerations

Sensitive Data Handling

• Threat models as code: Store in version control but review for sensitive architecture details
• Credentials and secrets: Never hardcode in threat models - use placeholders
• Data classification: Clearly label data flows with sensitivity levels (PII, PCI, PHI)
• Report distribution: Control access to threat reports revealing security architecture

Access Control

• Threat model repository: Restrict write access to security team and architects
• CI/CD integration: Protect threat modeling pipeline from tampering
• Diagram artifacts: Control distribution of DFDs showing system architecture
• Mitigation tracking: Integrate with secure issue tracking systems

Audit Logging

Log the following for security governance:

• Threat model creation and modification history
• Identified threats and severity assessments
• Mitigation implementation and validation
• Risk acceptance decisions with approval
• Threat model review and update cycles

Compliance Requirements

• NIST 800-30: Risk assessment methodology alignment
• ISO 27001: A.14.1.2 - Securing application services on public networks
• OWASP SAMM: Threat Assessment practice maturity
• PCI-DSS 6.3.1: Security threat identification in development
• SOC2 CC9.1: Risk assessment process for system changes

Bundled Resources

Scripts (scripts/)

• generate_report.py - Generate markdown/JSON threat reports with STRIDE categorization
• check_mitigations.py - Validate all identified threats have documented mitigations
• threat_classifier.py - Classify threats by severity using DREAD or custom risk matrix
• template_generator.py - Generate threat model templates for common architectures

References (references/)

• stride_methodology.md - Complete STRIDE methodology guide with threat examples
• risk_matrix.md - Risk assessment framework with likelihood and impact scoring
• component_library.md - Reusable pytm components for common patterns (APIs, databases, cloud services)
• mitigation_strategies.md - Common mitigation patterns mapped to STRIDE categories and OWASP controls

Assets (assets/)

• templates/web_application.py - Web application threat model template
• templates/microservices.py - Microservices architecture template
• templates/mobile_app.py - Mobile application threat model template
• templates/iot_system.py - IoT system threat model template
• dfd_styles.json - Custom graphviz styl