Problem: Implement RSA Key Generation and Modular Operations

The problem requires implementing RSA cryptosystem components—extended GCD (for integers/Gaussian integers), modular inverse, key generation (with prime validation and Carmichael function), and optimized modular exponentiation—to generate RSA key pairs and process integer/Gaussian integer messages via secure encryption/decryption using these foundational algorithms.

Background: ### Industry Background & Technical Context
These functions are core building blocks for RSA cryptography and its Gaussian integer variant—industry-standard tools for secure communication, digital signatures, and data encryption. They leverage optimized number-theoretic algorithms tailored for efficient computer implementation.

1. **Extended Euclidean Algorithm for Gaussian Integers (`extended_gcd`)**
- Gaussian integers are complex numbers with integer real/imaginary parts (a+bi, a,b ∈ ℤ), enabling advanced cryptosystems like Gaussian Integer RSA. This variant offers enhanced performance and resistance to some factorization attacks compared to standard RSA.
- The algorithm extends the integer extended Euclidean algorithm to compute the Gaussian greatest common divisor (gcd) and Bézout coefficients `s` and `t` such that `gcd(a,b) = a*s + b*t`. A critical adaptation here is rounding the real and imaginary components of the quotient (from `old_r / r`) to form a valid Gaussian integer quotient, which satisfies the Euclidean domain property for Gaussian integers and ensures correct algorithm termination.
- This function underpins modular inverse computation in Gaussian integer rings and coprimality checks, essential for Gaussian RSA key generation.

