運用 STRIDE 對 AI 基礎設施進行威脅建模

"""
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, end users
    DMZ = "dmz"                       # API gateways, load balancers
    INFERENCE = "inference"           # Model serving
    TRAINING = "training"            # Training cluster
    DATA = "data"                    # Training data storage
    CONTROL = "control"              # Pipeline orchestration
    MONITORING = "monitoring"        # Observability
    REGISTRY = "registry"           # Model/artifact registry
 
@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:
        self.elements[element.id] = element
 
    def add_data_flow(self, flow: DataFlow) -> None:
        self.data_flows.append(flow)
 
    def apply_stride(self) -> list[dict]:
        """Systematically apply STRIDE to all elements and data flows."""
        threats = []
        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))
        for flow in self.data_flows:
            threats.extend(self._stride_data_flow(flow))
        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})"
        threats.append({
            "element": elem.id, "category": "Spoofing",
            "threat": f"An attacker could impersonate {prefix}.",
            "mitigation": "Implement mutual TLS with workload identity (SPIFFE).",
        })
        threats.append({
            "element": elem.id, "category": "Tampering",
            "threat": f"An attacker could modify the behavior of {prefix}.",
            "mitigation": "Sign artifacts, use read-only filesystems, verify inputs.",
        })
        # Additional categories: Repudiation, Information Disclosure,
        # Denial of Service, Elevation of Privilege...
        return threats
 
    # Similar helpers for data stores, external entities, data flows
    # and cross-boundary trust zone crossings follow the same pattern.
 
    def generate_report(self) -> str:
        """Generate a structured threat model report as JSON."""
        if not self.threats:
            self.apply_stride()
        # Aggregate threats by category and AI-specific flag, return JSON.
        return json.dumps({"system": self.system_name, "threats": self.threats}, indent=2)

"""Threat model for an LLM inference pipeline on Kubernetes GPUs."""
 
def create_llm_inference_threat_model() -> "AIThreatModel":
    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."
        ),
    )
 
    # External entities: API clients, ML engineers
    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"],
    ))
 
    # Processes: API gateway (Kong/NGINX in DMZ), vLLM inference
    # server (GPU inference zone), deployment pipeline (control zone)
    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"],
    ))
 
    # Data stores: S3 model artifacts, MLflow registry, log store
    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"],
    ))
 
    # Data flows with encryption/authentication flags:
    # - user -> gateway: HTTPS, encrypted, authenticated
    # - gateway -> vLLM: HTTP/gRPC, often unencrypted (common gap)
    # - s3_models -> vLLM: HTTPS, encrypted, authenticated
    # - mlflow -> pipeline: webhook, often unauthenticated (gap)
    # - vLLM -> logs: HTTP, unencrypted
    tm.add_data_flow(DataFlow(
        id="df2", source_id="gateway", target_id="vllm",
        data_description="Authenticated inference requests",
        protocol="HTTP/gRPC",
        is_encrypted=False,  # Common gap: internal traffic unencrypted
        is_authenticated=False,
    ))
    return tm

"""DREAD scoring adapted for AI infrastructure threats."""
 
@dataclass
class AIThreatScore:
    threat_name: str
    atlas_technique: Optional[str]
    # Standard DREAD dimensions (1-10)
    damage: int
    reproducibility: int
    exploitability: int
    affected_users: int
    discoverability: int
    # AI-specific adjustment factors (1.0-2.0)
    model_impact: float = 1.0
    data_sensitivity: float = 1.0
    detection_difficulty: float = 1.0
 
    @property
    def base_dread_score(self) -> float:
        return (self.damage + self.reproducibility + self.exploitability
                + self.affected_users + self.discoverability) / 5.0
 
    @property
    def ai_adjusted_score(self) -> float:
        """Geometric mean of AI factors prevents extreme values."""
        base = self.base_dread_score
        ai_mult = (self.model_impact * self.data_sensitivity
                   * self.detection_difficulty) ** (1/3)
        return min(10.0, base * ai_mult)
 
    @property
    def priority(self) -> str:
        s = self.ai_adjusted_score
        if s >= 8.0: return "CRITICAL"
        if s >= 6.0: return "HIGH"
        if s >= 4.0: return "MEDIUM"
        if s >= 2.0: return "LOW"
        return "INFO"