Code.py

import numpy as np
from typing import Tuple

# ============ CONFIGURATION ============
# These parameters control security and key sizes
# Increase for more security, decrease for faster operations

DIMENSION = 8              # n: dimension of secret vector
SAMPLES = 10               # m: number of LWE samples
MODULUS = 3329            # q: modular arithmetic modulus (from ML-KEM standard)
ERROR_BOUND = 2           # β: bound on error values [-β, β]

# ============ HELPER FUNCTIONS ============

def generate_random_matrix(rows: int, cols: int, mod: int) -> np.ndarray:
    """Generate random matrix with entries in [0, mod)"""
    return np.random.randint(0, mod, size=(rows, cols))

def generate_error_vector(length: int, error_bound: int) -> np.ndarray:
    """Generate error/noise vector with small values in [-error_bound, error_bound]"""
    return np.random.randint(-error_bound, error_bound + 1, size=length)

def matrix_vector_mult_mod(A: np.ndarray, v: np.ndarray, mod: int) -> np.ndarray:
    """Multiply A*v with modular reduction"""
    result = np.dot(A, v)
    return result % mod

# ============ STEP 1: KEY GENERATION ============
# Alice creates a public and private key

def key_generation() -> Tuple[Tuple[np.ndarray, np.ndarray], np.ndarray]:
    """
    Generate LWE key pair.
    
    Returns:
        public_key: (A matrix, b vector) - can be shared publicly
        secret_key: s vector - must be kept secret
    """
    # Generate public matrix A
    A = generate_random_matrix(SAMPLES, DIMENSION, MODULUS)
    
    # Generate secret vector s (short values)
    s = generate_error_vector(DIMENSION, ERROR_BOUND)
    
    # Generate error vector e (small noise)
    e = generate_error_vector(SAMPLES, ERROR_BOUND)
    
    # Compute b = A*s + e (mod q)
    # This is the key equation in LWE
    b = (matrix_vector_mult_mod(A, s, MODULUS) + e) % MODULUS
    
    public_key = (A, b)
    secret_key = s
    
    return public_key, secret_key

# ============ STEP 2: ENCRYPTION ============
# Bob uses Alice's public key to encrypt a message

def encrypt(public_key: Tuple[np.ndarray, np.ndarray], message: int) -> Tuple[np.ndarray, int]:
    """
    Encrypt a single bit message using LWE public key.
    
    Args:
        public_key: (A, b) from Alice
        message: bit value (0 or 1)
    
    Returns:
        ciphertext: (u, v) pair
    """
    A, b = public_key
    
    # Generate random vector r (ephemeral randomness)
    r = generate_error_vector(SAMPLES, ERROR_BOUND)
    
    # Generate additional error vectors for security
    e1 = generate_error_vector(DIMENSION, ERROR_BOUND)
    e2 = np.random.randint(-ERROR_BOUND, ERROR_BOUND + 1)
    
    # Compute u = r^T * A + e1 (mod q)
    # This prevents attackers from recovering r via Gaussian elimination
    u = (matrix_vector_mult_mod(r, A, MODULUS) + e1) % MODULUS
    
    # Compute v = r^T * b + e2 + message * (q/2) (mod q)
    # Message encoding: 0 maps to 0, 1 maps to q/2 (approximately)
    # The q/2 offset helps with decryption despite the noise
    v = (np.dot(r, b) + e2 + message * (MODULUS // 2)) % MODULUS
    
    ciphertext = (u, v)
    return ciphertext

# ============ STEP 3: DECRYPTION ============
# Alice uses her secret key to decrypt Bob's message

def decrypt(secret_key: np.ndarray, ciphertext: Tuple[np.ndarray, int]) -> int:
    """
    Decrypt LWE ciphertext using secret key.
    
    Args:
        secret_key: s from Alice
        ciphertext: (u, v) from Bob
    
    Returns:
        message: recovered bit (0 or 1)
    """
    s = secret_key
    u, v = ciphertext
    
    # Compute m' = v - u^T * s (mod q)
    # If encryption was correct:
    #   m' = r^T*b + e2 + msg*(q/2) - (r^T*A + e1)^T * s (mod q)
    #   m' = r^T*(A*s + e) + e2 + msg*(q/2) - r^T*A*s - e1^T*s (mod q)
    #   m' = r^T*e + e2 - e1^T*s + msg*(q/2) (mod q)
    # The first three terms are small (noise), so m' ≈ msg*(q/2) with small error
    
    m_prime = (v - np.dot(u, s)) % MODULUS
    
    # Decode: if m' is closer to 0, message is 0; if closer to q/2, message is 1
    # We check which half of the modulus m_prime falls into
    threshold = MODULUS // 4
    
    if m_prime < threshold or m_prime > 3 * threshold:
        return 0
    else:
        return 1

# ============ PART 3: TESTING THE SCHEME ============

def test_lwe_encryption():
    """Test the complete LWE encryption scheme"""
    
    print("=" * 60)
    print("LWE ENCRYPTION SCHEME TEST")
    print("=" * 60)
    
    print(f"\nParameters:")
    print(f"  Dimension (n):        {DIMENSION}")
    print(f"  Samples (m):          {SAMPLES}")
    print(f"  Modulus (q):          {MODULUS}")
    print(f"  Error bound (β):      {ERROR_BOUND}")
    
    # Step 1: Key Generation
    print(f"\n[Step 1] Generating Keys...")
    public_key, secret_key = key_generation()
    A, b = public_key
    
    print(f"  ✓ Public key A shape: {A.shape}")
    print(f"  ✓ Public key b shape: {b.shape}")
    print(f"  ✓ Secret key s shape: {secret_key.shape}")
    
    # Step 2 & 3: Encryption and Decryption Tests
    print(f"\n[Step 2-3] Testing Encryption/Decryption...")
    
    test_cases = [0, 1, 1, 0, 1]
    all_passed = True
    
    for i, message in enumerate(test_cases, 1):
        ciphertext = encrypt(public_key, message)
        decrypted = decrypt(secret_key, ciphertext)
        
        status = "✓ PASS" if decrypted == message else "✗ FAIL"
        print(f"  Test {i}: Encrypt({message}) → Decrypt() = {decrypted} ... {status}")
        
        if decrypted != message:
            all_passed = False
    
    print(f"\n[Result] {'All tests passed!' if all_passed else 'Some tests failed - try adjusting parameters'}")
    
    # Show ciphertext sizes
    u, v = ciphertext
    print(f"\n[Ciphertext Info]")
    print(f"  u component size: {u.nbytes} bytes")
    print(f"  v component size: 1 integer (~4-8 bytes)")
    print(f"  Total ciphertext: ~{u.nbytes + 8} bytes")
    print(f"  Plaintext size: 1 bit")
    
    # Security Analysis
    print(f"\n[Security Notes]")
    print(f"  • This is a SIMPLIFIED scheme for learning purposes")
    print(f"  • Current parameters provide ~{64 + 8*DIMENSION} bits of entropy")
    print(f"  • For production use, see ML-KEM (Kyber) specifications")
    print(f"  • Production uses dimension=256, samples=256+, modulus=3329")
    
    return all_passed


# ============ MAIN EXECUTION ============

if __name__ == "__main__":
    # Run the basic test
    test_lwe_encryption()