ML-DSA (Dilithium): Post-Quantum Digital Signatures

# SageMath Setup for ML-DSA
import random
from sage.rings.polynomial.polynomial_ring import PolynomialRing_dense_finite_field

# ML-DSA Parameters (Dilithium-2 simplified)
class MLDSAParams:
    """ML-DSA-44 (formerly Dilithium-2) parameters"""
    def __init__(self):
        self.n = 256          # polynomial degree
        self.k = 4            # rows in A (verification key dimension)
        self.l = 4            # columns in A (signing key dimension)
        self.q = 8380417      # coefficient modulus
        self.eta = 2          # secret key noise bound
        self.gamma1 = 2**17   # mask range parameter
        self.gamma2 = (self.q - 1) // 88  # rejection threshold
        self.beta = self.eta * 256  # signature bound
        self.tau = 39         # challenge weight (number of ±1 coefficients)

# Initialize polynomial ring
params = MLDSAParams()
n, q = params.n, params.q

print(f"🔧 Initializing ML-DSA with parameters:")
print(f"   • Polynomial degree n = {n}")
print(f"   • Modulus q = {q}")
print(f"   • Matrix dimensions: {params.k}×{params.l}")
print(f"   • q ≡ 1 (mod 2n): {q % (2*n) == 1}")

# 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"✓ Polynomial ring R_q created successfully")
print(f"✓ Working in: Z_{q}[X]/(X^{n} + 1)")

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

def small_poly(bound):
    """Generate polynomial with coefficients in [-bound, bound]"""
    coeffs = [randint(-bound, bound) for _ in range(n)]
    return Rq(R(coeffs))

def infinity_norm(poly):
    """Compute infinity norm (max absolute coefficient) of polynomial"""
    coeffs = poly.lift().coefficients(sparse=False)
    if not coeffs:
        return 0
    return max(abs(int(c)) for c in coeffs)

def vector_infinity_norm(vec):
    """Compute infinity norm of vector of polynomials"""
    return max(infinity_norm(poly) for poly in vec)

# Test the setup
test_poly = uniform_poly()
test_small = small_poly(params.eta)

print(f"\n🧪 Testing polynomial operations:")
print(f"   • Random polynomial infinity norm: {infinity_norm(test_poly)}")
print(f"   • Small polynomial (eta={params.eta}) infinity norm: {infinity_norm(test_small)}")
print(f"   • Polynomial multiplication works: {type(test_poly * test_small)}")

print(f"\n✅ ML-DSA environment setup complete!")

# ML-DSA Key Generation Algorithm

def mldsa_keygen(params):
    """
    ML-DSA Key Generation
    
    Args:
        params: MLDSAParams object containing algorithm parameters
        
    Returns:
        tuple: (verification_key, signing_key)
            - verification_key: (A, t) - public matrix and vector
            - signing_key: (A, s1, s2, t) - includes secret vectors
    """
    k, l, eta = params.k, params.l, params.eta
    
    print(f"🔑 Generating ML-DSA key pair...")
    
    # Step 1: Generate public matrix A ∈ R_q^{k×l}
    # In practice, A is generated deterministically from a seed
    A = Matrix(Rq, k, l)
    for i in range(k):
        for j in range(l):
            A[i,j] = uniform_poly()
    
    print(f"   ✓ Generated public matrix A ({k}×{l})")
    
    # Step 2: Sample secret key vectors with small coefficients
    # s1 ∈ R_q^l with coefficients in [-η, η]
    s1 = vector([small_poly(eta) for _ in range(l)])
    
    # s2 ∈ R_q^k with coefficients in [-η, η]  
    s2 = vector([small_poly(eta) for _ in range(k)])
    
    print(f"   ✓ Generated secret vectors s1, s2 with ||·||∞ ≤ {eta}")
    
    # Step 3: Compute public vector t = As1 + s2
    t = A * s1 + s2
    
    print(f"   ✓ Computed public vector t = As1 + s2")
    
    # Verification key: (A, t)
    verification_key = (A, t)
    
    # Signing key: (A, s1, s2, t) - includes everything needed for signing
    signing_key = (A, s1, s2, t)
    
    # Verify key generation correctness
    verification_check = A * s1 + s2 - t
    all_zero = all(poly == 0 for poly in verification_check)
    
    print(f"   ✓ Key generation verification: {'PASS' if all_zero else 'FAIL'}")
    
    # Display key properties
    s1_norm = vector_infinity_norm(s1)
    s2_norm = vector_infinity_norm(s2)
    t_norm = vector_infinity_norm(t)
    
    print(f"   📊 Key statistics:")
    print(f"      • ||s1||∞ = {s1_norm}")
    print(f"      • ||s2||∞ = {s2_norm}")  
    print(f"      • ||t||∞ = {t_norm}")
    
    return verification_key, signing_key

