Cryptography - HW3.1 (dlog.py)

import math

# Find x such that g^x = h (mod p)
# 0 <= x <= max_x
def discrete_log(g, h, p, max_x):
    # let B be maximum range to iterate over (2^20)
    B = int(math.sqrt(max_x))

    # Equation to solve: h/(g^x1) = (g^B)^x0

    ''' Iterate Over Left Part of Equation [h/(g^x1)] '''
    # table to save left part equation solutions
    hash_table = {}
    # g^(-1) = g^(p-2)
    g_inv = pow(g, p - 2, p)
    h_g_x1 = h
    # Iterate and fill in hash table
    for x1 in range(B):
        # fill hash table
        hash_table[h_g_x1] = x1
        # h * ((g^-1) ^ x1)
        h_g_x1 *= g_inv
        h_g_x1 %= p

    ''' Iterate Over Right Part of Equation [(g^B)^x0] '''
    # g^B
    g_B = pow(g, B, p)
    # Iterate to find such x0 that (g^b)^x0 is in our hash table
    g_B_x0 = 1
    for x0 in range(B):
        # if hash table contains given number then we have foun solution
        if g_B_x0 in hash_table:
            x1 = hash_table[g_B_x0]
            return x0 * B + x1
        # (g^B)^x0
        g_B_x0 = (g_B_x0*g_B) % p


def main():
    p = 13407807929942597099574024998205846127479365820592393377723561443721764030073546976801874298166903427690031858186486050853753882811946569946433649006084171
    g = 11717829880366207009516117596335367088558084999998952205599979459063929499736583746670572176471460312928594829675428279466566527115212748467589894601965568
    h = 3239475104050450443565264378728065788649097520952449527834792452971981976143292558073856937958553180532878928001494706097394108577585732452307673444020333
    max_x = 1 << 40  # 2^40
    x = discrete_log(g, h, p, max_x)

    # check if discrete logarithm is correct
    ans = pow(g, x, p)
    if (ans == h):
        print("Discrete Logarithm is correct!")
        print("Result : x={}".format(x))
    else:
        print("Discrete Logarithm is incorrect!")


if __name__ == '__main__':
    main()