ML-KEM (Kyber): Post-Quantum Key Encapsulation

# 1) set the prime q
q = 3329

# 2) integer rounding function Round_q(x)
def Round_q_int(x):
    """
    Return 0 if x mod q lies in (-q/4, q/4), else return 1.
    Conventions as in slide: we interpret reps in [-(q-1)/2, (q-1)/2].
    """
    xm = Integer(x) % q
    # shift to symmetric rep in [-(q-1)/2, ..., (q-1)/2]
    if xm > (q-1)/2:
        xm -= q
    # now xm ∈ [-(q-1)/2, (q-1)/2]
    # check strict bounds −q/4 < xm < q/4
    if -q/4 < xm < q/4:
        return 0
    else:
        return 1

# test integer rounding
print("Integer tests:")
for test in [0, 832, -832, 833, 1664, 1665, 1666]:
    print(f"Round_{q}({test}) = {Round_q_int(test)}")

# 3) polynomial rounding: apply Round_q_int to each coefficient
R.<x> = PolynomialRing(Integers(q))
f = R(3000 + 1500*x + 2010*x^2 + 37*x^3)

def Round_q_poly(poly):
    coeffs = poly.list()    # [c0, c1, c2, ...]
    return sum( Round_q_int(c)*x^i for i,c in enumerate(coeffs) )

g = Round_q_poly(f)
print("\nPolynomial example:")
print("f =", f)
print("Round_q(f) =", g)

# SageMath setup for polynomial rings used in Kyber and Dilithium
import random
from sage.rings.polynomial.polynomial_ring import PolynomialRing_dense_finite_field

# Parameters for Kyber-512 (simplified)
n = 256  # Degree of polynomials
q = 3329  # Prime modulus (q ≡ 1 (mod 2n))

# Create the polynomial ring R_q = Z_q[X]/(X^n + 1)
Zq = IntegerModRing(q)
R.<x> = PolynomialRing(Zq)
Rq = R.quotient(x^n + 1)

print(f"Working in polynomial ring: Z_{q}[X]/(X^{n} + 1)")
print(f"Prime modulus q = {q}")
print(f"q ≡ 1 (mod 2n): {q % (2*n) == 1}")

# Example: create a random polynomial
def random_poly(degree_bound):
    """Generate a random polynomial with coefficients in [0, q-1]"""
    coeffs = [randint(0, q-1) for _ in range(degree_bound)]
    return Rq(R(coeffs))

# Create two random polynomials
f = random_poly(n)
g = random_poly(n)

print(f"\nExample polynomials (showing first 8 coefficients):")
print(f"f(x) = {list(f.lift().coefficients(sparse=False))[:8]}...")
print(f"g(x) = {list(g.lift().coefficients(sparse=False))[:8]}...")

# Polynomial arithmetic in R_q
h = f * g  # Multiplication modulo (X^n + 1) and q
print(f"f*g (first 8 coeffs) = {list(h.lift().coefficients(sparse=False))[:8]}...")

# Noise sampling for lattice-based cryptography

def centered_binomial_distribution(eta):
    """
    Sample from centered binomial distribution with parameter eta.
    Used in Kyber for noise generation.
    Returns a value in [-eta, eta] with binomial distribution.
    """
    a = sum(randint(0, 1) for _ in range(eta))
    b = sum(randint(0, 1) for _ in range(eta))
    return a - b

def sample_noise_poly(eta):
    """Generate a noise polynomial with coefficients from centered binomial distribution"""
    coeffs = [centered_binomial_distribution(eta) for _ in range(n)]
    return Rq(R(coeffs))

def uniform_poly():
    """Generate a uniformly random polynomial"""
    coeffs = [randint(0, q-1) for _ in range(n)]
    return Rq(R(coeffs))

# Example: Generate noise polynomials with different parameters
eta2 = 2  # Kyber-512 parameter
eta3 = 3  # Used in some variants

noise_poly_2 = sample_noise_poly(eta2)
noise_poly_3 = sample_noise_poly(eta3)

print("Noise polynomial coefficients (eta=2):")
noise_coeffs_2 = list(noise_poly_2.lift().coefficients(sparse=False))
print(f"First 16 coefficients: {noise_coeffs_2[:16]}")
print(f"Min coefficient: {min(noise_coeffs_2)}")
print(f"Max coefficient: {max(noise_coeffs_2)}")

print(f"\nNoise polynomial coefficients (eta=3):")
noise_coeffs_3 = list(noise_poly_3.lift().coefficients(sparse=False))
print(f"First 16 coefficients: {noise_coeffs_3[:16]}")
print(f"Min coefficient: {min(noise_coeffs_3)}")
print(f"Max coefficient: {max(noise_coeffs_3)}")