# Generate ML-DSA key pair
print("=" * 60)
verification_key, signing_key = mldsa_keygen(params)
A, t = verification_key
A_sk, s1, s2, t_sk = signing_key

print(f"\n🎉 ML-DSA key pair generated successfully!")
print(f"   • Verification key size: {params.k}×{params.l} matrix + {params.k} vector")
print(f"   • Signing key size: includes {params.l + params.k} secret polynomials")
print("=" * 60)

# ML-DSA Challenge Generation and Signing

def generate_challenge(message, commitment_vector, params):
    """
    Generate challenge polynomial from message and commitment
    
    Args:
        message: Message string to be signed
        commitment_vector: Vector w from commitment phase
        params: ML-DSA parameters
        
    Returns:
        Challenge polynomial c with exactly τ non-zero ±1 coefficients
    """
    tau = params.tau  # Number of ±1 coefficients
    
    # Simplified challenge generation (in practice, uses SHAKE-256)
    # Create a deterministic "hash" from message and commitment
    message_hash = hash(str(message) + str(commitment_vector[0]))
    
    # Use hash to seed random generator for reproducible challenges
    random.seed(message_hash % (2**32))
    
    # Create sparse polynomial with exactly τ non-zero ±1 coefficients
    coeffs = [0] * params.n
    positions = random.sample(range(params.n), tau)
    
    for pos in positions:
        coeffs[pos] = random.choice([-1, 1])
    
    return Rq(R(coeffs))

