Dreigingsmodellering voor AI-infrastructuur met STRIDE

"""
AI infrastructure data flow diagram generator.
Creates structured DFD representations for threat modeling
ML systems using STRIDE.
"""
 
import json
from dataclasses import dataclass, field
from typing import Optional
from enum import Enum
 
class ElementType(Enum):
    PROCESS = "process"
    DATA_STORE = "data_store"
    EXTERNAL_ENTITY = "external_entity"
    DATA_FLOW = "data_flow"
    TRUST_BOUNDARY = "trust_boundary"
 
class TrustZone(Enum):
    EXTERNAL = "external"             # Internet, eindgebruikers
    DMZ = "dmz"                       # API-gateways, load balancers
    INFERENCE = "inference"           # Model serving
    TRAINING = "training"            # Trainingscluster
    DATA = "data"                    # Opslag van trainingsdata
    CONTROL = "control"              # Pijplijnorkestratie
    MONITORING = "monitoring"        # Observability
    REGISTRY = "registry"           # Model-/artefactregister
 
@dataclass
class DFDElement:
    """An element in the data flow diagram."""
    id: str
    name: str
    element_type: ElementType
    trust_zone: TrustZone
    description: str
    technologies: list[str] = field(default_factory=list)
    threats: list[str] = field(default_factory=list)
 
@dataclass
class DataFlow:
    """A data flow between DFD elements."""
    id: str
    source_id: str
    target_id: str
    data_description: str
    protocol: str
    is_encrypted: bool = False
    is_authenticated: bool = False
    threats: list[str] = field(default_factory=list)
 
