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

# Script Name: dvmcp
# Description: Damn Vulnerable MCP Server launcher
# Usage: dvmcp start|stop|status|logs|build

VERSION="1.1.0"

# Colors
readonly RED='\033[0;31m'
readonly GREEN='\033[0;32m'
readonly YELLOW='\033[1;33m'
readonly CYAN='\033[0;36m'
readonly BOLD='\033[1m'
readonly NC='\033[0m'

# Container settings
CONTAINER_NAME="dvmcp"
IMAGE="dvmcp:latest"
PORT_RANGE="9001-9010"

# Config file paths
SETTINGS_DIR="${HOME}/.claude"
PROD_CONFIG="${SETTINGS_DIR}/settings.json"
CTF_CONFIG="${SETTINGS_DIR}/settings-ctf.json"
BACKUP_CONFIG="${SETTINGS_DIR}/settings-prod-backup.json"
CONFIG_STATE_FILE="${SETTINGS_DIR}/.dvmcp-config-state"

show_help() {
  echo -e "${BOLD}dvmcp${NC} - Damn Vulnerable MCP Server Launcher v${VERSION}"
  echo
  echo -e "${BOLD}USAGE:${NC}"
  echo "  dvmcp <command>"
  echo
  echo -e "${BOLD}COMMANDS:${NC}"
  echo -e "  ${CYAN}build${NC}      Build Docker image from Dockerfile"
  echo -e "  ${CYAN}start${NC}      Start DVMCP server"
  echo -e "  ${CYAN}stop${NC}       Stop DVMCP server"
  echo -e "  ${CYAN}restart${NC}    Restart DVMCP server"
  echo -e "  ${CYAN}status${NC}     Check if running"
  echo -e "  ${CYAN}logs${NC}       Show container logs"
  echo -e "  ${CYAN}shell${NC}      Open shell in container"
  echo
  echo -e "${BOLD}EXAMPLES:${NC}"
  echo "  dvmcp build    # Build image (first time setup)"
  echo "  dvmcp start    # Launch DVMCP server"
  echo "  dvmcp stop     # Stop DVMCP server"
  echo "  dvmcp logs     # View logs"
  echo
  echo -e "${BOLD}ACCESS:${NC}"
  echo "  Ports: ${BOLD}9001-9010${NC} (10 challenge instances)"
  echo "  Test: ${BOLD}curl http://localhost:9001${NC}"
  echo
  echo -e "${BOLD}ABOUT:${NC}"
  echo "  Damn Vulnerable MCP Server - Intentionally vulnerable MCP implementation"
  echo "  Perfect for testing MCP security vulnerabilities"
  echo "  GitHub: https://github.com/harishsg993010/damn-vulnerable-MCP-server"
  echo
  echo -e "${BOLD}SECURITY:${NC}"
  echo "  This script automatically manages Claude Code config isolation:"
  echo "  - ${GREEN}start${NC}: Backs up production config, loads CTF-only config"
  echo "  - ${GREEN}stop${NC}:  Restores production config automatically"
  echo "  - ${YELLOW}Your production MCP servers are protected${NC}"
  echo
  echo -e "${BOLD}FIRST TIME SETUP:${NC}"
  echo "  1. Clone repo: git clone https://github.com/harishsg993010/damn-vulnerable-MCP-server.git"
  echo "  2. cd damn-vulnerable-MCP-server/"
  echo "  3. Build: dvmcp build"
  echo "  4. Start: dvmcp start"
  echo "  5. ${BOLD}Restart Claude Code${NC} to load CTF config"
}

check_docker() {
  if ! command -v docker &>/dev/null; then
    echo -e "${RED}Error:${NC} Docker not installed"
    exit 1
  fi
}

check_image_exists() {
  if ! docker images --format '{{.Repository}}:{{.Tag}}' | grep -q "^${IMAGE}$"; then
    return 1
  fi
  return 0
}

create_ctf_config_template() {
  if [[ ! -f "$CTF_CONFIG" ]]; then
    echo -e "${CYAN}[*]${NC} Creating CTF config template at ${CTF_CONFIG}..."
    cat > "$CTF_CONFIG" <<'EOF'
{
  "mcpServers": {
    "Challenge 1": {
      "command": "npx",
      "args": ["mcp-remote", "http://127.0.0.1:9001/sse"]
    },
    "Challenge 2": {
      "command": "npx",
      "args": ["mcp-remote", "http://127.0.0.1:9002/sse"]
    },
    "Challenge 3": {
      "command": "npx",
      "args": ["mcp-remote", "http://127.0.0.1:9003/sse"]
    },
    "Challenge 4": {
      "command": "npx",
      "args": ["mcp-remote", "http://127.0.0.1:9004/sse"]
    },
    "Challenge 5": {
      "command": "npx",
      "args": ["mcp-remote", "http://127.0.0.1:9005/sse"]
    },
    "Challenge 6": {
      "command": "npx",
      "args": ["mcp-remote", "http://127.0.0.1:9006/sse"]
    },
    "Challenge 7": {
      "command": "npx",
      "args": ["mcp-remote", "http://127.0.0.1:9007/sse"]
    },
    "Challenge 8": {
      "command": "npx",
      "args": ["mcp-remote", "http://127.0.0.1:9008/sse"]
    },
    "Challenge 9": {
      "command": "npx",
      "args": ["mcp-remote", "http://127.0.0.1:9009/sse"]
    },
    "Challenge 10": {
      "command": "npx",
      "args": ["mcp-remote", "http://127.0.0.1:9010/sse"]
    }
  }
}
EOF
    echo -e "${GREEN}✓${NC} CTF config created"
  fi
}