# Kyber Key Generation Algorithm

class KyberParams:
    """Kyber-512 parameters (simplified)"""
    def __init__(self):
        self.n = 256      # polynomial degree
        self.k = 2        # module rank
        self.q = 3329     # coefficient modulus
        self.eta1 = 3     # noise parameter for secret key
        self.eta2 = 2     # noise parameter for error terms

def kyber_keygen(params):
    """
    Kyber Key Generation Algorithm
    
    Returns:
        pk: Public key (A, t)
        sk: Secret key (s)
    """
    k = params.k
    
    # 1. Generate the public matrix A ∈ R_q^{k×k}
    # In practice, A is generated deterministically from a seed
    A = Matrix(Rq, k, k)
    for i in range(k):
        for j in range(k):
            A[i,j] = uniform_poly()
    
    # 2. Sample secret key vector s ∈ R_q^k from noise distribution
    s = vector([sample_noise_poly(params.eta1) for _ in range(k)])
    
    # 3. Sample error vector e ∈ R_q^k from noise distribution  
    e = vector([sample_noise_poly(params.eta2) for _ in range(k)])
    
    # 4. Compute t = As + e
    t = A * s + e
    
    # Public key is (A, t), secret key is s
    pk = (A, t)
    sk = s
    
    return pk, sk

# Generate a Kyber key pair
params = KyberParams()
public_key, secret_key = kyber_keygen(params)

A, t = public_key
print(f"Generated Kyber-{params.k*256} key pair")
print(f"Public key matrix A: {A.nrows()}×{A.ncols()}")
print(f"Public key vector t: length {len(t)}")
print(f"Secret key vector s: length {len(secret_key)}")

# Show first few coefficients of first polynomial in t
t0_coeffs = list(t[0].lift().coefficients(sparse=False))
print(f"First polynomial in t (first 8 coefficients): {t0_coeffs[:8]}")

# Verify the key generation equation: t = As + e
# We can't verify exactly since we don't store e, but we can check dimensions
print(f"\nKey generation verification:")
print(f"A dimensions: {A.dimensions()}")
print(f"s length: {len(secret_key)}")
print(f"t length: {len(t)}")
print(f"A * s computed successfully: {type(A * secret_key)}")

# Kyber Encapsulation and Decapsulation Algorithms

def kyber_encaps(pk, params):
    """
    Kyber Encapsulation Algorithm
    
    Args:
        pk: Public key (A, t)
        params: Kyber parameters
    
    Returns:
        ciphertext: (u, v)
        shared_secret: Random session key
    """
    A, t = pk
    k = params.k
    
    # 1. Sample message m (in practice, this is a random value)
    # For simplicity, we'll use a random polynomial
    m = sample_noise_poly(1)  # Small message
    
    # 2. Sample random vector r ∈ R_q^k from noise distribution
    r = vector([sample_noise_poly(params.eta1) for _ in range(k)])
    
    # 3. Sample error terms e1 ∈ R_q^k, e2 ∈ R_q
    e1 = vector([sample_noise_poly(params.eta2) for _ in range(k)])
    e2 = sample_noise_poly(params.eta2)
    
    # 4. Compute u = A^T * r + e1
    u = A.transpose() * r + e1
    
    # 5. Compute v = t^T * r + e2 + encode(m)
    # For simplicity, we'll add m directly (in practice, encoding is more complex)
    v = sum(t[i] * r[i] for i in range(k)) + e2 + m
    
    # The shared secret is derived from m (simplified)
    shared_secret = m
    
    ciphertext = (u, v)
    return ciphertext, shared_secret

def kyber_decaps(ciphertext, sk, params):
    """
    Kyber Decapsulation Algorithm
    
    Args:
        ciphertext: (u, v)
        sk: Secret key s
        params: Kyber parameters
    
    Returns:
        shared_secret: Recovered session key
    """
    u, v = ciphertext
    s = sk
    k = params.k
    
    # 1. Compute m' = v - s^T * u
    m_prime = v - sum(s[i] * u[i] for i in range(k))
    
    # 2. The shared secret is derived from m' (simplified)
    shared_secret = m_prime
    
    return shared_secret

# Test the encapsulation/decapsulation
print("Testing Kyber Encapsulation/Decapsulation:")
print("=" * 50)

# Encapsulate
ciphertext, original_secret = kyber_encaps(public_key, params)
u, v = ciphertext

