AI 基礎設施的零信任架構

"""
SPIFFE-based workload identity verification for AI pipeline components.
Demonstrates how to verify that a training job, model registry, or
inference endpoint has a valid identity before allowing access.
"""
 
import ssl
import socket
import json
from dataclasses import dataclass
from typing import Optional
from urllib.parse import urlparse
from cryptography import x509
from cryptography.x509.oid import ExtensionOID, NameOID
from cryptography.hazmat.primitives import hashes
 
@dataclass
class WorkloadIdentity:
    """Parsed SPIFFE identity from an X.509 SVID."""
    spiffe_id: str
    trust_domain: str
    workload_path: str
    certificate_hash: str
    not_valid_after: str
 
    @property
    def component_type(self) -> str:
        """Extract AI component type from SPIFFE path."""
        parts = self.workload_path.strip("/").split("/")
        if parts:
            return parts[0]  # e.g., "training", "inference", "registry"
        return "unknown"
 
def extract_spiffe_id(cert: x509.Certificate) -> Optional[WorkloadIdentity]:
    """
    Extract SPIFFE ID from X.509 certificate SAN extension.
    SPIFFE IDs are encoded as URI SANs: spiffe://<trust-domain>/<workload>
    """
    try:
        san_ext = cert.extensions.get_extension_for_oid(
            ExtensionOID.SUBJECT_ALTERNATIVE_NAME
        )
        for uri in san_ext.value.get_values_for_type(
            x509.UniformResourceIdentifier
        ):
            if uri.startswith("spiffe://"):
                parsed = urlparse(uri)
                return WorkloadIdentity(
                    spiffe_id=uri,
                    trust_domain=parsed.hostname or "",
                    workload_path=parsed.path,
                    certificate_hash=cert.fingerprint(hashes.SHA256()).hex(),
                    not_valid_after=str(cert.not_valid_after_utc),
                )
    except x509.ExtensionNotFound:
        return None
    return None
 
class AIZeroTrustVerifier:
    """Enforce zero trust access policies for AI components."""
 
    ACCESS_POLICIES = {
        "training": {
            "allowed_targets": [
                "data-store", "registry",
                "experiment-tracker", "feature-store",
            ],
            "denied_targets": ["inference", "monitoring-admin"],
        },
        "inference": {
            "allowed_targets": ["registry", "feature-store"],
            "denied_targets": ["training", "data-store"],
        },
        "registry": {
            "allowed_targets": ["data-store"],
            "denied_targets": ["training", "inference"],
        },
        "pipeline": {
            "allowed_targets": [
                "training", "registry",
                "data-store", "feature-store",
            ],
            "denied_targets": ["inference"],
        },
    }
 
    def __init__(self, trust_domain: str):
        self.trust_domain = trust_domain
 
    def verify_access(
        self, source: WorkloadIdentity, target_component: str,
    ) -> tuple[bool, str]:
        if source.trust_domain != self.trust_domain:
            return False, f"Trust domain mismatch: {source.trust_domain}"
        policy = self.ACCESS_POLICIES.get(source.component_type)
        if policy is None:
            return False, f"No policy for {source.component_type}"
        if target_component in policy.get("denied_targets", []):
            return False, f"Denied: {source.component_type} -> {target_component}"
        if target_component in policy.get("allowed_targets", []):
            return True, "Access permitted by policy"
        return False, "No explicit allow (default deny)"

"""Audit Kubernetes network policies and pod configuration for zero
trust gaps in GPU cluster interconnects."""
 
import subprocess
import json
 
def audit_kubernetes_network_policies(
    namespace: str = "ml-platform",
) -> list[dict]:
    findings = []
    pods = json.loads(subprocess.run(
        ["kubectl", "get", "pods", "-n", namespace, "-o", "json"],
        capture_output=True, text=True, timeout=30,
    ).stdout)
    policies = json.loads(subprocess.run(
        ["kubectl", "get", "networkpolicies", "-n", namespace, "-o", "json"],
        capture_output=True, text=True, timeout=30,
    ).stdout)
 
    # 1. Missing default-deny is a HIGH severity finding:
    #    without it, all pods communicate freely.
    has_default_deny = any(
        p["metadata"]["name"].startswith("default-deny")
        for p in policies.get("items", [])
    )
    if not has_default_deny:
        findings.append({"severity": "HIGH",
                         "title": f"No default-deny in {namespace}"})
 
    # 2. Host networking bypasses all NetworkPolicies; flag HIGH.
    # 3. Privileged containers (common for GPU workloads) escape isolation.
    # 4. RDMA/InfiniBand device requests bypass kernel networking entirely.
    for pod in pods.get("items", []):
        spec = pod.get("spec", {})
        if spec.get("hostNetwork"):
            findings.append({"severity": "HIGH",
                             "title": f"hostNetwork: {pod['metadata']['name']}"})
        for c in spec.get("containers", []):
            if c.get("securityContext", {}).get("privileged"):
                findings.append({"severity": "HIGH",
                                 "title": f"Privileged: {c['name']}"})
            limits = {**c.get("resources", {}).get("limits", {}),
                      **c.get("resources", {}).get("requests", {})}
            for r in limits:
                if "rdma" in r or "infiniband" in r:
                    findings.append({"severity": "MEDIUM",
                                     "title": f"RDMA device: {c['name']}"})
    return findings
 
def check_service_mesh_coverage(namespace: str) -> list[dict]:
    """Verify Istio/Linkerd sidecars exist and PeerAuthentication is STRICT."""
    findings = []
    # Check each pod for istio-proxy / linkerd-proxy / envoy sidecar
    # and check for STRICT mTLS PeerAuthentication in the namespace.
    # PERMISSIVE mTLS allows plaintext — treat as HIGH severity.
    return findings

