#!/bin/bash
set -euo pipefail

# Verify the proposed SQL-injection fix for LiteLLM proxy combined_view token lookup.
# This script patches a vulnerable litellm==1.83.6 package with the proposed diff,
# inspects the code transformation, and runs a real LiteLLM proxy + PostgreSQL to prove
# that the SQL-injection payload is rejected after the patch.

ROOT="${PRUVA_ROOT:-$(cd "$(dirname "$0")/.." && pwd)}"
CODING_DIR="$ROOT/coding"
WORK_DIR="$CODING_DIR/verify_work"
LOG_DIR="$ROOT/logs"
PATCH_FILE="$CODING_DIR/proposed_fix.diff"
mkdir -p "$WORK_DIR" "$LOG_DIR"

cd "$ROOT"

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

# Load project cache if available.
CACHE_DIR="$ROOT/artifacts/litellm-cache"
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")"
fi
VENV="$CACHE_DIR/venvs/litellm_vuln_1_83_6"
PGDATA="$CACHE_DIR/pglitellm_repro"
DB_NAME="litellm_repro_vuln"
DB_USER="${USER:-$(whoami)}"

PROXY_PORT=""
PG_PID=""
PROXY_PID=""
RESULT="failed"
FIXED_TIME="0"
FIXED_CODE="000"
VALID_TIME="0"
VALID_CODE="000"

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
}

cleanup() {
    set +e
    if [[ -n "$PROXY_PID" ]]; then
        kill "$PROXY_PID" 2>/dev/null || true
        sleep 1
        kill -9 "$PROXY_PID" 2>/dev/null || true
    fi
    if [[ -f "$PGDATA/postmaster.pid" ]]; then
        PATH="/usr/lib/postgresql/18/bin:$PATH" pg_ctl -D "$PGDATA" stop -m fast >/dev/null 2>&1 || true
    fi
}
trap cleanup EXIT

# 1. Copy vulnerable litellm package and apply the proposed patch.
echo "[+] Copying vulnerable litellm package from $VENV"
rm -rf "$WORK_DIR/litellm"
cp -r "$VENV/lib/python3.14/site-packages/litellm" "$WORK_DIR/litellm"

echo "[+] Applying proposed patch to $WORK_DIR/litellm"
cd "$WORK_DIR"
if ! patch -p1 < "$PATCH_FILE"; then
    echo "[-] Patch did not apply cleanly"
    exit 2
fi

echo "[+] Verifying code transformation"
PATCHED_UTILS="$WORK_DIR/litellm/proxy/utils.py"
if grep -q "WHERE v.token = '{token}'" "$PATCHED_UTILS"; then
    echo "[-] Vulnerable raw SQL interpolation still present"
    exit 2
fi
if ! grep -q 'WHERE v.token = \$1' "$PATCHED_UTILS"; then
    echo "[-] Parameterized token lookup not found"
    exit 2
fi
if ! grep -A2 'async def _query_first_with_cached_plan_fallback(' "$PATCHED_UTILS" | grep -q '*args'; then
    echo "[-] _query_first_with_cached_plan_fallback signature not updated to accept *args"
    exit 2
fi
if ! grep -q 'self.db.query_first(sql_query, \*args)' "$PATCHED_UTILS"; then
    echo "[-] query_first not called with positional arguments"
    exit 2
fi
if ! grep -q 'sql_query, hashed_token' "$PATCHED_UTILS"; then
    echo "[-] hashed_token not passed as query parameter"
    exit 2
fi
echo "[+] Code transformation verified"

# 2. Start PostgreSQL from the cache data directory.
echo "[+] Starting PostgreSQL from $PGDATA"
export PATH="/usr/lib/postgresql/18/bin:$PATH"
if [[ -f "$PGDATA/postmaster.pid" ]]; then
    pg_ctl -D "$PGDATA" stop -m fast >/dev/null 2>&1 || true
    sleep 2
fi

