Lab Specification — Module S07: Threat Modeling Harnesses

from pydantic import BaseModel, Field
from typing import Literal

class DFDElement(BaseModel):
    id: str
    name: str
    type: Literal["external_actor", "process", "data_store", "service"]
    trust_boundary_id: str
    technologies: list[str] = Field(default_factory=list)
    exposes: list[str] = Field(default_factory=list)

class Flow(BaseModel):
    id: str
    source: str
    target: str
    data: list[str] = Field(default_factory=list)
    protocol: str = "HTTPS"
    crosses_boundary: bool = False

class TrustBoundary(BaseModel):
    id: str
    name: str
    description: str

class DFD(BaseModel):
    elements: list[DFDElement]
    flows: list[Flow]
    trust_boundaries: list[TrustBoundary]
    metadata: dict

import hcl2, json

def parse_terraform(tf_path: str) -> tuple[list[DFDElement], list[Flow]]:
    """Parse Terraform resources into DFD elements + flows."""
    with open(tf_path) as f:
        tf = hcl2.load(f)

    elements = []
    # Map resource types to DFD element types
    type_map = {
        "aws_s3_bucket": ("data_store", ["S3"]),
        "aws_lambda_function": ("process", ["lambda"]),
        "aws_api_gateway_rest_api": ("service", ["api_gateway"]),
        "aws_dynamodb_table": ("data_store", ["dynamodb"]),
        "aws_rds_cluster": ("data_store", ["postgres"]),
    }
    for resource_block in tf.get("resource", []):
        for rtype, instances in resource_block.items():
            for name, config in instances.items():
                elem_type, tech = type_map.get(rtype, ("process", [rtype]))
                boundary = classify_boundary(rtype, config)
                elements.append(DFDElement(
                    id=f"{rtype}.{name}",
                    name=config.get("name", name),
                    type=elem_type,
                    trust_boundary_id=boundary,
                    technologies=tech,
                    exposes=get_exposes(rtype, config),
                ))
    flows = extract_flows_from_references(tf, elements)
    return elements, flows

def classify_boundary(rtype: str, config: dict) -> str:
    """Classify an element's trust boundary from its config."""
    # Internet-facing: public IP, public LB, internet gateway
    if config.get("publicly_accessible") or config.get("ingress"):
        return "internet_facing"
    return "internal_vpc"

def mark_crossing_flows(flows: list[Flow], elements: list[DFDElement]) -> list[Flow]:
    elem_boundaries = {e.id: e.trust_boundary_id for e in elements}
    for f in flows:
        src_b = elem_boundaries.get(f.source, "unknown")
        tgt_b = elem_boundaries.get(f.target, "unknown")
        f.crosses_boundary = (src_b != tgt_b)
    return flows

STRIDE_APPLICABILITY = {
    "external_actor": ["spoofing", "repudiation"],
    "process": ["spoofing", "tampering", "repudiation", "info_disclosure", "denial_of_service", "elevation_of_privilege"],
    "service": ["spoofing", "tampering", "repudiation", "info_disclosure", "denial_of_service", "elevation_of_privilege"],
    "data_store": ["tampering", "repudiation", "info_disclosure", "denial_of_service"],
}

import json

async def generate_threats(element: DFDElement, category: str, flows: list[Flow]) -> list[Threat]:
    element_flows = [f for f in flows if f.source == element.id or f.target == element.id]
    prompt = f"""
    You are a threat modeling analyst applying STRIDE.

    Element: {element.name} (type: {element.type})
    Technologies: {element.technologies}
    Trust boundary: {element.trust_boundary_id}
    STRIDE category: {category}
    Flows crossing this element's boundary:
    {json.dumps([f.model_dump() for f in element_flows], indent=2)}

    Generate candidate threats in this category that are SPECIFIC to this element,
    its technologies, and the data that flows through it. For each threat:
    - A concrete attack narrative (how would an adversary do this?)
    - The data or asset at risk
    - A draft CVSS base vector (AV/AC/PR/UI/S/C/I/A)

    Do NOT produce generic threats. Ground each in the actual technologies and flows.
    Respond as JSON list: [{{"title", "attack_narrative", "asset_at_risk", "cvss_vector", "cvss_score"}}]
    """
    result = await llm_complete(prompt)
    return [Threat(**t, element_id=element.id, category=category) for t in json.loads(result)]