class AIThreatModel:
    """
    Threat model for an AI infrastructure deployment.
    """
 
    def __init__(self, system_name: str, description: str):
        self.system_name = system_name
        self.description = description
        self.elements: dict[str, DFDElement] = {}
        self.data_flows: list[DataFlow] = []
        self.threats: list[dict] = []
 
    def add_element(self, element: DFDElement) -> None:
        """Add a DFD element."""
        self.elements[element.id] = element
 
    def add_data_flow(self, flow: DataFlow) -> None:
        """Add a data flow between elements."""
        self.data_flows.append(flow)
 
    def apply_stride(self) -> list[dict]:
        """
        Systematically apply STRIDE to all elements and data flows.
        Returns identified threats.
        """
        threats = []
 
        # Pas STRIDE toe op elk element op basis van type
        for elem in self.elements.values():
            if elem.element_type == ElementType.PROCESS:
                threats.extend(self._stride_process(elem))
            elif elem.element_type == ElementType.DATA_STORE:
                threats.extend(self._stride_data_store(elem))
            elif elem.element_type == ElementType.EXTERNAL_ENTITY:
                threats.extend(self._stride_external_entity(elem))
 
        # Pas STRIDE toe op elke data flow
        for flow in self.data_flows:
            threats.extend(self._stride_data_flow(flow))
 
        # Pas grensoverschrijdende dreigingen toe
        threats.extend(self._cross_boundary_threats())
 
        self.threats = threats
        return threats
 
    def _stride_process(self, elem: DFDElement) -> list[dict]:
        """Apply all STRIDE categories to a process element."""
        threats = []
        prefix = f"{elem.name} ({elem.id})"
 
        # Spoofing
        threats.append({
            "element": elem.id,
            "category": "Spoofing",
            "threat": (
                f"An attacker could impersonate {prefix} to "
                f"interact with upstream or downstream components."
            ),
            "ai_specific": "model" in elem.name.lower() or "inference" in elem.name.lower(),
            "mitigation": "Implement mutual TLS with workload identity (SPIFFE).",
        })
 
        # Tampering
        threats.append({
            "element": elem.id,
            "category": "Tampering",
            "threat": (
                f"An attacker could modify the behavior of {prefix} "
                f"by altering its inputs, configuration, or code."
            ),
            "ai_specific": "training" in elem.name.lower(),
            "mitigation": "Sign artifacts, use read-only filesystems, verify inputs.",
        })
 
        # Repudiation
        threats.append({
            "element": elem.id,
            "category": "Repudiation",
            "threat": (
                f"Actions performed by or through {prefix} may not "
                f"be properly logged, allowing denial of activity."
            ),
            "ai_specific": False,
            "mitigation": "Implement comprehensive audit logging.",
        })
 
        # Information Disclosure
        if "model" in elem.name.lower() or "inference" in elem.name.lower():
            threats.append({
                "element": elem.id,
                "category": "Information Disclosure",
                "threat": (
                    f"{prefix} could leak model weights, architecture, "
                    f"or training data through inference API responses, "
                    f"error messages, or timing side channels."
                ),
                "ai_specific": True,
                "mitigation": (
                    "Implement rate limiting, response filtering, "
                    "differential privacy, and timing normalization."
                ),
            })
 
        # Denial of Service
        threats.append({
            "element": elem.id,
            "category": "Denial of Service",
            "threat": (
                f"{prefix} could be overwhelmed by excessive requests "
                f"or resource-intensive operations."
            ),
            "ai_specific": "gpu" in elem.description.lower(),
            "mitigation": "Implement resource limits, rate limiting, and circuit breakers.",
        })
 
        # Elevation of Privilege
        threats.append({
            "element": elem.id,
            "category": "Elevation of Privilege",
            "threat": (
                f"A vulnerability in {prefix} could allow an attacker "
                f"to gain the privileges of its service account."
            ),
            "ai_specific": False,
            "mitigation": "Apply least privilege, use non-root containers.",
        })
 
        return threats
 
    def _stride_data_store(self, elem: DFDElement) -> list[dict]:
        """Apply applicable STRIDE categories to a data store."""
        threats = []
 
        # Tampering (primaire dreiging voor datastores)
        threats.append({
            "element": elem.id,
            "category": "Tampering",
            "threat": (
                f"Data in {elem.name} could be modified. "
                f"{'This directly enables training data poisoning. ' if 'training' in elem.name.lower() or 'data' in elem.name.lower() else ''}"
                f"{'Model artifacts could be replaced with malicious versions. ' if 'model' in elem.name.lower() or 'registry' in elem.name.lower() else ''}"
            ),
            "ai_specific": True,
            "mitigation": "Enable versioning, access logging, integrity verification.",
        })
 
        # Information Disclosure
        threats.append({
            "element": elem.id,
            "category": "Information Disclosure",
            "threat": (
                f"Unauthorized read access to {elem.name} could "
                f"expose {'training data and PII' if 'data' in elem.name.lower() else 'model intellectual property'}."
            ),
            "ai_specific": True,
            "mitigation": "Encrypt at rest, enforce IAM policies, audit access.",
        })
 
        return threats
 
    def _stride_external_entity(self, elem: DFDElement) -> list[dict]:
        """Spoofing is the primary STRIDE threat for external entities."""
        return [{
            "element": elem.id,
            "category": "Spoofing",
            "threat": (
                f"An attacker could impersonate {elem.name} to "
                f"submit malicious requests or data."
            ),
            "ai_specific": "data" in elem.name.lower(),
            "mitigation": "Authenticate all external entities. Validate inputs.",
        }]
 
    def _stride_data_flow(self, flow: DataFlow) -> list[dict]:
        """Apply Tampering and Information Disclosure to data flows."""
        threats = []
        source = self.elements.get(flow.source_id)
        target = self.elements.get(flow.target_id)
 
        if not flow.is_encrypted:
            threats.append({
                "element": flow.id,
                "category": "Tampering",
                "threat": (
                    f"Data flow from {source.name if source else flow.source_id} "
                    f"to {target.name if target else flow.target_id} "
                    f"({flow.data_description}) is unencrypted. "
                    f"An attacker could modify data in transit."
                ),
                "ai_specific": "model" in flow.data_description.lower(),
                "mitigation": "Enable TLS/mTLS for all data flows.",
            })
 
            threats.append({
                "element": flow.id,
                "category": "Information Disclosure",
                "threat": (
                    f"Unencrypted {flow.data_description} could be "
                    f"intercepted between {source.name if source else '?'} "
                    f"and {target.name if target else '?'}."
                ),
                "ai_specific": True,
                "mitigation": "Encrypt all data in transit.",
            })
 
        if not flow.is_authenticated:
            threats.append({
                "element": flow.id,
                "category": "Spoofing",
                "threat": (
                    f"Data flow to {target.name if target else '?'} "
                    f"is not authenticated. An attacker could send "
                    f"forged {flow.data_description}."
                ),
                "ai_specific": "model" in flow.data_description.lower()
                              or "training" in flow.data_description.lower(),
                "mitigation": "Implement mutual authentication.",
            })
 
        return threats
 
    def _cross_boundary_threats(self) -> list[dict]:
        """Identify threats at trust boundary crossings."""
        threats = []
        for flow in self.data_flows:
            source = self.elements.get(flow.source_id)
            target = self.elements.get(flow.target_id)
            if source and target and source.trust_zone != target.trust_zone:
                threats.append({
                    "element": flow.id,
                    "category": "Trust Boundary Crossing",
                    "threat": (
                        f"Data flows from {source.trust_zone.value} zone "
                        f"({source.name}) to {target.trust_zone.value} zone "
                        f"({target.name}). Cross-boundary data "
                        f"({flow.data_description}) requires validation."
                    ),
                    "ai_specific": True,
                    "mitigation": (
                        "Validate all data crossing trust boundaries. "
                        "Apply defense-in-depth at boundary."
                    ),
                })
        return threats
 
    def generate_report(self) -> str:
        """Generate a structured threat model report."""
        if not self.threats:
            self.apply_stride()
 
        ai_threats = [t for t in self.threats if t.get("ai_specific")]
        traditional_threats = [t for t in self.threats if not t.get("ai_specific")]
 
        # Tel per categorie
        category_counts: dict[str, int] = {}
        for t in self.threats:
            cat = t["category"]
            category_counts[cat] = category_counts.get(cat, 0) + 1
 
        report = {
            "system": self.system_name,
            "description": self.description,
            "summary": {
                "total_threats": len(self.threats),
                "ai_specific_threats": len(ai_threats),
                "traditional_threats": len(traditional_threats),
                "by_category": category_counts,
            },
            "elements": len(self.elements),
            "data_flows": len(self.data_flows),
            "threats": self.threats,
        }
 
        return json.dumps(report, indent=2, default=str)

