# Arup Guha
# 10/28/2020
# Example of the Diffie-Hellman Key Exchange

# Note: for today, we assume that someone else knows a and p and can
#       just enter it!

# Returns base to the power power mod mod.
def modPow(base,power,mod):

    # Base case.
    if power == 0:
        return 1%mod

    # Time savings by getting square root and squaring.
    if power%2 == 0:
        tmp = modPow(base, power//2, mod)
        return (tmp*tmp)%mod

    # Regular case.
    return (modPow(base,power-1,mod)*base)%mod


def main():

    # Get public keys; I don't check them.
    p = int(input("Enter the prime number for Diffie-Hellman"))
    a = int(input("Enter the primitive root used for Diffie-Hellman"))

    # Get private keys.
    alice = int(input("Enter Alice's secret key in between 1 and "+str(p)))
    bob = int(input("Enter Bob's secret key in between 1 and "+str(p)))

    # Show what Alice sends
    C1 = modPow(a, alice, p)
    print("Alice sends to Bob",C1)

    # Show what Bob sends
    C2 = modPow(a, bob, p)
    print("Bob sends to Alice", C2)

    # Show what Alice computes.
    aliceAns = modPow(C2, alice, p)
    print("Alice calculates the secret key of", aliceAns)

    # And what bob computes.
    bobAns = modPow(C1, bob, p)
    print("Bob calculates the secret key of", bobAns)

main()