JWE Operations

def encrypt(plaintext, key, encryption='A256GCM', algorithm='dir', zip=None, cty=None, kid=None):
    """
    Encrypts plaintext and returns a JWE compact serialization string.
    
    Args:
        plaintext (bytes): The data to encrypt. Must be bytes object.
            Use encode() to convert strings to bytes
        key (str or bytes or dict): The encryption key. Supports:
            - Bytes for direct encryption (dir algorithm)
            - RSA public keys in PEM format for RSA key wrapping
            - JWK dictionaries with appropriate key material
        encryption (str): The content encryption algorithm. Defaults to 'A256GCM'.
            Supported: A128GCM, A192GCM, A256GCM, A128CBC-HS256, A192CBC-HS384, A256CBC-HS512
        algorithm (str): The key management algorithm. Defaults to 'dir'.
            Supported: dir, RSA1_5, RSA-OAEP, RSA-OAEP-256, A128KW, A192KW, A256KW
        zip (str, optional): The compression algorithm applied before encryption.
            Supported: 'DEF' for DEFLATE compression, None for no compression
        cty (str, optional): The content type of the secured content.
            Examples: 'application/json', 'text/plain'
        kid (str, optional): Key ID hint for the recipient to identify the key
    
    Returns:
        bytes: The JWE token in compact format (header.encryptedkey.iv.ciphertext.tag)
    
    Raises:
        JWEError: If encryption fails or parameters are invalid
        JWEAlgorithmUnsupportedError: If the specified algorithm is not supported
    """

from jose import jwe
from jose.constants import ALGORITHMS

# Basic AES-GCM encryption with direct key
key = b'This is a 32-byte key for AES256!!'  # 32 bytes for A256GCM
plaintext = b'Secret message'
encrypted = jwe.encrypt(plaintext, key, ALGORITHMS.A256GCM, ALGORITHMS.DIR)

# String to bytes conversion
message = "Confidential data"
encrypted = jwe.encrypt(message.encode('utf-8'), key, ALGORITHMS.A256GCM, ALGORITHMS.DIR)

# RSA key wrapping with OAEP
rsa_public_key = """-----BEGIN PUBLIC KEY-----
MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA...
-----END PUBLIC KEY-----"""

encrypted = jwe.encrypt(
    b'Sensitive data', 
    rsa_public_key, 
    encryption=ALGORITHMS.A256GCM,
    algorithm=ALGORITHMS.RSA_OAEP
)

# With compression
encrypted = jwe.encrypt(
    b'Large repetitive data that compresses well...',
    key,
    encryption=ALGORITHMS.A256GCM,
    algorithm=ALGORITHMS.DIR,
    zip='DEF'  # DEFLATE compression
)

# With content type and key ID
encrypted = jwe.encrypt(
    b'{"user": "john", "role": "admin"}',
    key,
    encryption=ALGORITHMS.A256GCM,
    algorithm=ALGORITHMS.DIR,
    cty='application/json',
    kid='encryption-key-1'
)

# AES key wrapping
aes_kek = b'Key Encryption Key 32 bytes!!!!!'  # 32 bytes for A256KW
encrypted = jwe.encrypt(
    b'Protected content',
    aes_kek,
    encryption=ALGORITHMS.A256GCM,
    algorithm=ALGORITHMS.A256KW
)

# CBC with HMAC authentication
encrypted = jwe.encrypt(
    b'Legacy compatibility data',
    key,
    encryption=ALGORITHMS.A256CBC_HS512,
    algorithm=ALGORITHMS.DIR
)

