Patched code sample

import base64
import json

# --- Custom Exception Classes for Clarity ---

class JWTError(Exception):
    """Base exception for all JWT-related errors in this simulation."""
    pass

class InvalidAlgorithmError(JWTError):
    """Raised when an algorithm is not allowed."""
    pass

class InvalidTokenError(JWTError):
    """Raised for general malformed token issues."""
    pass


# --- Helper Function for Base64 URL Decoding ---

def base64url_decode(input_str):
    """Decodes a base64url encoded string, adding padding if necessary."""
    rem = len(input_str) % 4
    if rem > 0:
        input_str += '=' * (4 - rem)
    return base64.urlsafe_b64decode(input_str)


# --- The Patched/Fixed JWT Decode Function ---

def fixed_jwt_decode(token, key, algorithms):
    """
    Decodes a JWT, but crucially, enforces that the token's algorithm
    MUST be present in the user-provided 'algorithms' list. This prevents
    the 'alg=none' vulnerability unless the developer explicitly allows it.

    This function simulates the fix for CVE-2025-61152.
    """
    if not isinstance(algorithms, list):
        raise TypeError("The 'algorithms' parameter must be a list of strings.")

    try:
        header_b64, payload_b64, signature_b64 = token.split('.')
    except ValueError:
        raise InvalidTokenError("Invalid token structure. Expected three parts separated by dots.")

    # Decode the header to find out the algorithm used
    try:
        header_data = json.loads(base64url_decode(header_b64))
        token_alg = header_data.get('alg')
    except Exception:
        raise InvalidTokenError("Invalid header. Could not decode or parse JSON.")

    # --- THE FIX ---
    # The core of the patch is this check. It ensures that the algorithm
    # specified in the token's header is one of the algorithms the developer
    # has explicitly stated they are willing to accept.
    if token_alg not in algorithms:
        raise InvalidAlgorithmError(f"The algorithm '{token_alg}' is not in the allowed list: {algorithms}")
    # --- END OF FIX ---

    # If 'none' was explicitly allowed, proceed without signature verification
    if token_alg == 'none':
        if signature_b64:
            raise InvalidTokenError("A token with 'alg=none' must not have a signature part.")
        # The token is considered valid *only because* 'none' was in the allowed list.
        return json.loads(base64url_decode(payload_b64))

    # (For demonstration) A real implementation would go on to verify signatures
    # for other algorithms like 'HS256', 'RS256', etc., using the provided key.
    # This part is omitted to focus on the 'alg=none' fix.
    
    # Placeholder for actual signature verification logic
    print(f"Proceeding to verify signature for algorithm: {token_alg} (Verification logic not implemented in this demo)")
    # ... verification would happen here ...
    
    return json.loads(base64url_decode(payload_b64))


# --- Demonstration of the Fix ---

if __name__ == '__main__':
    print("--- Demonstrating the fix for 'alg=none' vulnerability (CVE-2025-61152) ---")

    # A malicious actor crafts a token with an 'is_admin' claim.
    # Header: {"alg": "none", "typ": "JWT"}
    malicious_header = base64.urlsafe_b64encode(b'{"alg": "none", "typ": "JWT"}').rstrip(b'=')
    # Payload: {"user": "attacker", "is_admin": true}
    malicious_payload = base64.urlsafe_b64encode(b'{"user": "attacker", "is_admin": true}').rstrip(b'=')
    # Signature: (empty for alg=none)
    malicious_signature = b''

    malicious_none_token = b'.'.join([malicious_header, malicious_payload, malicious_signature]).decode('utf-8')
    
    print(f"\nMaliciously crafted token: {malicious_none_token}\n")

    # --- Scenario: A secure application that only expects 'HS256' tokens ---
    # The developer correctly specifies the algorithm(s) they use.
    # The list of allowed algorithms does NOT include 'none'.
    securely_allowed_algorithms = ['HS256']
    
    print(f"Attempting to validate the token with a secure configuration: allowed_algorithms={securely_allowed_algorithms}")
    
    try:
        # The fixed_jwt_decode function is called.
        decoded_payload = fixed_jwt_decode(
            token=malicious_none_token,
            key="any-secret-key",  # Key is irrelevant for this attack but required by the function signature
            algorithms=securely_allowed_algorithms
        )
        # This code block should NOT be reached.
        print("\n[FAIL] VULNERABILITY PRESENT: The forged 'alg=none' token was accepted.")
        print(f"Decoded Payload: {decoded_payload}")
    except InvalidAlgorithmError as e:
        # This is the expected, correct behavior. The fix works.
        print("\n[SUCCESS] VULNERABILITY FIXED: The forged 'alg=none' token was correctly rejected.")
        print(f"Reason: {e}")

    print("\n------------------------------------------------------------------")

    # --- Scenario: An application that EXPLICITLY needs to accept 'none' tokens ---
    # This is generally a bad practice, but the fixed library allows it if the
    # developer makes a conscious choice.
    insecurely_allowed_algorithms = ['HS256', 'none']
    
    print(f"\nAttempting to validate again with an explicit (but insecure) configuration: allowed_algorithms={insecurely_allowed_algorithms}")
    
    try:
        decoded_payload = fixed_jwt_decode(
            token=malicious_none_token,
            key="any-secret-key",
            algorithms=insecurely_allowed_algorithms
        )
        print("\n[SUCCESS] Token accepted as expected because 'none' was explicitly allowed by the developer.")
        print(f"Decoded Payload: {decoded_payload}")
    except JWTError as e:
        print(f"\n[FAIL] Token was unexpectedly rejected: {e}")