code.py

import random
def is_prime(num):
    for i in range(1000):
        a = random.randint(2, num - 1)
        if pow(a, num - 1, num) != 1:
            return False
    return True

def dh_exchange(p):
    g = random.randint(1, p - 1)

    # Or secret
    a = random.randint(1, p - 1)

    # Shir secret
    b = random.randint(1, p - 1)
    
    x = pow(g, a, p)
    y = pow(g, b, p)
    key_a = pow(y, a, p)
    key_b = pow(x, b, p)
    if key_a != key_b:
        print("Waaaat???")
    return g, a, b, x, y, key_a

def find_prime(n):
    while True:
        candidate = random.randrange(2 ** (n-1), 2 ** n)
        if is_prime(candidate):
            return candidate

def crack_dh(p, g, x, y):
    for a in range(1, p - 1):
        if pow(g, a, p) == x:
            return a
    return None