print(f"Encapsulation completed")
print(f"Ciphertext u: vector of length {len(u)}")
print(f"Ciphertext v: polynomial")
print(f"Original shared secret (first 8 coeffs): {list(original_secret.lift().coefficients(sparse=False))[:8]}")

# Decapsulate
recovered_secret = kyber_decaps(ciphertext, secret_key, params)

print(f"\nDecapsulation completed")
print(f"Recovered shared secret (first 8 coeffs): {list(recovered_secret.lift().coefficients(sparse=False))[:8]}")

# Check if decapsulation worked
difference = original_secret - recovered_secret
diff_coeffs = list(difference.lift().coefficients(sparse=False))
max_error = max(abs(c) for c in diff_coeffs) if diff_coeffs else 0

print(f"\nVerification:")
print(f"Maximum coefficient difference: {max_error}")
print(f"Decapsulation successful: {max_error < 10}")  # Small threshold for noise

# ML-KEM Performance Benchmarking and Analysis

import time

def benchmark_mlkem(params, num_trials=10):
    """Benchmark ML-KEM operations"""
    
    print(f"⚡ Benchmarking ML-KEM with {num_trials} trials...")
    print(f"   📊 Note: Educational implementation, not optimized for performance")
    
    # Key Generation
    keygen_times = []
    for _ in range(num_trials):
        start = time.time()
        pk, sk = kyber_keygen(params)
        end = time.time()
        keygen_times.append(end - start)
    
    # Encapsulation
    encaps_times = []
    pk, sk = kyber_keygen(params)  # Generate once for encaps/decaps
    for _ in range(num_trials):
        start = time.time()
        ct, secret = kyber_encaps(pk, params)
        end = time.time()
        encaps_times.append(end - start)
    
    # Decapsulation
    decaps_times = []
    ct, _ = kyber_encaps(pk, params)  # Generate once for decaps
    for _ in range(num_trials):
        start = time.time()
        recovered = kyber_decaps(ct, sk, params)
        end = time.time()
        decaps_times.append(end - start)
    
    return {
        'keygen': keygen_times,
        'encaps': encaps_times, 
        'decaps': decaps_times
    }

def print_benchmark_results(name, results):
    """Print benchmark results in a nice format"""
    print(f"\n🏆 {name} Benchmark Results:")
    print("=" * 40)
    
    for operation, times in results.items():
        if times:
            avg_time = sum(times) / len(times) * 1000  # Convert to milliseconds
            min_time = min(times) * 1000
            max_time = max(times) * 1000
            print(f"{operation.capitalize():12} - Avg: {avg_time:.2f}ms, Min: {min_time:.2f}ms, Max: {max_time:.2f}ms")
        else:
            print(f"{operation.capitalize():12} - No successful operations")

def analyze_mlkem_security(params):
    """Analyze ML-KEM security properties"""
    print(f"\n🔒 ML-KEM Security Analysis:")
    print("=" * 40)
    
    # Generate sample keys to analyze
    pk, sk = kyber_keygen(params)
    A, t = pk
    s = sk
    
    # Analyze secret key properties
    s_norms = [infinity_norm(poly) for poly in s]
    max_s_norm = max(s_norms)
    avg_s_norm = sum(s_norms) / len(s_norms)
    
    print(f"Secret key analysis:")
    print(f"   • Max ||s_i||∞: {max_s_norm}")
    print(f"   • Avg ||s_i||∞: {avg_s_norm:.1f}")
    print(f"   • Expected range: [-{params.eta1}, {params.eta1}]")
    
    # Analyze public key properties  
    t_norms = [infinity_norm(poly) for poly in t]
    max_t_norm = max(t_norms)
    avg_t_norm = sum(t_norms) / len(t_norms)
    
    print(f"\nPublic key analysis:")
    print(f"   • Max ||t_i||∞: {max_t_norm}")
    print(f"   • Avg ||t_i||∞: {avg_t_norm:.1f}")
    print(f"   • Theoretical max: ~{params.q//2}")
    
    # Test encapsulation/decapsulation correctness
    print(f"\nCorrectness verification:")
    num_tests = 10
    successes = 0
    
    for i in range(num_tests):
        ct, original = kyber_encaps(pk, params)
        recovered = kyber_decaps(ct, sk, params)
        if (original - recovered) == 0:
            successes += 1
    
    success_rate = successes / num_tests
    print(f"   • Correctness rate: {success_rate:.1%} ({successes}/{num_tests})")
    
    return {
        'secret_key_max_norm': max_s_norm,
        'public_key_max_norm': max_t_norm,
        'correctness_rate': success_rate
    }