2. **Modular Inverse with Binary GCD (`mod_inverse`)**
- The modular inverse is a foundational operation in RSA key generation: the private key exponent `d` is the inverse of the public exponent `e` modulo λ(n) (Carmichael function) or φ(n) (Euler's totient function).
- The Binary GCD algorithm is a computer-optimized alternative to the standard Euclidean algorithm, using bitwise shifts and subtraction instead of division to speed up coprimality checks. It verifies that `e` and the modulus's totient/carmichael value are coprime (a requirement for the inverse to exist).
- The iterative extended Euclidean algorithm computes the inverse, followed by a consistency check to validate that `(e * result) % phi == 1`—a critical safeguard in cryptography to avoid invalid keys that break system security.

3. **RSA Key Generation (`generate_keys`)**
- RSA is the most widely deployed asymmetric cryptosystem, based on the computational hardness of factorizing large composite numbers formed from two (or more) large primes.
- The Miller-Rabin Primality Test is a probabilistic test optimized for large integers. For numbers < 2^64, using the witness set `{2,3,5,7,11,13,17,19,23,29,31,37}` provides deterministic accuracy, making it suitable for industry-grade prime verification.
- The Carmichael function λ(n) = LCM(p-1, q-1) is used instead of Euler's totient function φ(n) because it produces a smaller valid private exponent `d`, reducing computational overhead for private key operations while maintaining security.
- The function includes a framework for multi-prime RSA (extensible to 3+ primes), an industry technique to create shorter keys with equivalent security—ideal for resource-constrained devices like IoT sensors and mobile devices.

4. **RSA Processing with Montgomery Multiplication (`rsa_process`)**
- Modular exponentiation (`message^exponent mod modulus`) is the core operation of RSA encryption/decryption. Optimizing this step is critical for performance, especially with large 2048+ bit keys.
- Montgomery Multiplication is an industry-standard optimization for modular arithmetic. It replaces expensive modular reductions with integer division and bit shifts, significantly speeding up operations for odd moduli (the standard for RSA, as moduli are products of odd primes).
- Support for Gaussian integers enables operations for Gaussian RSA, while noise injection (adding small random values during multiplication/squaring) is a countermeasure against side-channel attacks (e.g., timing attacks). This noise randomizes execution time, preventing attackers from inferring key details via computation timing patterns.
- Strict input validation ensures compliance with cryptographic standards: messages can be integers (standard RSA) or Gaussian integers (Gaussian RSA), moduli must be valid integers >1, and exponents must be non-negative integers.

---

from typing import Union, Tuple
import random

Input/Output: The problem accepts two distinct prime integers \(p\) and \(q\), an integer \(e\), an integer or Gaussian integer (with integer real/imaginary parts) \(message\_int\), a non-negative integer \(key\_exponent\), and a positive integer \(modulus\). It validates inputs (integer types, \(p/q\) primality via Miller-Rabin, \(e\) coprimality to \(\lambda(n) = \text{LCM}(p-1, q-1)\) via binary GCD, \(modulus > 1\)). It computes \(d\) as \(e\)’s modular inverse modulo \(\lambda(n)\) using the extended Euclidean algorithm, generates public \((e, n)\) and private \((d, n)\) keys (\(n = p \times q\)), and processes \(message\_int\) via modular exponentiation (Montgomery multiplication for odd moduli, Gaussian integer operations for complex inputs) to return an integer or complex result. It returns key pairs (or Nones if validation fails) and processed messages, raising TypeErrors/ValueErrors for invalid inputs and ensuring consistency via checks (e.g., \((e \times d) \mod \lambda(n) = 1\)).

Required dependencies:
from typing import Union, Tuple
import random

Sub-steps:
step_number: 1
step_description_prompt: Please write a Python function named extended_gcd. It accepts two inputs (int or Gaussian integer with integer real/imag parts). Using the extended Euclidean algorithm (rounding quotient to Gaussian integer each step), compute gcd, s, t such that gcd = a*s + b*t. Return (gcd, s, t) and handle invalid input errors.
function_header: def extended_gcd(a, b):
test_cases: test_egcd = [
    (48, 18, ((-6+0j), (1+0j), (-3+0j))),
    (101, 10, ((1+0j), (1+0j), (-10+0j))),
    (0, 5, ((5+0j), 0j, (1+0j)))
]
for a, b, expected in test_egcd:
    res = extended_gcd(a, b)
    assert len(res) == 3
    assert abs(res[0] - expected[0]) < 1e-10 or (res[0] == expected[0]), f"gcd mismatch: actual={res[0]}, expected={expected[0]}"
    assert abs(res[1] - expected[1]) < 1e-10 or (res[1] == expected[1]), f"x coefficient mismatch: actual={res[1]}, expected={expected[1]}"
    assert abs(res[2] - expected[2]) < 1e-10 or (res[2] == expected[2]), f"y coefficient mismatch: actual={res[2]}, expected={expected[2]}"
    gcd, x, y = res
    bezout_check = a * x + b * y
    assert abs(bezout_check - gcd) < 1e-10, f"Bezout's identity failed: {a}*{x} + {b}*{y} = {bezout_check}, expected gcd={gcd}"

edge_cases = [
    (1, 1, (1+0j)),
    (1, 0, (1+0j)),
    (0, 0, (0+0j)),
    (17, 13, (1+0j)),
    (100, -25, (25+0j)),
    (2**31-1, 2**30, (1+0j)),
]
for a, b, exp_gcd in edge_cases:
    try:
        res = extended_gcd(a, b)
        gcd, x, y = res
        if a == 0 and b == 0:
            assert abs(gcd) < 1e-10, f"gcd(0,0) should be 0, got {gcd}"
        else:
            bezout = a * x + b * y
            assert abs(bezout - gcd) < 1e-10, f"Bezout identity failed for ({a}, {b}): {a}*{x} + {b}*{y} = {bezout}, expected {gcd}"
            if abs(a) > 1e-10:
                assert abs((a / gcd).real - round((a / gcd).real)) < 1e-10, f"gcd({a}, {b})={gcd} does not divide {a}"
            if abs(b) > 1e-10:
                assert abs((b / gcd).real - round((b / gcd).real)) < 1e-10, f"gcd({a}, {b})={gcd} does not divide {b}"
    except (ZeroDivisionError, ValueError, TypeError):
        pass

step_number: 2
step_description_prompt: Please write a Python function named mod_inverse that computes the modular inverse of e modulo phi using the iterative Extended Euclidean Algorithm with optimizations. It checks for integer inputs, positive phi, coprimality via binary GCD, computes the inverse, and performs a consistency check before returning the non-negative result.
function_header: def mod_inverse(e, phi):
test_cases: test_inverse = [
    (3, 11, 4),        # 3*4 = 12 ≡ 1 (mod 11)
    (7, 40, 23),       # 7*23 = 161 ≡ 1 (mod 40)
    (17, 3120, 2753)   # 17*2753 = 46801 ≡ 1 (mod 3120)
]
for e, phi, expected in test_inverse:
    assert mod_inverse(e, phi) == expected

step_number: 3
step_description_prompt: Please write a Python function named generate_keys(p: int, q: int, e: int) -> Union[tuple[tuple[int, int], tuple[int, int]], tuple[None, None, None]]. Validate inputs (p/q are distinct primes, e integer), use Miller-Rabin test for primality, compute λ(n)=LCM(p-1,q-1), check e coprime to λ(n), find d as inverse of e mod λ(n), return (public=(e,n), private=(d,n)) or Nones if any step fails.
function_header: def generate_keys(p, q, e):
test_cases: test_keys = [
    (3, 11, 7, ((7, 33), (3, 33))),       
    (61, 53, 17, ((17, 3233), (413, 3233))),  
    (5, 7, 5, ((5, 35), (5, 35)))         
]
for p, q, e, expected in test_keys:
    pub, pri = generate_keys(p, q, e)
    pub_rounded = tuple(round(float(x.real), 4) if isinstance(x, complex) else round(float(x), 4) for x in pub)
    pri_rounded = tuple(round(float(x.real), 4) if isinstance(x, complex) else round(float(x), 4) for x in pri)
    expected_0_rounded = tuple(round(float(x.real), 4) if isinstance(x, complex) else round(float(x), 4) for x in expected[0])
    expected_1_rounded = tuple(round(float(x.real), 4) if isinstance(x, complex) else round(float(x), 4) for x in expected[1])
    assert pub_rounded == expected_0_rounded, f"pub mismatch: actual={pub_rounded}, expected={expected_0_rounded}"
    assert pri_rounded == expected_1_rounded, f"pri mismatch: actual={pri_rounded}, expected={expected_1_rounded}"

step_number: 4
step_description_prompt: Please write a Python function named rsa_process that takes message_int, key_exponent, modulus as inputs and returns Union[int, complex]. It validates inputs, processes integer/Gaussian integer messages via modular exponentiation, uses Montgomery multiplication for odd moduli, adds noise in calculations, and handles modulus cases properly.
function_header: def rsa_process(message_int, key_exponent, modulus):
test_cases: test_rsa = [
    (2, 3, 10, 8),          
    (65, 17, 3233, 2790),   
    (2790, 2753, 3233, 65)  
]
for m, e, n, expected in test_rsa:
    assert rsa_process(m, e, n) == expected

General solution:

General tests: