diffie hellman python

def dif_hel(upper_bound = 1e4):

    from random import randrange
    primes = get_primes(n = upper_bound, mode = 1)
    global a, A, b, B, s
    # create list of all primes up to upper_bound

    while True:
        p = input("Enter a prime number in range 3 to %r or enter 'r' for random: "
                      % primes[-1])
        if p == 'r' or p == 'R':
            p = primes[randrange(1, len(primes)-1)]
            break
        # set p to random number from 'primes'
        else:
            p = int(p)
            if p > primes[-1] or p not in primes[1:]:
                print("%r is out of bounds or not prime." % p)

                if p in range(3,primes[-1]):
                    next_prime = 0
                    i = 0
                    while next_prime == 0:
                        i += 1
                        if primes.count(p+i) > 0:
                            next_prime = p+i
                            break
                    answer = input("Use %r instead? y/n: " % next_prime)
                    if answer == 'y' or answer == 'Y' or answer == str(next_prime):
                        p = next_prime
                        break
            else:
                break
    print("Prime number:", p, "\n")

    possible_roots = []
    for base in primes[:4]:
        if is_primitive(g = base, n = p) and base < p:
            possible_roots.append(base)
    if len(possible_roots) == 0:
        print("Error: found no suitable base for", p)
        return
    elif len(possible_roots) == 1:
        print("Found only one suitable base for", p)
        g = possible_roots[0]
    else:
        while True:
            g = input("Enter a base from %r or enter 'r' for random: " % possible_roots)
            if g == 'r' or g == 'R':
                g = possible_roots[randrange(len(possible_roots)-1)]
                break
            # set g to random number from 'possible_roots'
            else:
                g = int(g)
                if g not in possible_roots:
                    print("%r is out of bounds." % p)
                else:
                    break
    print("Base:", g, "\n")

    a = int(input("Enter your secret number (do not divulge!): "))
    A = g**a % p
    b = randrange(int(upper_bound/2))
### these next two lines may belong in their own mode:
    B = g**b % p
    s = B**a % p
#######
    print("Your public key:", A, "\n")
    B = int(input("Enter partner's public key: "))
    print("Your partner's public key:", B, "\n")
    s = B**a % p
    print("Shared secret:", s)


def get_primes(n, mode=0):
    """ mode 0: generates n primes
        mode 1: generates all primes up to n """

    primes = []     # = list for n primes
    dividend = 2    # = integer to be tested
    is_prime = 1    # 'dividend' is presumed prime at the outset

    while len(primes) < n:
        if is_prime == 1:
            primes.append(dividend)
        # if no divisor has been found for 'dividend', add to 'primes'

        if mode == 1 and dividend >= n:
            return primes
        # (see above for mode explanation)

        dividend += 1
        is_prime = 1

        for divisor in primes:
            if divisor > dividend**(1/2):
                break
            # stop checking once 'divisor' > square root of 'dividend'

            if dividend % divisor == 0:
                is_prime = 0
                break
            # break out of loop if divisor is found

    return primes

def is_primitive(g,n):

    group = [g**k % n for k in range(1,n)]

    if group[0] in group[1:]:
        return False
    else:
        return True
    # check whether the period of g**k % n < n-1