"""
Complete threat model for an LLM inference pipeline.
Demonstrates applying STRIDE to a realistic architecture.
"""
 
from typing import Any  # Maakt gebruik van de hierboven gedefinieerde klassen
 
def create_llm_inference_threat_model() -> "AIThreatModel":
    """
    Build a threat model for a typical LLM inference pipeline
    deployed on Kubernetes with GPU nodes.
    """
    tm = AIThreatModel(
        system_name="LLM Inference Pipeline",
        description=(
            "Production LLM inference system serving API requests "
            "through a load balancer, with model artifacts stored "
            "in S3 and served via vLLM on GPU nodes."
        ),
    )
 
    # Externe entiteiten
    tm.add_element(DFDElement(
        id="user", name="API Client",
        element_type=ElementType.EXTERNAL_ENTITY,
        trust_zone=TrustZone.EXTERNAL,
        description="External application consuming the LLM API",
        technologies=["HTTPS", "REST"],
    ))
 
    tm.add_element(DFDElement(
        id="admin", name="ML Engineer",
        element_type=ElementType.EXTERNAL_ENTITY,
        trust_zone=TrustZone.EXTERNAL,
        description="Engineer managing model deployments",
        technologies=["kubectl", "mlflow CLI"],
    ))
 
    # Processen
    tm.add_element(DFDElement(
        id="gateway", name="API Gateway",
        element_type=ElementType.PROCESS,
        trust_zone=TrustZone.DMZ,
        description="Authenticates requests, rate limits, routes to inference",
        technologies=["Kong", "NGINX"],
    ))
 
    tm.add_element(DFDElement(
        id="vllm", name="vLLM Inference Server",
        element_type=ElementType.PROCESS,
        trust_zone=TrustZone.INFERENCE,
        description="GPU-accelerated LLM inference with PagedAttention",
        technologies=["vLLM", "CUDA", "PyTorch"],
    ))
 
    tm.add_element(DFDElement(
        id="pipeline", name="Model Deployment Pipeline",
        element_type=ElementType.PROCESS,
        trust_zone=TrustZone.CONTROL,
        description="Automated pipeline that loads new model versions",
        technologies=["Argo CD", "Helm"],
    ))
 
    # Datastores
    tm.add_element(DFDElement(
        id="s3_models", name="S3 Model Artifact Store",
        element_type=ElementType.DATA_STORE,
        trust_zone=TrustZone.REGISTRY,
        description="Stores trained model weights and configurations",
        technologies=["AWS S3", "SSE-KMS"],
    ))
 
    tm.add_element(DFDElement(
        id="mlflow", name="MLflow Model Registry",
        element_type=ElementType.DATA_STORE,
        trust_zone=TrustZone.REGISTRY,
        description="Tracks model versions and deployment stages",
        technologies=["MLflow", "PostgreSQL"],
    ))
 
    tm.add_element(DFDElement(
        id="logs", name="Inference Log Store",
        element_type=ElementType.DATA_STORE,
        trust_zone=TrustZone.MONITORING,
        description="Stores inference requests and responses for monitoring",
        technologies=["Elasticsearch", "S3"],
    ))
 
    # Dataflows
    tm.add_data_flow(DataFlow(
        id="df1", source_id="user", target_id="gateway",
        data_description="Inference requests (prompts)",
        protocol="HTTPS",
        is_encrypted=True, is_authenticated=True,
    ))
 
    tm.add_data_flow(DataFlow(
        id="df2", source_id="gateway", target_id="vllm",
        data_description="Authenticated inference requests",
        protocol="HTTP/gRPC",
        is_encrypted=False,  # Veelvoorkomend gat: intern verkeer onversleuteld
        is_authenticated=False,
    ))
 
    tm.add_data_flow(DataFlow(
        id="df3", source_id="s3_models", target_id="vllm",
        data_description="Model weights (multi-GB)",
        protocol="HTTPS (S3 API)",
        is_encrypted=True, is_authenticated=True,
    ))
 
    tm.add_data_flow(DataFlow(
        id="df4", source_id="admin", target_id="mlflow",
        data_description="Model stage transitions",
        protocol="HTTPS",
        is_encrypted=True, is_authenticated=True,
    ))
 
    tm.add_data_flow(DataFlow(
        id="df5", source_id="mlflow", target_id="pipeline",
        data_description="Model deployment triggers",
        protocol="Webhook HTTP",
        is_encrypted=False,
        is_authenticated=False,
    ))
 
    tm.add_data_flow(DataFlow(
        id="df6", source_id="vllm", target_id="logs",
        data_description="Inference requests and responses",
        protocol="HTTP",
        is_encrypted=False,
        is_authenticated=False,
    ))
 
    return tm
 
