Threat Modeling

# Python: simple threat model data structure
from dataclasses import dataclass, field
from enum import Enum
from typing import List

class StrideCategory(Enum):
    SPOOFING = "Spoofing"
    TAMPERING = "Tampering"
    REPUDIATION = "Repudiation"
    INFORMATION_DISCLOSURE = "Information Disclosure"
    DENIAL_OF_SERVICE = "Denial of Service"
    ELEVATION_OF_PRIVILEGE = "Elevation of Privilege"

@dataclass
class Threat:
    category: StrideCategory
    component: str
    description: str
    mitigation: str
    risk_rating: str    # Critical / High / Medium / Low
    status: str         # Open / Mitigated / Accepted

@dataclass
class ThreatModel:
    system_name: str
    threats: List[Threat] = field(default_factory=list)

    def add_threat(self, threat: Threat):
        self.threats.append(threat)

    def open_threats(self) -> List[Threat]:
        return [t for t in self.threats if t.status == "Open"]

    def summary(self):
        by_rating = {}
        for t in self.open_threats():
            by_rating.setdefault(t.risk_rating, []).append(t)
        for rating in ['Critical', 'High', 'Medium', 'Low']:
            items = by_rating.get(rating, [])
            if items:
                print(f"{rating}: {len(items)} open threats")

model = ThreatModel("Payment API")
model.add_threat(Threat(
    category=StrideCategory.SPOOFING,
    component="POST /api/v1/transfer",
    description="Attacker forges JWT to impersonate account holder",
    mitigation="Validate RS256 JWT signature against JWKS endpoint; enforce exp claim",
    risk_rating="Critical",
    status="Open",
))
model.summary()

# Python: represent an attack tree
from dataclasses import dataclass, field
from typing import List, Optional
from enum import Enum

class NodeType(Enum):
    OR = "OR"     # any child being true makes parent true
    AND = "AND"   # all children must be true

@dataclass
class AttackNode:
    description: str
    node_type: NodeType = NodeType.OR
    likelihood: str = "Medium"   # High / Medium / Low
    children: List["AttackNode"] = field(default_factory=list)

    def is_feasible(self) -> bool:
        if not self.children:
            return self.likelihood in ("High", "Medium")
        if self.node_type == NodeType.OR:
            return any(c.is_feasible() for c in self.children)
        return all(c.is_feasible() for c in self.children)

root = AttackNode(
    "Exfiltrate customer PII",
    NodeType.OR,
    children=[
        AttackNode("SQL injection in search", NodeType.AND, children=[
            AttackNode("Find injectable parameter", likelihood="High"),
            AttackNode("Dump users table via UNION", likelihood="High"),
        ]),
        AttackNode("Compromise developer account", NodeType.AND, children=[
            AttackNode("Phish developer credentials", likelihood="Medium"),
            AttackNode("Use credentials without MFA", likelihood="Low"),
        ]),
    ]
)

print(f"Attack feasible: {root.is_feasible()}")

# save as payment_model.py, then run:
#   python3 payment_model.py --list        # print STRIDE threat list
#   python3 payment_model.py --dfd         # generate DFD (requires graphviz)
#   python3 payment_model.py --json out.json  # machine-readable output

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

tm = TM("Payment API Threat Model")
tm.description = "Threat model for the payment processing API"
tm.isOrdered = True

# Boundaries
internet = Boundary("Internet")
dmz = Boundary("DMZ")
internal = Boundary("Internal Network")

# Components
user = Actor("Customer")
user.inBoundary = internet

api = Server("Payment API")
api.inBoundary = dmz
api.usesSessionTokens = True   # enables session-hijacking threats
api.minTLSVersion = TLSVersion.TLSv12  # triggers outdated-TLS threats below this

database = Datastore("PostgreSQL")
database.inBoundary = internal
database.isSQL = True          # enables SQL injection threat rules
database.storesSensitiveData = True  # enables data-exposure threat rules

# Data flows
Dataflow(user, api, "HTTPS POST /api/v1/transfer")
Dataflow(api, database, "Parameterized SQL query")
Dataflow(database, api, "Query results")
Dataflow(api, user, "HTTPS JSON response")

tm.process()   # reads sys.argv; behaviour depends on the flag passed above

from dataclasses import dataclass

@dataclass
class DREADScore:
    damage: int          # 1-3
    reproducibility: int
    exploitability: int
    affected_users: int
    discoverability: int

    def total(self) -> int:
        return self.damage + self.reproducibility + self.exploitability \
               + self.affected_users + self.discoverability

    def rating(self) -> str:
        score = self.total()
        if score >= 12:
            return "Critical"
        if score >= 8:
            return "High"
        if score >= 5:
            return "Medium"
        return "Low"

sqli = DREADScore(damage=3, reproducibility=3, exploitability=2,
                  affected_users=3, discoverability=2)
print(f"SQL injection: {sqli.total()}/15 — {sqli.rating()}")
# SQL injection: 13/15 — Critical

# Parse CVSS scores from Trivy JSON output
import json
import subprocess

result = subprocess.run(
    ['trivy', 'image', '--format', 'json', 'nginx:latest'],
    capture_output=True, text=True
)
data = json.loads(result.stdout)

critical_vulns = []
for result in data.get('Results', []):
    for vuln in result.get('Vulnerabilities', []):
        score = vuln.get('CVSS', {}).get('nvd', {}).get('V3Score', 0)
        if score >= 9.0:
            critical_vulns.append({
                'id': vuln['VulnerabilityID'],
                'package': vuln['PkgName'],
                'score': score,
                'description': vuln.get('Description', '')[:100],
            })

for v in sorted(critical_vulns, key=lambda x: x['score'], reverse=True):
    print(f"{v['score']:.1f} {v['id']} ({v['package']}): {v['description']}")