def decrypt(jwe_str, key):
    """
    Decrypts a JWE token and returns the original plaintext.
    
    Args:
        jwe_str (str or bytes): The JWE token to decrypt in compact format
        key (str or bytes or dict): The decryption key. Must match the key
            used for encryption and be appropriate for the key management algorithm:
            - Bytes for direct encryption (dir algorithm)  
            - RSA private keys in PEM format for RSA key wrapping
            - AES keys for AES key wrapping
            - JWK dictionaries with appropriate key material
    
    Returns:
        bytes: The decrypted plaintext as bytes. Use decode() to convert to string
    
    Raises:
        JWEError: If decryption fails or token format is invalid
        JWEParseError: If the JWE token cannot be parsed
        JWEInvalidAuth: If authentication verification fails
    """

from jose import jwe
from jose.exceptions import JWEError, JWEParseError, JWEInvalidAuth

try:
    # Basic decryption
    plaintext = jwe.decrypt(encrypted_token, key)
    message = plaintext.decode('utf-8')
    print(f"Decrypted: {message}")
    
except JWEInvalidAuth:
    print("Authentication verification failed - data may be tampered")
except JWEParseError:
    print("Invalid JWE token format")
except JWEError as e:
    print(f"Decryption failed: {e}")

# RSA key unwrapping
rsa_private_key = """-----BEGIN PRIVATE KEY-----
MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQC...
-----END PRIVATE KEY-----"""

plaintext = jwe.decrypt(encrypted_token, rsa_private_key)

# AES key unwrapping
aes_kek = b'Key Encryption Key 32 bytes!!!!!'
plaintext = jwe.decrypt(encrypted_token, aes_kek)

# Direct key decryption
direct_key = b'This is a 32-byte key for AES256!!'
plaintext = jwe.decrypt(encrypted_token, direct_key)

# JWK dictionary key
jwk_key = {
    'kty': 'oct',
    'k': 'VGhpcyBpcyBhIDMyLWJ5dGUga2V5IGZvciBBRVMyNTYhIQ'  # base64url encoded
}
plaintext = jwe.decrypt(encrypted_token, jwk_key)

def get_unverified_header(jwe_str):
    """
    Get JWE header without verification or decryption.
    
    Args:
        jwe_str (str or bytes): The JWE token
    
    Returns:
        dict: The JWE header dictionary
    
    Raises:
        JWEParseError: If the JWE token format is invalid
    """

# Inspect JWE header
header = jwe.get_unverified_header(encrypted_token)
print(f"Key Management Algorithm: {header['alg']}")
print(f"Content Encryption Algorithm: {header['enc']}")
print(f"Key ID: {header.get('kid')}")
print(f"Content Type: {header.get('cty')}")
print(f"Compression: {header.get('zip')}")

# Key selection based on header
header = jwe.get_unverified_header(encrypted_token)
key_id = header.get('kid')
if key_id == 'rsa-key-1':
    key = rsa_private_key_1
elif key_id == 'rsa-key-2':
    key = rsa_private_key_2
else:
    key = default_key

plaintext = jwe.decrypt(encrypted_token, key)

# Standard JWE header fields
header = {
    'alg': 'dir',             # Key management algorithm (required)
    'enc': 'A256GCM',         # Content encryption algorithm (required)
    'zip': 'DEF',             # Compression algorithm (optional)
    'kid': 'key-1',           # Key ID for key identification (optional)
    'cty': 'application/json', # Content type (optional)
    'typ': 'JWE'              # Type (optional)
}

from jose.constants import ALGORITHMS

# Modern authenticated encryption (recommended)
encrypted = jwe.encrypt(data, key, ALGORITHMS.A256GCM, ALGORITHMS.DIR)

# Legacy compatibility
encrypted = jwe.encrypt(data, key, ALGORITHMS.A256CBC_HS512, ALGORITHMS.DIR)

# Different key sizes
encrypted = jwe.encrypt(data, key_16, ALGORITHMS.A128GCM, ALGORITHMS.DIR)  # 16-byte key
encrypted = jwe.encrypt(data, key_24, ALGORITHMS.A192GCM, ALGORITHMS.DIR)  # 24-byte key
encrypted = jwe.encrypt(data, key_32, ALGORITHMS.A256GCM, ALGORITHMS.DIR)  # 32-byte key