def dedup_threats(threats: list[Threat], similarity_threshold: float = 0.85) -> list[Threat]:
    """Dedup threats via semantic similarity; merge near-duplicates, retain highest severity."""
    embeddings = [embed(t.title + " " + t.attack_narrative) for t in threats]
    kept = []
    for i, t in enumerate(threats):
        is_dup = any(
            cosine_sim(embeddings[i], embeddings[j]) > similarity_threshold
            for j, k in enumerate(kept)
        )
        if not is_dup:
            kept.append(i)
    # Prioritize by CVSS score descending
    return sorted([threats[i] for i in kept], key=lambda t: t.cvss_score, reverse=True)

async def generate_mitigation(threat: Threat) -> Mitigation:
    prompt = f"""
    Threat: {threat.attack_narrative}
    Element: {threat.element_id}
    CVSS draft: {threat.cvss_vector} (score: {threat.cvss_score})

    Provide a mitigation that:
    1. Maps to the relevant OWASP ASVS control.
    2. References the relevant CWE remediation.
    3. References the specific cloud provider best practice (AWS/Azure/GCP).
    4. Is implementable as a concrete change (a config, an IaC update, a code change) — not a platitude.

    Respond as JSON: {{"owasp_control", "cwe_mitigation", "cloud_best_practice", "implementation"}}
    """
    result = await llm_complete(prompt)
    return Mitigation(**json.loads(result), threat_id=threat.id, status="open")

import subprocess

def create_github_issue(threat: Threat, mitigation: Mitigation, repo: str) -> str:
    body = f"""
## Security Threat (auto-generated from threat model)

**Element**: {threat.element_id}
**STRIDE category**: {threat.category}
**CVSS draft**: {threat.cvss_vector} (score: {threat.cvss_score})

**Attack narrative**:
{threat.attack_narrative}

**Mitigation**:
{mitigation.implementation}

**References**:
- OWASP: {mitigation.owasp_control}
- CWE: {mitigation.cwe_mitigation}
- Cloud best practice: {mitigation.cloud_best_practice}

---
Generated by threat modeling harness.
"""
    result = subprocess.run(
        ["gh", "issue", "create",
         "--repo", repo,
         "--title", f"[Security] {threat.title}",
         "--body", body,
         "--label", f"security,stride:{threat.category}"],
        capture_output=True, text=True
    )
    return result.stdout.strip()  # issue URL

# Lab Specification — Module S07: Threat Modeling Harnesses

**Course**: 2A — Building AI Harnesses for Cybersecurity
**Module**: S07 — Threat Modeling Harnesses
**Duration**: 90–120 minutes (three labs, one per sub-section)
**Environment**: Python 3.11+, Pydantic. Sample Terraform config and draw.io XML (provided). An LLM API key for threat and mitigation generation. `gh` CLI for issue creation (lab 3).

---

## Learning objectives

1. Build an architecture-to-DFD ingestion pipeline that parses Terraform and draw.io into a structured DFD with trust boundaries.
2. Implement a STRIDE engine that applies applicable categories per element and generates grounded threats with dedup and prioritization.
3. Generate per-threat mitigations mapped to OWASP/CWE/cloud best practices and auto-create GitHub issues.

---

## Phase 1 — Architecture Ingestion (35 min)

### 1.1 Define the DFD schema

```python
from pydantic import BaseModel, Field
from typing import Literal

class DFDElement(BaseModel):
    id: str
    name: str
    type: Literal["external_actor", "process", "data_store", "service"]
    trust_boundary_id: str
    technologies: list[str] = Field(default_factory=list)
    exposes: list[str] = Field(default_factory=list)

class Flow(BaseModel):
    id: str
    source: str
    target: str
    data: list[str] = Field(default_factory=list)
    protocol: str = "HTTPS"
    crosses_boundary: bool = False

class TrustBoundary(BaseModel):
    id: str
    name: str
    description: str

class DFD(BaseModel):
    elements: list[DFDElement]
    flows: list[Flow]
    trust_boundaries: list[TrustBoundary]
    metadata: dict
```

### 1.2 Implement the Terraform parser

