"""
pubkey.py

simulates a public key transaction
May 21, 2018
"""

import random

# arbitrary bounds
keymin = 1
keymax = 255

def nth_prime_number(n):
    # initial prime number list
    prime_list = [2]
    # first number to test if prime
    num = 3
    # keep generating primes until we get to the nth one
    while len(prime_list) < n:

        # check if num is divisible by any prime before it
        for p in prime_list:
            # if there is no remainder dividing the number
            # then the number is not a prime
            if num % p == 0:
                # break to stop testing more numbers, we know it's not a prime
                break

        # if it is a prime, then add it to the list
        # after a for loop, else runs if the "break" command has not been given
        else:
            # append to prime list
            prime_list.append(num)

        # same optimization you had, don't check even numbers
        num += 2

    # return the last prime number generated
    return prime_list[-1]

def main():

   # get primes
	index = random.randint(keymin, keymax)
	gp = nth_prime_number(index)
	index = random.randint(keymin, keymax)
  pp = nth_prime_number(index)
	print("g = " + str(gp))
	print("P = " + str(pp))

	alice_private = random.randint(keymin,keymax)
	print("Alice's private key = {{ %d }}"%(alice_private))

	bob_private = random.randint(keymin,keymax)
	print("Bob's private key = {{ %d }}"%(bob_private))

	ashared = pow(gp,alice_private)%pp
	print("Alice's sends to Bob: %d"%(ashared))

	bshared = pow(gp,bob_private)%pp
	print("Bob's sends to Alice: %d"%(bshared))

	asecret = pow(bshared,alice_private)%pp
	bsecret = pow(ashared,bob_private)%pp
	print("Bob and Alice both secretly use: {{ %d }} and {{ %d }}"%(asecret, bsecret))

if __name__ == '__main__':
    main()