if __name__ == "__main__":
    tm = create_llm_inference_threat_model()
    threats = tm.apply_stride()
    report = tm.generate_report()
    print(report)
    print(f"\nTotal threats identified: {len(threats)}")
    print(f"AI-specific threats: {sum(1 for t in threats if t.get('ai_specific'))}")

"""
DREAD risk scoring adapted for AI infrastructure threats.
Adds AI-specific weighting factors to standard DREAD scoring.
"""
 
from dataclasses import dataclass
from typing import Optional
 
@dataclass
class AIThreatScore:
    """DREAD score for an AI infrastructure threat."""
    threat_name: str
    atlas_technique: Optional[str]
 
    # Standaard DREAD-dimensies (schaal 1-10)
    damage: int           # Hoe erg is het als de aanval slaagt?
    reproducibility: int  # Hoe makkelijk is het te reproduceren?
    exploitability: int   # Hoe makkelijk is het om de aanval uit te voeren?
    affected_users: int   # Hoeveel gebruikers worden getroffen?
    discoverability: int  # Hoe makkelijk is het om de kwetsbaarheid te vinden?
 
    # AI-specifieke aanpassingsfactoren
    model_impact: float = 1.0       # Beïnvloedt dit het modelgedrag? (1.0-2.0)
    data_sensitivity: float = 1.0   # Is trainings-/inferentiedata gevoelig? (1.0-2.0)
    detection_difficulty: float = 1.0  # Hoe moeilijk is het AI-specifieke aspect te detecteren? (1.0-2.0)
 
    @property
    def base_dread_score(self) -> float:
        """Calculate standard DREAD score (0-10)."""
        return (
            self.damage + self.reproducibility + self.exploitability
            + self.affected_users + self.discoverability
        ) / 5.0
 
    @property
    def ai_adjusted_score(self) -> float:
        """
        DREAD score adjusted for AI-specific factors.
        AI attacks that corrupt model behavior or are hard to
        detect get higher priority.
        """
        base = self.base_dread_score
        ai_multiplier = (
            self.model_impact * self.data_sensitivity * self.detection_difficulty
        ) ** (1/3)  # Geometrisch gemiddelde om extreme waarden te voorkomen
        return min(10.0, base * ai_multiplier)
 
    @property
    def priority(self) -> str:
        """Human-readable priority based on AI-adjusted score."""
        score = self.ai_adjusted_score
        if score >= 8.0:
            return "CRITICAL"
        elif score >= 6.0:
            return "HIGH"
        elif score >= 4.0:
            return "MEDIUM"
        elif score >= 2.0:
            return "LOW"
        return "INFO"
 
