JSON Web Keys (JWK)

class PyJWK:
    def __init__(self, jwk_data: dict, algorithm: str = None):
        """
        Initialize PyJWK from JWK dictionary.

        Args:
            jwk_data (dict): JWK data according to RFC 7517
            algorithm (str): Force specific algorithm (optional)

        Raises:
            InvalidKeyError: Invalid key format or unsupported key type
            PyJWKError: General JWK processing error
            MissingCryptographyError: Cryptography library required
        """

    @staticmethod
    def from_dict(obj: dict, algorithm: str = None) -> 'PyJWK':
        """
        Create PyJWK from dictionary.

        Args:
            obj (dict): JWK dictionary
            algorithm (str): Optional algorithm override

        Returns:
            PyJWK: Initialized key object
        """

    @staticmethod  
    def from_json(data: str, algorithm: str = None) -> 'PyJWK':
        """
        Create PyJWK from JSON string.

        Args:
            data (str): JWK JSON string
            algorithm (str): Optional algorithm override

        Returns:
            PyJWK: Initialized key object
        """

    @property
    def key_type(self) -> str:
        """Key type (kty parameter): 'RSA', 'EC', 'oct', 'OKP'."""

    @property
    def key_id(self) -> str:
        """Key ID (kid parameter) for key identification."""

    @property
    def public_key_use(self) -> str:
        """Public key use (use parameter): 'sig', 'enc', or None."""

    @property
    def algorithm_name(self) -> str:
        """Detected or assigned algorithm name."""

    @property
    def Algorithm(self):
        """Algorithm implementation object."""

    @property
    def key(self):
        """Underlying cryptographic key object."""

import jwt
from jwt import PyJWK

# RSA public key from JWK
rsa_jwk_data = {
    "kty": "RSA",
    "kid": "rsa-key-1",
    "use": "sig",
    "n": "base64url-encoded-modulus",
    "e": "AQAB"
}

jwk = PyJWK.from_dict(rsa_jwk_data)
print(f"Key type: {jwk.key_type}")
print(f"Key ID: {jwk.key_id}")
print(f"Algorithm: {jwk.algorithm_name}")

# Use with JWT decoding
token = "eyJ..."
payload = jwt.decode(token, jwk.key, algorithms=[jwk.algorithm_name])

# ECDSA key with explicit algorithm
ec_jwk_data = {
    "kty": "EC",
    "kid": "ec-key-1",
    "crv": "P-256",
    "x": "base64url-encoded-x",
    "y": "base64url-encoded-y"
}

jwk = PyJWK.from_dict(ec_jwk_data, algorithm="ES256")

# Symmetric key (HMAC)
oct_jwk_data = {
    "kty": "oct",
    "kid": "hmac-key-1",
    "k": "base64url-encoded-secret"
}

jwk = PyJWK.from_dict(oct_jwk_data)  # Defaults to HS256

class PyJWKSet:
    def __init__(self, keys: list):
        """
        Initialize JWK Set from list of JWK dictionaries.

        Args:
            keys (list): List of JWK dictionaries

        Raises:
            PyJWKSetError: Invalid JWK Set or no usable keys
            MissingCryptographyError: Required for asymmetric keys
        """

    @staticmethod
    def from_dict(obj: dict) -> 'PyJWKSet':
        """
        Create PyJWKSet from JWKS dictionary.

        Args:
            obj (dict): JWKS dictionary with 'keys' array

        Returns:
            PyJWKSet: Initialized key set
        """

    @staticmethod
    def from_json(data: str) -> 'PyJWKSet':
        """
        Create PyJWKSet from JWKS JSON string.

        Args:
            data (str): JWKS JSON string

        Returns:
            PyJWKSet: Initialized key set
        """

    def __getitem__(self, kid: str) -> PyJWK:
        """
        Get key by key ID.

        Args:
            kid (str): Key ID to look up

        Returns:
            PyJWK: Matching key

        Raises:
            KeyError: No key found with given ID
        """

    @property
    def keys(self) -> list:
        """List of PyJWK objects in this set."""

import jwt
from jwt import PyJWKSet