swap_to_ctf_config() {
  if [[ ! -f "$PROD_CONFIG" ]]; then
    echo -e "${RED}Error:${NC} Production config not found at ${PROD_CONFIG}"
    exit 1
  fi

  # Check if we're already in CTF mode
  if [[ -f "$CONFIG_STATE_FILE" ]]; then
    echo -e "${YELLOW}⚠${NC} Already in CTF mode"
    return 0
  fi

  # Create CTF config if it doesn't exist
  create_ctf_config_template

  echo -e "${CYAN}[*]${NC} Backing up production config..."
  cp "$PROD_CONFIG" "$BACKUP_CONFIG"

  echo -e "${CYAN}[*]${NC} Switching to CTF config..."
  cp "$CTF_CONFIG" "$PROD_CONFIG"

  # Mark that we're in CTF mode
  echo "CTF_MODE_ACTIVE" > "$CONFIG_STATE_FILE"

  echo -e "${GREEN}✓${NC} Switched to CTF config"
  echo -e "${YELLOW}Note:${NC} Production MCP servers are disabled until you stop DVMCP"
}

restore_prod_config() {
  if [[ ! -f "$CONFIG_STATE_FILE" ]]; then
    echo -e "${YELLOW}⚠${NC} Already using production config"
    return 0
  fi

  if [[ ! -f "$BACKUP_CONFIG" ]]; then
    echo -e "${RED}Error:${NC} Backup config not found at ${BACKUP_CONFIG}"
    echo -e "${YELLOW}Warning:${NC} Cannot restore production config!"
    exit 1
  fi

  echo -e "${CYAN}[*]${NC} Restoring production config..."
  cp "$BACKUP_CONFIG" "$PROD_CONFIG"

  # Clean up backup and state file
  rm -f "$BACKUP_CONFIG" "$CONFIG_STATE_FILE"

  echo -e "${GREEN}✓${NC} Production config restored"
}

cleanup_on_exit() {
  # If script exits unexpectedly while in CTF mode, restore production config
  if [[ -f "$CONFIG_STATE_FILE" ]]; then
    echo
    echo -e "${YELLOW}[!]${NC} Script interrupted, restoring production config..."
    restore_prod_config
  fi
}

# Register cleanup trap
trap cleanup_on_exit EXIT INT TERM

build_dvmcp() {
  # Check if Dockerfile exists in current directory
  if [[ ! -f "Dockerfile" ]]; then
    echo -e "${RED}Error:${NC} Dockerfile not found in current directory"
    echo
    echo -e "${YELLOW}Expected location:${NC} ./Dockerfile"
    echo
    echo -e "${CYAN}To build the image:${NC}"
    echo "  1. cd ~/path/to/damn-vulnerable-MCP-server/"
    echo "  2. Run: dvmcp build"
    echo
    echo -e "${CYAN}Or clone the repo:${NC}"
    echo "  git clone https://github.com/harishsg993010/damn-vulnerable-MCP-server.git"
    exit 1
  fi

  echo -e "${CYAN}[*]${NC} Building DVMCP Docker image..."
  echo -e "${YELLOW}Note:${NC} This may take several minutes..."

  if docker build -t "$IMAGE" . ; then
    echo -e "${GREEN}✓${NC} DVMCP image built successfully"
    echo -e "${CYAN}[*]${NC} You can now run: ${BOLD}dvmcp start${NC}"
  else
    echo -e "${RED}✗${NC} Build failed"
    exit 1
  fi
}