# Voorbeeld: Scoor veelvoorkomende dreigingen voor AI-infrastructuur
AI_THREAT_SCORES = [
    AIThreatScore(
        threat_name="Training data poisoning via S3 write access",
        atlas_technique="AML.T0020",
        damage=9, reproducibility=8, exploitability=6,
        affected_users=10, discoverability=4,
        model_impact=2.0, data_sensitivity=1.8, detection_difficulty=1.8,
    ),
    AIThreatScore(
        threat_name="Model extraction through inference API",
        atlas_technique="AML.T0024",
        damage=8, reproducibility=9, exploitability=7,
        affected_users=8, discoverability=6,
        model_impact=1.5, data_sensitivity=1.3, detection_difficulty=1.5,
    ),
    AIThreatScore(
        threat_name="RCE via TorchServe management API",
        atlas_technique="AML.T0010",
        damage=10, reproducibility=9, exploitability=8,
        affected_users=10, discoverability=7,
        model_impact=1.0, data_sensitivity=1.0, detection_difficulty=1.0,
    ),
    AIThreatScore(
        threat_name="GPU memory side-channel between tenants",
        atlas_technique="AML.T0024",
        damage=7, reproducibility=5, exploitability=4,
        affected_users=6, discoverability=3,
        model_impact=1.5, data_sensitivity=1.7, detection_difficulty=2.0,
    ),
]
 
def generate_risk_report(threats: list[AIThreatScore]) -> str:
    """Generate a prioritized risk report from scored threats."""
    sorted_threats = sorted(
        threats, key=lambda t: t.ai_adjusted_score, reverse=True,
    )
 
    lines = ["AI Infrastructure Threat Risk Report", "=" * 45, ""]
    for t in sorted_threats:
        lines.append(
            f"[{t.priority}] {t.threat_name}"
        )
        lines.append(
            f"  Base DREAD: {t.base_dread_score:.1f} | "
            f"AI-Adjusted: {t.ai_adjusted_score:.1f}"
        )
        if t.atlas_technique:
            lines.append(f"  ATLAS: {t.atlas_technique}")
        lines.append("")
 
    return "\n".join(lines)
 
if __name__ == "__main__":
    print(generate_risk_report(AI_THREAT_SCORES))