PG_PORT="$(grep -E '^port' "$PGDATA/postgresql.conf" | tail -1 | sed -e 's/^port = //' -e 's/[^0-9]//g')"
if [[ -z "$PG_PORT" ]]; then
    PG_PORT="$(find_free_port)"
    sed -i '/^port = /d' "$PGDATA/postgresql.conf" 2>/dev/null || true
    echo "port = $PG_PORT" >> "$PGDATA/postgresql.conf"
else
    if ! is_port_free "$PG_PORT"; then
        PG_PORT="$(find_free_port)"
        sed -i '/^port = /d' "$PGDATA/postgresql.conf" 2>/dev/null || true
        echo "port = $PG_PORT" >> "$PGDATA/postgresql.conf"
    fi
fi

pg_ctl -D "$PGDATA" -l "$LOG_DIR/verify_postgres.log" start
PG_PID="$(cat "$PGDATA/postmaster.pid" 2>/dev/null | head -1 || 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 ready on 127.0.0.1:$PG_PORT"
        break
    fi
    sleep 1
    if [[ "$i" == 30 ]]; then
        echo "[-] PostgreSQL did not become ready"
        exit 2
    fi
done

# Ensure the database exists and has the schema pushed.
if ! psql -h 127.0.0.1 -p "$PG_PORT" -U "$DB_USER" -d postgres -Atc "SELECT 1 FROM pg_database WHERE datname='$DB_NAME'" | grep -q 1; then
    createdb -h 127.0.0.1 -p "$PG_PORT" -U "$DB_USER" "$DB_NAME"
fi

SCHEMA_FILE="$(find "$VENV" -path '*/site-packages/litellm/proxy/schema.prisma' -print -quit)"
if [[ -z "$SCHEMA_FILE" ]]; then
    echo "[-] schema.prisma not found"
    exit 2
fi

export DATABASE_URL="postgresql://$DB_USER@127.0.0.1:$PG_PORT/$DB_NAME"
if psql -h 127.0.0.1 -p "$PG_PORT" -U "$DB_USER" -d "$DB_NAME" -Atc 'SELECT to_regclass('"'"'"LiteLLM_VerificationToken"'"'"')' 2>/dev/null | grep -q 'LiteLLM_VerificationToken'; then
    echo "[+] Prisma schema already present in $DB_NAME"
else
    echo "[+] Pushing Prisma schema to $DB_NAME"
    PATH="$VENV/bin:$PATH" "$VENV/bin/prisma" db push --schema "$SCHEMA_FILE" --accept-data-loss >/dev/null
fi
psql -h 127.0.0.1 -p "$PG_PORT" -U "$DB_USER" -d "$DB_NAME" -c 'DELETE FROM "LiteLLM_VerificationToken";' >/dev/null 2>&1 || true

# 3. Write a minimal LiteLLM config.
CONFIG_FILE="$WORK_DIR/litellm_config.yaml"
cat > "$CONFIG_FILE" <<'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

# 4. Start the patched LiteLLM proxy.
PROXY_PORT="$(find_free_port)"
SCHEMA_DIR="$(find "$VENV" -path '*/site-packages/litellm/proxy' -type d -print -quit)"

echo "[+] Starting patched LiteLLM proxy on 127.0.0.1:$PROXY_PORT"
(
    export PATH="$VENV/bin:$PATH"
    export PYTHONOPTIMIZE=1
    export PYTHONUNBUFFERED=1
    export PYTHONPATH="$WORK_DIR"
    export DATABASE_URL="postgresql://$DB_USER@127.0.0.1:$PG_PORT/$DB_NAME"
    export LITELLM_MASTER_KEY="sk-verify-master-key"
    export OPENAI_API_KEY="sk-verify-fake-openai"
    cd "$SCHEMA_DIR"
    exec stdbuf -oL -eL litellm --config "$CONFIG_FILE" --host 127.0.0.1 --port "$PROXY_PORT" > "$LOG_DIR/verify_proxy.log" 2>&1
) &
PROXY_PID=$!

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