# Run ML-KEM benchmarks and analysis
print("🚀 ML-KEM Performance and Security Analysis")
print("=" * 60)

# Performance benchmark
mlkem_results = benchmark_mlkem(params, num_trials=5)
print_benchmark_results("ML-KEM", mlkem_results)

# Security analysis
security_results = analyze_mlkem_security(params)

print(f"\n📊 ML-KEM vs Classical Comparison:")
print("=" * 40)
print(f"{'Algorithm':<12} {'Key Size':<10} {'CT Size':<10} {'Security':<15}")
print("-" * 50)
print(f"{'RSA-2048':<12} {'~2KB':<10} {'256B':<10} {'Classical':<15}")
print(f"{'ECDH P-256':<12} {'32B':<10} {'64B':<10} {'Classical':<15}")
print(f"{'ML-KEM-512':<12} {'~800B':<10} {'768B':<10} {'Post-Quantum':<15}")
print(f"{'ML-KEM-768':<12} {'~1.2KB':<10} {'1.1KB':<10} {'Post-Quantum':<15}")

print(f"\n✅ ML-KEM analysis complete!")
print(f"   For digital signatures, see the ML-DSA tutorial: [MLDSA.ipynb](MLDSA)")
print("=" * 60)

# Interactive ML-KEM Exploration and Exercises

print("🎓 ML-KEM (Kyber) Tutorial Complete!")
print("=" * 50)

print("\n📚 What you've learned:")
print("• Post-quantum cryptography fundamentals")
print("• Module Learning With Errors (M-LWE) problem")
print("• ML-KEM key encapsulation mechanism")
print("• Key generation, encapsulation, and decapsulation")
print("• Security analysis and performance characteristics")

print("\n🔬 Try these ML-KEM exercises:")
print("1. Modify noise parameters (eta1, eta2) and observe correctness")
print("2. Test different module ranks k and analyze security/performance trade-offs")
print("3. Implement a simple hybrid encryption scheme (ML-KEM + AES)")
print("4. Compare encapsulation/decapsulation times with different parameters")
print("5. Analyze the effect of different modulus sizes on operations")

print("\n🔧 Experimental parameters to try:")
print("• Module ranks: k=2 (Kyber-512), k=3 (Kyber-768), k=4 (Kyber-1024)")
print("• Noise bounds: eta1 ∈ {2, 3}, eta2 ∈ {2, 3}")
print("• Polynomial degree: n ∈ {128, 256, 512}")
print("• Different moduli q (must satisfy q ≡ 1 (mod 2n))")

# Example: Quick parameter experiment
print("\n🧪 Quick experiment - Effect of module rank k:")
for k_test in [2, 3, 4]:
    params_test = KyberParams()
    params_test.k = k_test
    
    # Estimate key sizes (simplified)
    pk_size = k_test * params_test.n * 12 // 8  # Approximate bits to bytes
    ct_size = (k_test + 1) * params_test.n * 12 // 8
    
    print(f"k={k_test}: ~{pk_size}B public key, ~{ct_size}B ciphertext")

print("\n🔗 Next steps:")
print("• Explore ML-DSA (Dilithium) signatures in [MLDSA.ipynb](MLDSA)")
print("• Study lattice reduction algorithms in [LLL.ipynb](LLL)")
print("• Learn about LWE problem fundamentals in [LWE.ipynb](LWE)")

# Final verification that ML-KEM works
print("\n🔍 Final ML-KEM verification:")
try:
    # Test complete ML-KEM workflow
    pk_final, sk_final = kyber_keygen(params)
    ct_final, secret_orig = kyber_encaps(pk_final, params)
    secret_recv = kyber_decaps(ct_final, sk_final, params)
    
    # Check correctness
    difference = secret_orig - secret_recv
    is_correct = difference == 0
    
    print(f"✓ ML-KEM end-to-end test: {'✅ PASSED' if is_correct else '❌ FAILED'}")
    
    if is_correct:
        print(f"  • Key encapsulation and decapsulation successful")
        print(f"  • Shared secret correctly recovered")
    else:
        print(f"  • Error in shared secret recovery")
        
except Exception as e:
    print(f"❌ Error in final verification: {e}")

print(f"\n🌟 Congratulations! You've mastered ML-KEM fundamentals!")
print(f"📖 Continue your post-quantum journey with ML-DSA signatures!")
print(f"✨ Remember: Use production libraries for real applications!")

print("\n" + "=" * 50)
print("🚀 ML-KEM Tutorial completed successfully! 🎉")
print("=" * 50)