Sample diffie-helman implementation in Python

dh.py

# Importing necessary libraries
import random

# Function to compute the modular exponentiation (a^b mod p)
def mod_exp(a, b, p):
    return pow(a, b, p)

def diffie_hellman(p, g):
    # Alice chooses a private key
    a = random.randint(1, p - 1)
    # Alice computes her public key
    A = mod_exp(g, a, p)
    
    # Bob chooses a private key
    b = random.randint(1, p - 1)
    # Bob computes his public key
    B = mod_exp(g, b, p)
    
    # Alice and Bob exchange public keys and compute the shared secret
    shared_secret_alice = mod_exp(B, a, p)
    shared_secret_bob = mod_exp(A, b, p)
    
    # Both shared secrets should be equal
    assert shared_secret_alice == shared_secret_bob
    
    return shared_secret_alice

p = 23
g = 5

# Running the Diffie-Hellman key exchange
shared_secret = diffie_hellman(p, g)
print(f'The shared secret is: {shared_secret}')