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

n_length = 512

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 = []
for i in xrange(100):
    base_message = "A random message for Foundations of Cybersecurity"
    message = base_message + 'i'*i
    start = time.time()
    cipher = encrypt(public, message)
    time_taken = time.time() - start
    timing.append(time_taken)
    message_len.append(len(message))
    print time_taken

plt.xlabel("Message Lenght")
plt.ylabel("Time")
plt.plot(message_len, timing)
plt.show()