from Crypto.Util import number
import random
from fractions import gcd
import time
import matplotlib.pyplot as plt

# Key Length Choice 64, 128, 256, 512
n_length = 64


def egcd(a, b):
    if a == 0:
        return (b, 0, 1)
    else:
        g, y, x = egcd(b % a, a)
        return (g, x - (b // a) * y, y)


def modinv(a, m):
    g, x, y = egcd(a, m)
    if g != 1:
        raise Exception('modular inverse does not exist')
    else:
        return x % m


def get_keys():
    p = number.getPrime(n_length)
    q = number.getPrime(n_length)

    n = p * q
    phi = (p - 1) * (q - 1)

    e = random.randrange(1, phi)

    g = gcd(e, phi)
    while g != 1:
        e = random.randrange(1, phi)
        g = gcd(e, phi)

    d = modinv(e, phi)

    return ((e, n), (d, n))


def encrypt(pk, plaintext):
    key, n = pk
    cipher = [pow(ord(char), key, n) for char in plaintext]
    return cipher


def decrypt(pk, ciphertext):
    key, n = pk
    plain = [chr(pow(char, key, n)) for char in ciphertext]
    print(plain)
    return ''.join(plain)


public, private = get_keys()
print(public, private)

timing = []
message_len = []
# Parameter tuning on range and base message
for i in range(101):
    base_message = "FoundationsOFCybersecurity BaseMessage of Length50"
    print(len(base_message))
    message = base_message + 'x' * i
    print(message)
    start = time.time()
    cipher = encrypt(public, message)
    time_taken = time.time() - start
    timing.append(time_taken)
    message_len.append(len(message))
    print(time_taken)
    print("Message Len: ", len(message))

plt.xlabel("Message Length")
plt.ylabel("Time in sec")
plt.plot(message_len, timing)
plt.show()