Patched code sample

import json
from authlib.jose import jws, jwk
from authlib.common.errors import InvalidSignatureError
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives import serialization

# This script demonstrates the fix for a JWK Header Injection vulnerability.
#
# The vulnerability (described in CVE-2026-27962, though the number may be a typo for a real CVE)
# occurs when a JWS deserialization function is called with `key=None`.
# In this state, 'authlib' prior to v1.0.0 would trust and use the public key
# provided in the 'jwk' header of the JWT itself.
#
# The fix is to ensure the application developer *never* calls the decode function
# with a null key. Instead, the developer must provide a specific, trusted
# public key to verify the signature against. This script shows both the
# vulnerable scenario and the corrected, secure implementation.

# --- 1. Setup: Create keys for the server and an attacker ---

# A function to generate RSA key pairs for demonstration
def generate_rsa_key_pair():
    private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    public_key = private_key.public_key()
    return private_key, public_key

# The legitimate server's keys. The server trusts its own keys.
SERVER_PRIVATE_KEY, SERVER_PUBLIC_KEY = generate_rsa_key_pair()

# The attacker's keys. The server should NOT trust these.
ATTACKER_PRIVATE_KEY, ATTACKER_PUBLIC_KEY = generate_rsa_key_pair()


# --- 2. Attacker: Craft a malicious JWT ---

# The attacker wants to forge a token with admin privileges.
payload = {'sub': 'user-123', 'iss': 'auth-provider', 'is_admin': True}

# The attacker converts their OWN public key into the JWK format.
attacker_public_jwk = jwk.dumps(ATTACKER_PUBLIC_KEY, kty='RSA')

# The attacker creates a protected header that includes their public key.
# This is the "JWK Header Injection".
protected_header = {
    'alg': 'RS256',
    'jwk': attacker_public_jwk
}

# The attacker signs the token using their OWN private key.
malicious_jws_token = jws.encode(
    protected_header,
    json.dumps(payload).encode('utf-8'),
    ATTACKER_PRIVATE_KEY
)

print("--- Attacker has crafted a malicious token ---")
print(f"Token: {malicious_jws_token.decode('utf-8')[:100]}...")
print("\n")


# --- 3. Vulnerable Implementation (The "Before" picture) ---
# A vulnerable server might try to decode the token without specifying a key.
# It incorrectly assumes the library will handle it securely.
print("--- VULNERABLE SERVER: Verifying token with `key=None` ---")
try:
    # VULNERABLE CALL: `key=None` causes authlib to use the key from the 'jwk' header.
    decoded_token = jws.decode(malicious_jws_token, key=None)
    
    print("\033[91m[VULNERABLE] SUCCESS: Token signature was VERIFIED!\033[0m")
    print("The server trusted the attacker's embedded public key.")
    print(f"Decoded Payload: {decoded_token.payload.decode('utf-8')}")

except InvalidSignatureError:
    print("\033[92m[VULNERABLE] FAILED: Token signature was invalid.\033[0m")
    print("This block should not be reached in the vulnerable scenario.")

print("\n" + "="*60 + "\n")


# --- 4. Fixed Implementation (The "After" picture / The Fix) ---
# The fix is at the application level: ALWAYS provide a trusted key for verification.
# The server must use a key it trusts (e.g., its own public key, or one from a
# trusted JWKS endpoint), not one from the token header.
print("--- FIXED SERVER: Verifying token with a trusted `key` ---")
try:
    # FIXED CALL: The server provides the public key it trusts.
    # It will try to verify the attacker's token against the SERVER's public key.
    # This will fail, because the token was signed with the ATTACKER's private key.
    decoded_token = jws.decode(malicious_jws_token, key=SERVER_PUBLIC_KEY)

    print("\033[91m[FIXED] SUCCESS: Token signature was VERIFIED!\033[0m")
    print("This block should not be reached. If it is, the fix failed.")

except InvalidSignatureError:
    print("\033[92m[FIXED] FAILED AS EXPECTED: InvalidSignatureError was raised.\033[0m")
    print("The server correctly rejected the token because the signature does not match the trusted server key.")
    print("The vulnerability is mitigated.")