# 5. Seed a valid virtual key.
echo "[+] Generating a valid virtual key through the patched proxy"
curl -sS -X POST "http://127.0.0.1:$PROXY_PORT/key/generate" \
    -H "Authorization: Bearer sk-verify-master-key" \
    -H "Content-Type: application/json" \
    -d '{"models":["gpt-4o"],"metadata":{"purpose":"verify-fix"}}' \
    --max-time 20 \
    -o "$LOG_DIR/verify_key_generate.json"
VALID_KEY="$(jq -r '.key' "$LOG_DIR/verify_key_generate.json" 2>/dev/null || true)"
if [[ -z "$VALID_KEY" || "$VALID_KEY" == "null" ]]; then
    echo "[-] Failed to generate valid key"
    cat "$LOG_DIR/verify_key_generate.json"
    exit 2
fi

# 6. Test the SQL-injection payload.
PAYLOAD="' OR EXISTS(SELECT 1 FROM pg_sleep(3)) --"
request_with_token() {
    local port="$1"
    local token="$2"
    local outfile="$3"
    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 "[+] Sending SQL-injection payload (Authorization Bearer) to patched proxy"
METRICS="$(request_with_token "$PROXY_PORT" "$PAYLOAD" "$LOG_DIR/verify_injection.txt")"
FIXED_TIME="${METRICS%,*}"
FIXED_CODE="${METRICS#*,}"
echo "[+] Patched injection response: time=${FIXED_TIME}s status=${FIXED_CODE}"

echo "[+] Sending valid key request to patched proxy"
METRICS="$(request_with_token "$PROXY_PORT" "$VALID_KEY" "$LOG_DIR/verify_valid_key.txt")"
VALID_TIME="${METRICS%,*}"
VALID_CODE="${METRICS#*,}"
echo "[+] Valid key response: time=${VALID_TIME}s status=${VALID_CODE}"

# 7. Verify results.
python3 - "$FIXED_TIME" "$FIXED_CODE" "$VALID_TIME" "$VALID_CODE" "$LOG_DIR" <<'PY'
import json, os, sys
fixed_time, fixed_code, valid_time, valid_code = sys.argv[1:5]
logs = sys.argv[5]
with open(os.path.join(logs, "verify_injection.txt"), encoding="utf-8", errors="replace") as f:
    inj_body = f.read()
with open(os.path.join(logs, "verify_valid_key.txt"), encoding="utf-8", errors="replace") as f:
    valid_body = f.read()
with open(os.path.join(logs, "verify_postgres.log"), encoding="utf-8", errors="replace") as f:
    pglog = f.read()

checks = {
    "injection_rejected": fixed_code == "401" and "token_not_found_in_db" in inj_body,
    "injection_not_delayed": float(fixed_time) < 1.5,
    "valid_key_accepted": valid_code == "200" and '"data"' in valid_body,
    "pglog_no_injected_sleep": "pg_sleep(3)" not in pglog or "Parameters: $1" in pglog,
}
print(json.dumps(checks, indent=2))
if not all(checks.values()):
    sys.exit(1)
PY

RESULT="passed"
echo "[+] FIX VERIFIED: patched litellm==1.83.6 rejects the SQL-injection payload (HTTP 401, <1.5s) while still accepting a valid key."

# 8. Record results.
{
    echo "result=$RESULT"
    echo "injection_time=$FIXED_TIME"
    echo "injection_code=$FIXED_CODE"
    echo "valid_key_time=$VALID_TIME"
    echo "valid_key_code=$VALID_CODE"
    echo "proxy_port=$PROXY_PORT"
    echo "pg_port=$PG_PORT"
    echo "patched_package=$WORK_DIR/litellm"
    echo "patch_applied=$PATCH_FILE"
} > "$WORK_DIR/verify_result.txt"

cp "$PGDATA/log" "$LOG_DIR/verify_postgres.log" 2>/dev/null || true

echo "[+] Verification complete"
exit 0