# Direct encryption (CEK = key)
encrypted = jwe.encrypt(plaintext, direct_key, algorithm=ALGORITHMS.DIR)

# RSA key wrapping (recommended OAEP-256)
encrypted = jwe.encrypt(plaintext, rsa_public_key, algorithm=ALGORITHMS.RSA_OAEP_256)

# AES key wrapping
encrypted = jwe.encrypt(plaintext, aes_kek, algorithm=ALGORITHMS.A256KW)

# Direct encryption key requirements
key_128 = b'16-byte key here'    # For A128GCM
key_192 = b'24-byte key here!!!!!!!!'  # For A192GCM  
key_256 = b'This is a 32-byte key for AES256!!'  # For A256GCM

# RSA keys (PEM format)
rsa_public_key = """-----BEGIN PUBLIC KEY-----
... (RSA public key in PEM format for encryption)
-----END PUBLIC KEY-----"""

rsa_private_key = """-----BEGIN PRIVATE KEY-----
... (RSA private key in PEM format for decryption)  
-----END PRIVATE KEY-----"""

# AES Key Encryption Keys (KEK)
aes_kek_128 = b'16-byte KEK here'           # For A128KW
aes_kek_192 = b'24-byte KEK here!!!!!!!'   # For A192KW
aes_kek_256 = b'32-byte KEK for wrapping!!!!!!!!'  # For A256KW

# With compression (useful for large or repetitive data)
large_data = b'{"users": [' + b'{"name": "user", "role": "admin"},' * 1000 + b']}'
encrypted = jwe.encrypt(
    large_data,
    key,
    encryption=ALGORITHMS.A256GCM,
    algorithm=ALGORITHMS.DIR,
    zip='DEF'  # DEFLATE compression
)

# Compression is transparent during decryption
decrypted = jwe.decrypt(encrypted, key)  # Automatically decompressed

class JWEError(JOSEError):
    """Base exception for JWE-related errors."""

class JWEParseError(JWEError):
    """JWE token parsing failed."""

class JWEInvalidAuth(JWEError):
    """JWE authentication verification failed."""

class JWEAlgorithmUnsupportedError(JWEError):
    """JWE algorithm not supported."""

from jose.exceptions import JWEError, JWEParseError, JWEInvalidAuth, JWEAlgorithmUnsupportedError

try:
    plaintext = jwe.decrypt(encrypted_token, key)
except JWEInvalidAuth:
    print("Authentication failed - data may be corrupted or tampered")
except JWEParseError:
    print("Invalid JWE token format")
except JWEAlgorithmUnsupportedError:
    print("Encryption algorithm not supported by current backend")
except JWEError as e:
    print(f"JWE operation failed: {e}")

# Encryption errors
try:
    encrypted = jwe.encrypt(plaintext, key, 'UNSUPPORTED_ALG', 'dir')
except JWEAlgorithmUnsupportedError:
    print("Unsupported encryption algorithm")
except JWEError as e:
    print(f"Encryption failed: {e}")

from jose import jwe, jwk

# Use JWK for key management
key = jwk.construct({'kty': 'oct', 'k': 'encoded-key'})
encrypted = jwe.encrypt(plaintext, key, ALGORITHMS.A256GCM, ALGORITHMS.DIR)

# Nested JWT (signed then encrypted)
from jose import jwt
# First sign with JWT
signed_token = jwt.encode(claims, signing_key, algorithm='HS256')
# Then encrypt the JWT
encrypted_jwt = jwe.encrypt(signed_token.encode(), encryption_key, ALGORITHMS.A256GCM, ALGORITHMS.DIR)

# Decrypt then verify
decrypted_jwt = jwe.decrypt(encrypted_jwt, encryption_key)
claims = jwt.decode(decrypted_jwt.decode(), signing_key, algorithms=['HS256'])