"""Audit ML pipeline artifacts for trust boundary violations
across stages (Kubeflow, Airflow, custom systems)."""
 
import hashlib
import json
from pathlib import Path
 
def audit_pipeline_artifacts(pipeline_run_dir: str) -> list[dict]:
    findings = []
    run_path = Path(pipeline_run_dir)
    if not run_path.exists():
        return [{"severity": "ERROR", "title": "run dir missing"}]
 
    # Walk all metadata.json / artifact_info.json produced by stages.
    for meta_file in (list(run_path.rglob("**/metadata.json"))
                      + list(run_path.rglob("**/artifact_info.json"))):
        try:
            metadata = json.loads(meta_file.read_text())
        except (json.JSONDecodeError, IOError):
            continue
        stage_name = metadata.get("stage", meta_file.parent.name)
        for artifact in metadata.get("output_artifacts", []):
            art_path = artifact.get("path", "")
            has_hash = "sha256" in artifact or "hash" in artifact
            has_sig = "signature" in artifact
            if not has_hash:
                findings.append({
                    "severity": "HIGH",
                    "title": f"No hash: {stage_name}/{art_path}",
                    "detail": "Upstream stage can substitute poisoned model.",
                })
            if not has_sig:
                findings.append({
                    "severity": "MEDIUM",
                    "title": f"No signature: {stage_name}/{art_path}",
                })
            # If hash exists, recompute and compare — mismatch is CRITICAL.
            if has_hash:
                expected = artifact.get("sha256") or artifact.get("hash")
                full = run_path / art_path
                if full.exists():
                    actual = hashlib.sha256(full.read_bytes()).hexdigest()
                    if actual != expected:
                        findings.append({"severity": "CRITICAL",
                                         "title": "Hash mismatch"})
 
    # Credentials accidentally stored in artifact directories are a
    # HIGH-severity finding: secrets persist into artifact storage.
    for env_file in (list(run_path.rglob("**/.env"))
                     + list(run_path.rglob("**/secrets.*"))):
        findings.append({"severity": "HIGH",
                         "title": f"Credentials in artifacts: {env_file}"})
    return findings

"""Audit Kubernetes AI workloads for credential-related zero trust
violations: excessive scope, long validity, hardcoded secrets."""
 
import subprocess
import json
 
def audit_ai_credentials(namespace: str = "ml-platform") -> list[dict]:
    findings = []
    sas = json.loads(subprocess.run(
        ["kubectl", "get", "serviceaccounts", "-n", namespace, "-o", "json"],
        capture_output=True, text=True, timeout=30,
    ).stdout)
 
    for sa in sas.get("items", []):
        # Bound static secrets -> MEDIUM: prefer projected tokens.
        if sa.get("secrets"):
            findings.append({"severity": "MEDIUM",
                             "title": f"Static secrets bound: "
                                      f"{sa['metadata']['name']}"})
        # Workload identity annotations (GKE / EKS IRSA) are informational;
        # the operator must still verify IAM scope is minimal.
        ann = sa.get("metadata", {}).get("annotations", {})
        if ann.get("iam.gke.io/gcp-service-account"):
            findings.append({"severity": "INFO", "title": "GKE WI binding"})
        if ann.get("eks.amazonaws.com/role-arn"):
            findings.append({"severity": "INFO", "title": "EKS IRSA binding"})
 
    pods = json.loads(subprocess.run(
        ["kubectl", "get", "pods", "-n", namespace, "-o", "json"],
        capture_output=True, text=True, timeout=30,
    ).stdout)
    patterns = ["KEY", "SECRET", "PASSWORD", "TOKEN", "CREDENTIAL"]
    for pod in pods.get("items", []):
        for c in pod.get("spec", {}).get("containers", []):
            for env in c.get("env", []):
                name = env.get("name", "").upper()
                if any(p in name for p in patterns):
                    # Hardcoded plaintext value = CRITICAL
                    # Referenced via secretKeyRef = LOW (verify rotation)
                    if env.get("value"):
                        findings.append({"severity": "CRITICAL",
                                         "title": f"Hardcoded env: {name}"})
                    elif "secretKeyRef" in env.get("valueFrom", {}):
                        findings.append({"severity": "LOW",
                                         "title": f"Secret-backed env: {name}"})
    return findings

#!/usr/bin/env bash
# Zero trust compliance audit for AI infrastructure on Kubernetes
set -euo pipefail
 
NAMESPACE="${1:-ml-platform}"
echo "=== Zero Trust Audit: Namespace $NAMESPACE ==="
 
# 1. Default-deny NetworkPolicy must exist
DENY=$(kubectl get networkpolicies -n "$NAMESPACE" \
    -o jsonpath='{range .items[*]}{.metadata.name}{"\n"}{end}' \
    | grep -c "default-deny" || true)
[ "$DENY" -eq 0 ] && echo "[FAIL] No default-deny policy" \
                  || echo "[PASS] Default-deny exists"
 
# 2. PeerAuthentication mTLS mode must be STRICT
STRICT=$(kubectl get peerauthentication -n "$NAMESPACE" \
    -o jsonpath='{range .items[*]}{.spec.mtls.mode}{"\n"}{end}' \
    | grep -c "STRICT" || true)
[ "$STRICT" -eq 0 ] && echo "[FAIL] No STRICT mTLS" \
                    || echo "[PASS] STRICT mTLS enforced"
 
# 3. Legacy non-expiring service account token volumes
# 4. Privileged containers and root UID (zero trust violation)
# 5. LoadBalancer / NodePort services expose workloads externally
# (see full script for jq-style Python parsers)