# JWKS from dictionary
jwks_data = {
    "keys": [
        {
            "kty": "RSA",
            "kid": "rsa-key-1",
            "use": "sig",
            "n": "modulus...",
            "e": "AQAB"
        },
        {
            "kty": "EC", 
            "kid": "ec-key-1",
            "crv": "P-256",
            "x": "x-coordinate...",
            "y": "y-coordinate..."
        }
    ]
}

jwk_set = PyJWKSet.from_dict(jwks_data)

# Get specific key by ID
key = jwk_set['rsa-key-1']
print(f"Found key: {key.algorithm_name}")

# Iterate through all keys
for jwk in jwk_set.keys:
    print(f"Key {jwk.key_id}: {jwk.algorithm_name}")

# From JSON string
jwks_json = '''{"keys": [...]}'''
jwk_set = PyJWKSet.from_json(jwks_json)

# Use with JWT verification
def verify_jwt_with_jwks(token, jwk_set):
    # Extract key ID from token header
    unverified_header = jwt.get_unverified_header(token)
    kid = unverified_header.get('kid')
    
    if not kid:
        raise ValueError("Token missing key ID")
    
    # Get the appropriate key
    key = jwk_set[kid]
    
    # Verify token
    return jwt.decode(token, key.key, algorithms=[key.algorithm_name])

# RSA public key
rsa_public_jwk = {
    "kty": "RSA",
    "kid": "rsa-pub-1",
    "use": "sig",
    "n": "modulus-base64url",     # Required
    "e": "exponent-base64url"    # Required
}

# RSA private key (includes all public parameters plus private)
rsa_private_jwk = {
    "kty": "RSA",
    "kid": "rsa-priv-1", 
    "use": "sig",
    "n": "modulus-base64url",
    "e": "exponent-base64url",
    "d": "private-exponent-base64url",  # Private key
    # Optional CRT parameters: p, q, dp, dq, qi
}

# ECDSA P-256 public key
ec_p256_jwk = {
    "kty": "EC",
    "kid": "ec-p256-1",
    "crv": "P-256",              # Curve: P-256, P-384, P-521, secp256k1
    "x": "x-coordinate-base64url",
    "y": "y-coordinate-base64url"
}

# ECDSA private key (includes 'd' parameter)
ec_private_jwk = {
    "kty": "EC",
    "kid": "ec-priv-1",
    "crv": "P-256",
    "x": "x-coordinate-base64url",
    "y": "y-coordinate-base64url", 
    "d": "private-key-base64url"  # Private key
}

# HMAC symmetric key
oct_jwk = {
    "kty": "oct",
    "kid": "hmac-1",
    "k": "symmetric-key-base64url"  # The symmetric key material
}

# Ed25519 public key
ed25519_jwk = {
    "kty": "OKP",
    "kid": "ed25519-1",
    "crv": "Ed25519",            # Or "Ed448"
    "x": "public-key-base64url"
}

# Ed25519 private key
ed25519_private_jwk = {
    "kty": "OKP", 
    "kid": "ed25519-priv-1",
    "crv": "Ed25519",
    "x": "public-key-base64url",
    "d": "private-key-base64url"  # Private key
}

# Algorithm detection examples
rsa_jwk = PyJWK.from_dict({"kty": "RSA", "n": "...", "e": "AQAB"})
print(rsa_jwk.algorithm_name)  # "RS256"

ec_jwk = PyJWK.from_dict({"kty": "EC", "crv": "P-384", "x": "...", "y": "..."})
print(ec_jwk.algorithm_name)  # "ES384"

# Override automatic detection
forced_jwk = PyJWK.from_dict(rsa_jwk_data, algorithm="PS256")
print(forced_jwk.algorithm_name)  # "PS256"

import jwt
from jwt.exceptions import PyJWKError, InvalidKeyError, MissingCryptographyError

try:
    jwk = PyJWK.from_dict(invalid_jwk_data)
except InvalidKeyError as e:
    print(f"Invalid key format: {e}")
except MissingCryptographyError as e:
    print(f"Install cryptography: {e}")
except PyJWKError as e:
    print(f"JWK error: {e}")

try:
    key = jwk_set['nonexistent-kid']
except KeyError as e:
    print(f"Key not found: {e}")