diffie_helman.py

# Diffie-Hellman Key Exchange Implementation
# This script implements the Diffie-Hellman key exchange algorithm for secure key establishment between two parties over an insecure channel.

import random
from copy import deepcopy

# Function to calculate modular exponentiation (a^n mod m)
def mod_pow(a, n, m):
    """
    Calculate a^n mod m using the binary exponentiation algorithm.
    
    Parameters:
        a (int): Base.
        n (int): Exponent.
        m (int): Modulus.
        
    Returns:
        int: Result of a^n mod m.
    """
    a = a % m
    result = 1
    while n > 0:
       if n % 2 == 1:
          result = (result * a) % m
       a = (a * a) % m
       n = n // 2
    return result

# Function to calculate the extended greatest common divisor (GCD) of two numbers
def extended_gcd(a, b):
    """
    Calculate the extended greatest common divisor (GCD) of two numbers using the Euclidean algorithm.
    
    Parameters:
        a (int): First number.
        b (int): Second number.
        
    Returns:
        tuple: (gcd, x, y) where gcd is the GCD of a and b, and x, y are the Bezout coefficients such that ax + by = gcd.
    """
    if b == 0:
       return (a, 1, 0)
    d, m, n = extended_gcd(b, a % b)
    return (d, n, m - a // b * n)

# Function to calculate modular multiplicative inverse
def mod_inv(a, x):
    """
    Calculate the modular multiplicative inverse of 'a' modulo 'x'.
    
    Parameters:
        a (int): Integer.
        x (int): Modulus.
        
    Returns:
        int: Modular multiplicative inverse of 'a' modulo 'x'.
    """
    d, m, n = extended_gcd(a, x)
    if d != 1:
       print("Values 'a' and 'x' are not mutually prime!")
    else:
       return m % x    

# Function to perform the Miller-Rabin primality test
def miller_rabin(n, k):
    """
    Perform the Miller-Rabin primality test to determine whether 'n' is prime.
    
    Parameters:
        n (int): Number to test for primality.
        k (int): Number of iterations for the test.
        
    Returns:
        bool: True if 'n' is probably prime, False otherwise.
    """
    if n <= 3:
       if n == 1:
          return False
       return True

    d = n - 1
    r = 0
    while d % 2 == 0:
       r = r + 1
       d = d // 2
    for i in range(k):
       a = random.randrange(2, n - 1)
       x = mod_pow(a, d, n)
       
       if x == 1 or x == n - 1:
          continue
       
       witness = True
       
       for j in range(r - 1):
           x = mod_pow(x, 2, n)
           if x == 1:
              return False
           if x == n - 1: # n - 1 = -1 (mod n)
              witness = False
              break
       if witness:
          return False
    return True

# Function to generate a prime number using the Miller-Rabin primality test
def get_prime(limit, k=20):
    """
    Generate a random prime number within the specified limit using the Miller-Rabin primality test.
    
    Parameters:
        limit (int): Upper limit for generating the prime number.
        k (int): Number of iterations for the Miller-Rabin primality test.
        
    Returns:
        int: A prime number.
    """
    is_prime = False
    while not is_prime:
        n = random.randrange(limit)
        is_prime = miller_rabin(n, k)
    return n

# Class implementing the Diffie-Hellman key exchange protocol
class Diffie_Helman:
    def __init__(self, q, g):
        """
        Initialize Diffie-Hellman parameters and generate public and private keys.
        
        Parameters:
            q (int): Prime modulus.
            g (int): Generator.
        """
        self.q = q
        self.g = g
        self.private_key = random.randrange(2, self.q)
        self.public_key = mod_pow(self.g, self.private_key, self.q)
    
    def get_key(self, pub_key):
        """
        Compute the shared secret key using the received public key.
        
        Parameters:
            pub_key (int): Public key received from the other party.
        
        Returns:
            int: Shared secret key.
        """
        return mod_pow(pub_key, self.private_key, self.q)

# Main function to demonstrate the Diffie-Hellman key exchange
def main():
    # Set the upper limit for generating prime numbers
    limit = 2 ** 256
    
    # Generate a prime number 'q' and a random generator 'g'
    q = get_prime(limit)
    g = random.randrange(2, q)
    
    # Create two instances of Diffie_Helman for two parties
    A = Diffie_Helman(q, g)
    B = Diffie_Helman(q, g)
    
    # Print private and public keys for both parties
    print("Party A: \n\tPrivate Key =", A.private_key, "\n\tPublic Key =", A.public_key)
    print("Party B: \n\tPrivate Key =", B.private_key, "\n\tPublic Key =", B.public_key)
    
    # Compute and print shared secret keys for both parties
    print("Shared Secret Key (A):", A.get_key(B.public_key))
    print("Shared Secret Key (B):", B.get_key(A.public_key))

if __name__ == "__main__":
   main()