import os

def rand_int(nbytes):
    return int.from_bytes(os.urandom(nbytes), byteorder='little')

def rand_less_than(upper_bound):
    while True:
        r = rand_int(((upper_bound).bit_length() + 7) // 8)
        if r < upper_bound:
            return r

def is_prime(p):
    for _ in range(5):
        a = rand_less_than(p)
        if not pow(a, p - 1, p) == 1:
            return False
    return True

def choose_parameters():
    while True:
        modulus = rand_int(1024 // 8)
        if is_prime(modulus):
            break
    base = rand_less_than(modulus)
    return base, modulus

class Peep:
    def __init__(self, base, modulus):
        self.base = base
        self.modulus = modulus
        self.private = rand_less_than(modulus)
        self.public = pow(base, self.private, modulus)

    def send(self):
        return self.public

    def receive(self, B):
        self.shared_secret = pow(B, self.private, self.modulus)

if __name__ == '__main__':
    base, modulus = choose_parameters()
    alice = Peep(base, modulus)
    bob = Peep(base, modulus)

    A = alice.send()
    B = bob.send()

    alice.receive(B)
    bob.receive(A)

    assert alice.shared_secret == bob.shared_secret