Cryptographic Utilities

from cryptography.hazmat.primitives import constant_time, keywrap

def bytes_eq(a: bytes, b: bytes) -> bool:
    """
    Compare two byte strings in constant time.
    
    Args:
        a (bytes): First byte string to compare
        b (bytes): Second byte string to compare
        
    Returns:
        bool: True if the byte strings are equal, False otherwise
        
    Raises:
        TypeError: If either argument is not bytes
        
    Note:
        This function prevents timing attacks by ensuring the comparison
        takes the same amount of time regardless of where the first
        difference occurs or if the strings are identical.
    """

from cryptography.hazmat.primitives import constant_time

# Constant-time comparison for security-sensitive operations
expected_token = b"secret_authentication_token"
received_token = b"user_provided_token_value"

# Use constant_time.bytes_eq instead of == for security
if constant_time.bytes_eq(expected_token, received_token):
    print("Authentication successful")
else:
    print("Authentication failed")

# Don't use regular comparison for security-sensitive data:
# if expected_token == received_token:  # ❌ Vulnerable to timing attacks

def aes_key_wrap(wrapping_key: bytes, key_to_wrap: bytes, backend=None) -> bytes:
    """
    Wrap a key using AES Key Wrap (RFC 3394).
    
    Args:
        wrapping_key (bytes): AES key (16, 24, or 32 bytes)
        key_to_wrap (bytes): Key material to wrap (min 16 bytes, multiple of 8)
        backend: Cryptographic backend (usually None)
        
    Returns:
        bytes: Wrapped key data
        
    Raises:
        ValueError: If key lengths are invalid
    """

def aes_key_unwrap(wrapping_key: bytes, wrapped_key: bytes, backend=None) -> bytes:
    """
    Unwrap a key using AES Key Wrap (RFC 3394).
    
    Args:
        wrapping_key (bytes): AES key used for wrapping
        wrapped_key (bytes): Wrapped key data (min 24 bytes, multiple of 8)
        backend: Cryptographic backend (usually None)
        
    Returns:
        bytes: Original key material
        
    Raises:
        ValueError: If key lengths are invalid
        InvalidUnwrap: If unwrapping fails (integrity check failure)
    """

def aes_key_wrap_with_padding(wrapping_key: bytes, key_to_wrap: bytes, backend=None) -> bytes:
    """
    Wrap a key using AES Key Wrap with Padding (RFC 5649).
    
    Args:
        wrapping_key (bytes): AES key (16, 24, or 32 bytes)
        key_to_wrap (bytes): Key material to wrap (any length)
        backend: Cryptographic backend (usually None)
        
    Returns:
        bytes: Wrapped key data with padding
        
    Raises:
        ValueError: If wrapping key length is invalid
    """

def aes_key_unwrap_with_padding(wrapping_key: bytes, wrapped_key: bytes, backend=None) -> bytes:
    """
    Unwrap a key using AES Key Wrap with Padding (RFC 5649).
    
    Args:
        wrapping_key (bytes): AES key used for wrapping
        wrapped_key (bytes): Wrapped key data (min 16 bytes)
        backend: Cryptographic backend (usually None)
        
    Returns:
        bytes: Original key material without padding
        
    Raises:
        ValueError: If wrapping key length is invalid
        InvalidUnwrap: If unwrapping fails (integrity or padding check failure)
    """

class InvalidUnwrap(Exception):
    """Raised when key unwrapping fails due to integrity check failure."""

from cryptography.hazmat.primitives import keywrap
from cryptography.hazmat.primitives.ciphers import algorithms
import os

# Generate keys
wrapping_key = os.urandom(32)  # 256-bit AES key
key_to_wrap = os.urandom(32)   # Key material to protect

# AES Key Wrap (RFC 3394) - requires key to be multiple of 8 bytes
wrapped_key = keywrap.aes_key_wrap(wrapping_key, key_to_wrap)
unwrapped_key = keywrap.aes_key_unwrap(wrapping_key, wrapped_key)

assert unwrapped_key == key_to_wrap

# AES Key Wrap with Padding (RFC 5649) - works with any key length
arbitrary_key = os.urandom(17)  # Not a multiple of 8
wrapped_padded = keywrap.aes_key_wrap_with_padding(wrapping_key, arbitrary_key)
unwrapped_padded = keywrap.aes_key_unwrap_with_padding(wrapping_key, wrapped_padded)

assert unwrapped_padded == arbitrary_key