```python
import hcl2, json

def parse_terraform(tf_path: str) -> tuple[list[DFDElement], list[Flow]]:
    """Parse Terraform resources into DFD elements + flows."""
    with open(tf_path) as f:
        tf = hcl2.load(f)

    elements = []
    # Map resource types to DFD element types
    type_map = {
        "aws_s3_bucket": ("data_store", ["S3"]),
        "aws_lambda_function": ("process", ["lambda"]),
        "aws_api_gateway_rest_api": ("service", ["api_gateway"]),
        "aws_dynamodb_table": ("data_store", ["dynamodb"]),
        "aws_rds_cluster": ("data_store", ["postgres"]),
    }
    for resource_block in tf.get("resource", []):
        for rtype, instances in resource_block.items():
            for name, config in instances.items():
                elem_type, tech = type_map.get(rtype, ("process", [rtype]))
                boundary = classify_boundary(rtype, config)
                elements.append(DFDElement(
                    id=f"{rtype}.{name}",
                    name=config.get("name", name),
                    type=elem_type,
                    trust_boundary_id=boundary,
                    technologies=tech,
                    exposes=get_exposes(rtype, config),
                ))
    flows = extract_flows_from_references(tf, elements)
    return elements, flows

def classify_boundary(rtype: str, config: dict) -> str:
    """Classify an element's trust boundary from its config."""
    # Internet-facing: public IP, public LB, internet gateway
    if config.get("publicly_accessible") or config.get("ingress"):
        return "internet_facing"
    return "internal_vpc"
```

### 1.3 Implement trust boundary assignment and mark crossing flows

```python
def mark_crossing_flows(flows: list[Flow], elements: list[DFDElement]) -> list[Flow]:
    elem_boundaries = {e.id: e.trust_boundary_id for e in elements}
    for f in flows:
        src_b = elem_boundaries.get(f.source, "unknown")
        tgt_b = elem_boundaries.get(f.target, "unknown")
        f.crosses_boundary = (src_b != tgt_b)
    return flows
```

### 1.4 Verify against the sample

1. Run the parser on the provided `sample/main.tf` (contains an API gateway, a Lambda, an S3 bucket, and a DynamoDB table).
2. Verify: 4 elements extracted with correct types; trust boundaries assigned; flows between them marked `crosses_boundary` where applicable.

### Deliverable
- [ ] Terraform parsed into DFD elements with correct types
- [ ] Trust boundaries assigned; crossing flows marked
- [ ] DFD object validates against the schema

---

## Phase 2 — STRIDE Analysis Engine (35 min)

### 2.1 Implement the STRIDE applicability matrix

```python
STRIDE_APPLICABILITY = {
    "external_actor": ["spoofing", "repudiation"],
    "process": ["spoofing", "tampering", "repudiation", "info_disclosure", "denial_of_service", "elevation_of_privilege"],
    "service": ["spoofing", "tampering", "repudiation", "info_disclosure", "denial_of_service", "elevation_of_privilege"],
    "data_store": ["tampering", "repudiation", "info_disclosure", "denial_of_service"],
}
```

### 2.2 Implement grounded threat generation

```python
import json

async def generate_threats(element: DFDElement, category: str, flows: list[Flow]) -> list[Threat]:
    element_flows = [f for f in flows if f.source == element.id or f.target == element.id]
    prompt = f"""
    You are a threat modeling analyst applying STRIDE.

    Element: {element.name} (type: {element.type})
    Technologies: {element.technologies}
    Trust boundary: {element.trust_boundary_id}
    STRIDE category: {category}
    Flows crossing this element's boundary:
    {json.dumps([f.model_dump() for f in element_flows], indent=2)}

    Generate candidate threats in this category that are SPECIFIC to this element,
    its technologies, and the data that flows through it. For each threat:
    - A concrete attack narrative (how would an adversary do this?)
    - The data or asset at risk
    - A draft CVSS base vector (AV/AC/PR/UI/S/C/I/A)

    Do NOT produce generic threats. Ground each in the actual technologies and flows.
    Respond as JSON list: [{{"title", "attack_narrative", "asset_at_risk", "cvss_vector", "cvss_score"}}]
    """
    result = await llm_complete(prompt)
    return [Threat(**t, element_id=element.id, category=category) for t in json.loads(result)]
```