start_dvmcp() {
  # Check if image exists
  if ! check_image_exists; then
    echo -e "${RED}Error:${NC} DVMCP image not found"
    echo -e "${YELLOW}Run:${NC} ${BOLD}dvmcp build${NC} first"
    exit 1
  fi

  # Swap to CTF config BEFORE starting container
  swap_to_ctf_config
  echo

  if docker ps -a --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
    if docker ps --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
      echo -e "${YELLOW}⚠${NC} DVMCP already running"
      echo -e "${CYAN}[*]${NC} Access at: ${BOLD}http://localhost:9001-9010${NC}"
      return 0
    else
      echo -e "${CYAN}[*]${NC} Starting existing container..."
      docker start "$CONTAINER_NAME"
    fi
  else
    echo -e "${CYAN}[*]${NC} Starting DVMCP server..."
    echo -e "${YELLOW}Note:${NC} Mapping ports 9001-9010 for challenge instances"
    docker run -d --name "$CONTAINER_NAME" \
      -p 9001:9001 \
      -p 9002:9002 \
      -p 9003:9003 \
      -p 9004:9004 \
      -p 9005:9005 \
      -p 9006:9006 \
      -p 9007:9007 \
      -p 9008:9008 \
      -p 9009:9009 \
      -p 9010:9010 \
      "$IMAGE"
  fi

  sleep 3

  echo -e "${GREEN}✓${NC} DVMCP server started"
  echo
  echo -e "${BOLD}Challenge Instances:${NC}"
  echo -e "  ${CYAN}Port 9001:${NC} http://localhost:9001"
  echo -e "  ${CYAN}Port 9002:${NC} http://localhost:9002"
  echo -e "  ${CYAN}Port 9003:${NC} http://localhost:9003"
  echo -e "  ${CYAN}Port 9004:${NC} http://localhost:9004"
  echo -e "  ${CYAN}Port 9005:${NC} http://localhost:9005"
  echo -e "  ${CYAN}Port 9006:${NC} http://localhost:9006"
  echo -e "  ${CYAN}Port 9007:${NC} http://localhost:9007"
  echo -e "  ${CYAN}Port 9008:${NC} http://localhost:9008"
  echo -e "  ${CYAN}Port 9009:${NC} http://localhost:9009"
  echo -e "  ${CYAN}Port 9010:${NC} http://localhost:9010"
  echo
  echo -e "${YELLOW}Tip:${NC} Test with: ${BOLD}curl http://localhost:9001${NC}"
  echo -e "${YELLOW}Security:${NC} ${GREEN}Production MCP servers isolated${NC}"
  echo
  echo -e "${BOLD}⚠ IMPORTANT:${NC} Restart Claude Code to load CTF config"
}

stop_dvmcp() {
  if docker ps --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
    echo -e "${CYAN}[*]${NC} Stopping DVMCP server..."
    docker stop "$CONTAINER_NAME"
    echo -e "${GREEN}✓${NC} DVMCP server stopped"
  else
    echo -e "${YELLOW}⚠${NC} DVMCP not running"
  fi

  # Restore production config AFTER stopping container
  echo
  restore_prod_config
  echo
  echo -e "${BOLD}⚠ IMPORTANT:${NC} Restart Claude Code to load production config"
}

restart_dvmcp() {
  stop_dvmcp
  sleep 2
  start_dvmcp
}

show_status() {
  if docker ps --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
    echo -e "${GREEN}●${NC} DVMCP server is ${GREEN}running${NC}"
    echo
    echo -e "${BOLD}Active Ports:${NC}"
    docker port "$CONTAINER_NAME" 2>/dev/null | while IFS= read -r line; do
      echo -e "  ${CYAN}${line}${NC}"
    done
    echo
    docker ps --format "table {{.Names}}\t{{.Status}}\t{{.Ports}}" | grep -E "(NAMES|${CONTAINER_NAME})"
  else
    echo -e "${RED}●${NC} DVMCP server is ${RED}stopped${NC}"

    if ! check_image_exists; then
      echo
      echo -e "${YELLOW}Note:${NC} Image not built yet. Run: ${BOLD}dvmcp build${NC}"
    fi
  fi

  echo
  echo -e "${BOLD}Config Status:${NC}"
  if [[ -f "$CONFIG_STATE_FILE" ]]; then
    echo -e "  ${CYAN}Mode:${NC} CTF (isolated)"
    echo -e "  ${CYAN}Active:${NC} ${CTF_CONFIG}"
    echo -e "  ${CYAN}Backup:${NC} ${BACKUP_CONFIG}"
  else
    echo -e "  ${CYAN}Mode:${NC} Production"
    echo -e "  ${CYAN}Active:${NC} ${PROD_CONFIG}"
  fi
}

show_logs() {
  if docker ps --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
    docker logs -f "$CONTAINER_NAME"
  else
    echo -e "${RED}Error:${NC} DVMCP server not running"
    exit 1
  fi
}

open_shell() {
  if docker ps --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
    echo -e "${CYAN}[*]${NC} Opening shell in DVMCP container..."
    docker exec -it "$CONTAINER_NAME" /bin/bash || docker exec -it "$CONTAINER_NAME" /bin/sh
  else
    echo -e "${RED}Error:${NC} DVMCP server not running"
    exit 1
  fi
}

# Main
check_docker

if [[ $# -eq 0 ]] || [[ "$1" =~ ^(-h|--help|help)$ ]]; then
  show_help
  exit 0
fi

case "$1" in
  build)
    build_dvmcp
    ;;
  start|up)
    start_dvmcp
    ;;
  stop|down)
    stop_dvmcp
    ;;
  restart)
    restart_dvmcp
    ;;
  status)
    show_status
    ;;
  logs)
    show_logs
    ;;
  shell|sh|bash)
    open_shell
    ;;
  *)
    echo -e "${RED}Error:${NC} Unknown command: $1"
    echo "Run 'dvmcp --help' for usage"
    exit 1
    ;;
esac