def mldsa_sign(message, signing_key, params, max_attempts=1000):
    """
    ML-DSA Signing Algorithm with Rejection Sampling
    
    Args:
        message: Message to sign (string)
        signing_key: (A, s1, s2, t) from key generation
        params: ML-DSA parameters
        max_attempts: Maximum rejection sampling attempts
        
    Returns:
        Signature (c, z) or None if signing failed
    """
    A, s1, s2, t = signing_key
    k, l = params.k, params.l
    gamma1, beta = params.gamma1, params.beta
    
    print(f"📝 Signing message: '{message}'")
    print(f"   🎯 Target: ||z||∞ < {gamma1 - beta} (γ₁ - β)")
    
    attempts = 0
    
    for attempt in range(max_attempts):
        attempts += 1
        
        # Step 1: Sample masking vector y ∈ R_q^l uniformly from [-γ₁, γ₁]
        # For computational efficiency, we use a smaller range
        mask_range = min(gamma1 // 1000, 100)  # Scaled down for demo
        y = vector([small_poly(mask_range) for _ in range(l)])
        
        # Step 2: Compute commitment w = Ay
        w = A * y
        
        if attempt == 0:
            print(f"   🎲 First attempt - mask range: [-{mask_range}, {mask_range}]")
            print(f"   📐 Commitment ||w||∞ = {vector_infinity_norm(w)}")
        
        # Step 3: Generate challenge c from message and commitment
        c = generate_challenge(message, w, params)
        
        if attempt == 0:
            challenge_weight = sum(1 for coeff in c.lift().coefficients(sparse=False) if coeff != 0)
            print(f"   🎯 Challenge weight: {challenge_weight} (target: {params.tau})")
        
        # Step 4: Compute response z = y + cs1
        z = y + vector([c * s1_poly for s1_poly in s1])
        
        # Step 5: Rejection sampling check
        z_norm = vector_infinity_norm(z)
        threshold = (gamma1 - beta) // 1000  # Scaled threshold
        
        if z_norm < threshold:
            # Additional check: w - cs2 should also be small
            w_cs2 = w - vector([c * s2_poly for s2_poly in s2])
            w_cs2_norm = vector_infinity_norm(w_cs2)
            
            if w_cs2_norm < params.gamma2 // 1000:  # Scaled threshold
                print(f"   ✅ Signature accepted after {attempts} attempts")
                print(f"      • ||z||∞ = {z_norm} < {threshold}")
                print(f"      • ||w - cs2||∞ = {w_cs2_norm}")
                return (c, z)
        
        if attempt > 0 and attempt % 100 == 0:
            print(f"   ⏳ Attempt {attempt}: ||z||∞ = {z_norm} ≥ {threshold} (rejected)")
    
    print(f"   ❌ Signing failed after {max_attempts} attempts")
    return None

# Test the signing algorithm
print("=" * 60)
message = "Hello, Post-Quantum World! 🌍"
signature = mldsa_sign(message, signing_key, params, max_attempts=200)

if signature:
    c, z = signature
    challenge_norm = infinity_norm(c)
    response_norm = vector_infinity_norm(z)
    
    print(f"\n🎉 Signature generated successfully!")
    print(f"   • Challenge ||c||∞ = {challenge_norm}")
    print(f"   • Response ||z||∞ = {response_norm}")
    print(f"   • Signature components: (c, z)")
else:
    print(f"\n❌ Signature generation failed")

print("=" * 60)

# ML-DSA Signature Verification

def mldsa_verify(message, signature, verification_key, params):
    """
    ML-DSA Signature Verification
    
    Args:
        message: Original message (string)
        signature: (c, z) signature pair
        verification_key: (A, t) public key
        params: ML-DSA parameters
        
    Returns:
        Boolean indicating signature validity
    """
    if signature is None:
        return False
        
    c, z = signature
    A, t = verification_key
    gamma1, beta = params.gamma1, params.beta
    
    print(f"🔍 Verifying signature for message: '{message}'")
    
    # Step 1: Check response bound ||z||∞ < γ₁ - β
    z_norm = vector_infinity_norm(z)
    threshold = (gamma1 - beta) // 1000  # Scaled threshold
    
    print(f"   📏 Checking response bound: ||z||∞ = {z_norm}")
    
    if z_norm >= threshold:
        print(f"   ❌ Response too large: {z_norm} ≥ {threshold}")
        return False
    
    print(f"   ✓ Response bound check passed: {z_norm} < {threshold}")
    
    # Step 2: Compute w' = Az - ct
    w_prime = A * z - vector([c * t_poly for t_poly in t])
    
    print(f"   🔄 Computed verification commitment w' = Az - ct")
    print(f"      • ||w'||∞ = {vector_infinity_norm(w_prime)}")
    
    # Step 3: Regenerate challenge c' from message and w'
    c_prime = generate_challenge(message, w_prime, params)
    
    # Step 4: Check if c = c' (polynomial equality)
    challenge_match = (c == c_prime)
    
    print(f"   🎯 Challenge comparison:")
    print(f"      • Original challenge ||c||∞ = {infinity_norm(c)}")
    print(f"      • Recomputed challenge ||c'||∞ = {infinity_norm(c_prime)}")
    print(f"      • Challenges match: {challenge_match}")
    
    if challenge_match:
        print(f"   ✅ Signature verification: VALID")
    else:
        print(f"   ❌ Signature verification: INVALID")
    
    return challenge_match

def test_signature_validity(message, signature, verification_key, params):
    """Test signature with correct and incorrect messages"""
    
    print(f"\n🧪 Testing signature validity...")
    
    # Test 1: Correct message
    print(f"\n📝 Test 1: Original message")
    is_valid = mldsa_verify(message, signature, verification_key, params)
    
    # Test 2: Modified message
    modified_message = message + " [MODIFIED]"
    print(f"\n📝 Test 2: Modified message")
    print(f"   Original: '{message}'")
    print(f"   Modified: '{modified_message}'")
    is_valid_modified = mldsa_verify(modified_message, signature, verification_key, params)
    
    # Test 3: Corrupted signature (if signature exists)
    if signature:
        c, z = signature
        # Modify one coefficient in the response
        z_corrupted = vector([z[0] + Rq(R([1]))] + list(z[1:]))
        corrupted_signature = (c, z_corrupted)
        
        print(f"\n📝 Test 3: Corrupted signature")
        is_valid_corrupted = mldsa_verify(message, corrupted_signature, verification_key, params)
    
    return is_valid, is_valid_modified, is_valid_corrupted if signature else False

# Verify the signature we generated
print("=" * 60)

if signature:
    # Verify the signature
    is_valid = mldsa_verify(message, signature, verification_key, params)
    
    # Run comprehensive tests
    test_results = test_signature_validity(message, signature, verification_key, params)
    valid_original, valid_modified, valid_corrupted = test_results
    
    print(f"\n📊 Verification Summary:")
    print(f"   • Original message: {'✅ VALID' if valid_original else '❌ INVALID'}")
    print(f"   • Modified message: {'✅ VALID' if valid_modified else '❌ INVALID'}")
    print(f"   • Corrupted signature: {'✅ VALID' if valid_corrupted else '❌ INVALID'}")
    
    # Expected results
    expected_results = valid_original and not valid_modified and not valid_corrupted
    print(f"\n🎯 Security test: {'✅ PASSED' if expected_results else '❌ FAILED'}")
    print(f"   (Valid signature should only verify for original message)")
    
else:
    print(f"❌ Cannot verify - no signature was generated")

print("=" * 60)

# Performance Analysis and Benchmarking

import time
from statistics import mean, stdev

def benchmark_mldsa(params, num_trials=5):
    """
    Benchmark ML-DSA operations
    
    Args:
        params: ML-DSA parameters
        num_trials: Number of trials for each operation
        
    Returns:
        Dictionary with timing results
    """
    print(f"⚡ Benchmarking ML-DSA operations ({num_trials} trials each)")
    print(f"   📊 Note: This is an educational implementation, not optimized")
    
    results = {
        'keygen_times': [],
        'sign_times': [],
        'verify_times': [],
        'sign_attempts': [],
        'sign_success_rate': 0
    }
    
    # Benchmark Key Generation
    print(f"\n🔑 Benchmarking Key Generation...")
    for trial in range(num_trials):
        start_time = time.time()
        vk, sk = mldsa_keygen(params)
        end_time = time.time()
        results['keygen_times'].append(end_time - start_time)
    
    keygen_avg = mean(results['keygen_times']) * 1000  # Convert to ms
    print(f"   ✓ Average key generation time: {keygen_avg:.2f} ms")
    
    # Benchmark Signing (with attempt counting)
    print(f"\n📝 Benchmarking Signing...")
    test_messages = [f"Test message {i}" for i in range(num_trials)]
    successful_signs = 0
    
    for i, message in enumerate(test_messages):
        start_time = time.time()
        
        # Count attempts by modifying the signing function temporarily
        max_attempts = 100
        attempts = 0
        
        # Simplified signing with attempt counting
        A, s1, s2, t = sk
        k, l = params.k, params.l
        gamma1, beta = params.gamma1, params.beta
        
        signature = None
        for attempt in range(max_attempts):
            attempts += 1
            
            # Sample mask and compute commitment
            mask_range = min(gamma1 // 1000, 100)
            y = vector([small_poly(mask_range) for _ in range(l)])
            w = A * y
            
            # Generate challenge and response
            c = generate_challenge(message, w, params)
            z = y + vector([c * s1_poly for s1_poly in s1])
            
            # Check rejection criteria
            z_norm = vector_infinity_norm(z)
            threshold = (gamma1 - beta) // 1000
            
            if z_norm < threshold:
                w_cs2 = w - vector([c * s2_poly for s2_poly in s2])
                if vector_infinity_norm(w_cs2) < params.gamma2 // 1000:
                    signature = (c, z)
                    break
        
        end_time = time.time()
        
        if signature:
            results['sign_times'].append(end_time - start_time)
            results['sign_attempts'].append(attempts)
            successful_signs += 1
        
        if i == 0:  # Report first attempt details
            print(f"   📈 First signing attempt: {attempts} iterations")
    
    results['sign_success_rate'] = successful_signs / num_trials
    
    if results['sign_times']:
        sign_avg = mean(results['sign_times']) * 1000
        attempts_avg = mean(results['sign_attempts'])
        print(f"   ✓ Average signing time: {sign_avg:.2f} ms")
        print(f"   ✓ Average attempts per signature: {attempts_avg:.1f}")
        print(f"   ✓ Success rate: {results['sign_success_rate']:.1%}")
    else:
        print(f"   ❌ No successful signatures generated")
    
    # Benchmark Verification (using last successful signature)
    if signature:
        print(f"\n🔍 Benchmarking Verification...")
        for trial in range(num_trials):
            start_time = time.time()
            is_valid = mldsa_verify(test_messages[0], signature, vk, params)
            end_time = time.time()
            results['verify_times'].append(end_time - start_time)
        
        verify_avg = mean(results['verify_times']) * 1000
        print(f"   ✓ Average verification time: {verify_avg:.2f} ms")
    
    return results

def analyze_rejection_sampling(params, num_trials=50):
    """Analyze rejection sampling statistics"""
    print(f"\n📊 Analyzing Rejection Sampling Statistics ({num_trials} trials)")
    
    # Generate a key pair for testing
    vk, sk = mldsa_keygen(params)
    A, s1, s2, t = sk
    
    attempts_list = []
    successful_signatures = 0
    
    for trial in range(num_trials):
        message = f"Analysis message {trial}"
        
        # Count attempts for this signature
        max_attempts = 200
        for attempt in range(max_attempts):
            # Sample and test
            mask_range = min(params.gamma1 // 1000, 100)
            y = vector([small_poly(mask_range) for _ in range(params.l)])
            w = A * y
            c = generate_challenge(message, w, params)
            z = y + vector([c * s1_poly for s1_poly in s1])
            
            z_norm = vector_infinity_norm(z)
            threshold = (params.gamma1 - params.beta) // 1000
            
            if z_norm < threshold:
                attempts_list.append(attempt + 1)
                successful_signatures += 1
                break
        else:
            # Failed after max_attempts
            attempts_list.append(max_attempts)
    
    if attempts_list:
        avg_attempts = mean(attempts_list)
        std_attempts = stdev(attempts_list) if len(attempts_list) > 1 else 0
        success_rate = successful_signatures / num_trials
        
        print(f"   📈 Average attempts per signature: {avg_attempts:.1f} ± {std_attempts:.1f}")
        print(f"   🎯 Success rate: {success_rate:.1%}")
        print(f"   📊 Min attempts: {min(attempts_list)}")
        print(f"   📊 Max attempts: {max(attempts_list)}")
        
        # Theoretical expectation
        # In practice, rejection probability ≈ 1/e for well-chosen parameters
        theoretical_attempts = 2.718  # Approximate
        print(f"   🔬 Theoretical expectation: ~{theoretical_attempts:.1f} attempts")
        
        return {
            'average_attempts': avg_attempts,
            'success_rate': success_rate,
            'std_attempts': std_attempts
        }
    
    return None

# Run benchmarks
print("=" * 70)
print("🚀 ML-DSA Performance Analysis")
print("=" * 70)

# Performance benchmark
benchmark_results = benchmark_mldsa(params, num_trials=3)

# Rejection sampling analysis
rejection_stats = analyze_rejection_sampling(params, num_trials=20)

print(f"\n" + "=" * 70)
print("📋 Performance Summary")
print("=" * 70)

if benchmark_results['keygen_times']:
    keygen_avg = mean(benchmark_results['keygen_times']) * 1000
    print(f"🔑 Key Generation: {keygen_avg:.2f} ms average")

if benchmark_results['sign_times']:
    sign_avg = mean(benchmark_results['sign_times']) * 1000
    print(f"📝 Signing: {sign_avg:.2f} ms average")

if benchmark_results['verify_times']:
    verify_avg = mean(benchmark_results['verify_times']) * 1000
    print(f"🔍 Verification: {verify_avg:.2f} ms average")

print(f"🎯 Overall success rate: {benchmark_results['sign_success_rate']:.1%}")

print("\n⚠️  Note: These timings are for educational implementations only!")
print("   Production implementations are orders of magnitude faster.")
print("=" * 70)