### 2.3 Implement dedup and prioritization

```python
def dedup_threats(threats: list[Threat], similarity_threshold: float = 0.85) -> list[Threat]:
    """Dedup threats via semantic similarity; merge near-duplicates, retain highest severity."""
    embeddings = [embed(t.title + " " + t.attack_narrative) for t in threats]
    kept = []
    for i, t in enumerate(threats):
        is_dup = any(
            cosine_sim(embeddings[i], embeddings[j]) > similarity_threshold
            for j, k in enumerate(kept)
        )
        if not is_dup:
            kept.append(i)
    # Prioritize by CVSS score descending
    return sorted([threats[i] for i in kept], key=lambda t: t.cvss_score, reverse=True)
```

### 2.4 Run against the sample DFD

1. For each element, apply the applicable STRIDE categories.
2. Generate grounded threats per (element, category).
3. Dedup and prioritize.
4. Verify: output is a ranked list of specific, grounded threats (not generic "attacker might spoof" statements).

### Deliverable
- [ ] STRIDE applicability matrix enforced (data stores not analyzed for spoofing)
- [ ] Grounded threats generated with attack narratives and CVSS drafts
- [ ] Dedup collapses near-duplicates; output ranked by CVSS

---

## Phase 3 — Mitigation Generation and Issue Tracking (30 min)

### 3.1 Implement mitigation generation

```python
async def generate_mitigation(threat: Threat) -> Mitigation:
    prompt = f"""
    Threat: {threat.attack_narrative}
    Element: {threat.element_id}
    CVSS draft: {threat.cvss_vector} (score: {threat.cvss_score})

    Provide a mitigation that:
    1. Maps to the relevant OWASP ASVS control.
    2. References the relevant CWE remediation.
    3. References the specific cloud provider best practice (AWS/Azure/GCP).
    4. Is implementable as a concrete change (a config, an IaC update, a code change) — not a platitude.

    Respond as JSON: {{"owasp_control", "cwe_mitigation", "cloud_best_practice", "implementation"}}
    """
    result = await llm_complete(prompt)
    return Mitigation(**json.loads(result), threat_id=threat.id, status="open")
```

### 3.2 Auto-create GitHub issues from high-severity open threats

```python
import subprocess

def create_github_issue(threat: Threat, mitigation: Mitigation, repo: str) -> str:
    body = f"""
## Security Threat (auto-generated from threat model)

**Element**: {threat.element_id}
**STRIDE category**: {threat.category}
**CVSS draft**: {threat.cvss_vector} (score: {threat.cvss_score})

**Attack narrative**:
{threat.attack_narrative}

**Mitigation**:
{mitigation.implementation}

**References**:
- OWASP: {mitigation.owasp_control}
- CWE: {mitigation.cwe_mitigation}
- Cloud best practice: {mitigation.cloud_best_practice}

---
Generated by threat modeling harness.
"""
    result = subprocess.run(
        ["gh", "issue", "create",
         "--repo", repo,
         "--title", f"[Security] {threat.title}",
         "--body", body,
         "--label", f"security,stride:{threat.category}"],
        capture_output=True, text=True
    )
    return result.stdout.strip()  # issue URL
```

### 3.3 Run end-to-end

1. Generate mitigations for the threat list from Phase 2.
2. For each high-severity open threat (CVSS ≥ 7.0), auto-create a GitHub issue.
3. Verify: 3 issues created with full threat context, mitigation mapping, and traceability.

### Deliverable
- [ ] Mitigations generated with 3-layer mapping (OWASP, CWE, cloud)
- [ ] Implementation field is concrete (IaC/config/code), not a platitude
- [ ] GitHub issues auto-created for high-severity open threats with full traceability

---

## Stretch goals

1. **Implement model version diffing**: store the DFD from this run; on the next run (after modifying the Terraform), diff the two models and re-run STRIDE only on new/changed elements. Verify removed elements retire their threats.
2. **Add draw.io ingestion**: parse the provided `sample/architecture.drawio` XML into DFD elements and merge with the Terraform-derived DFD. Resolve conflicts (IaC authoritative).
3. **Bidirectional status sync**: poll open GitHub issues; when one closes, update the corresponding threat's status to "mitigated" in the threat model store.