#!/usr/bin/env python3
import matplotlib.pyplot as plt
from matplotlib.pyplot import rcParams
from matplotlib.ticker import FormatStrFormatter
from cryptography.fernet import Fernet
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
import numpy as np
import base64
import time
import string
import random
import aes
import numpy as np
from ctypes import *
import _ctypes
from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_OAEP
import csv
#from multiprocessing import Pool

# C stuff

def u_encrypt(arg, num):
    if len(str(num)) < 2:
        num = str(num).zfill(2)
    num = str(num)
    so_file = 'EES401/URG_encrypt.so'
    u_enc = CDLL(so_file)
    u_enc.main.restype = c_double
    u_enc.main.argtype = [c_char_p, c_char_p]
    c_return = u_enc.main(arg.encode('utf-8'), num.encode('utf-8'))
    _ctypes.dlclose(u_enc._handle)
    print("C Return ", c_return)
    return c_return

def u_encrypt593(arg, num):
    if len(str(num)) < 2:
        num = str(num).zfill(2)
    num = str(num)
    so_file = 'EES593/URG_encrypt.so'
    u_enc = CDLL(so_file)
    u_enc.main.restype = c_double
    u_enc.main.argtype = [c_char_p, c_char_p]
    c_return = u_enc.main(arg.encode('utf-8'), num.encode('utf-8'))
    _ctypes.dlclose(u_enc._handle)
    print("C Return ", c_return)
    return c_return

def Fernet_process():
    password="urg123"
    pass_byte=password.encode()
    salt=b'salt_'
    kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100,
            backend=default_backend()
            )
    key = base64.urlsafe_b64encode(kdf.derive(pass_byte))
    return key

def randString(length):
    letters = string.ascii_lowercase
    return ''.join(random.choice(letters) for i in range(length))

def file_handler(f_name, num):
    num = str(num).zfill(2)
    f = open("./plain/"+f_name+"_"+num+".txt", "r")
    return f.read()

def rsa_file_handler(f_name, num, cipher):
    num = str(num).zfill(2)
    f = open("./cipher/"+f_name+"_"+num+".dat", "wb")
    f.write(cipher)


NTRU_var        = []
NTRU593_var     = []
RSA_var         = []
RSA2_var        = []
AES_var         = []
Fernet_var      = []

aes_key = "1234123412341234"

base_name = "plain"
base_name_RSA3 = "RSA_3072"
base_name_RSA7 = "RSA_7680"
base_name_aes = "AES_128"
base_name_Fernet = "Fernet_128"

# i = 3
# while i <= 5 :
#     plain_length.append(i*1)
#     i = i + 1
n = input("Enter number of file ")
n = int(n)
for i in range(1, n+1):

    plain = file_handler(base_name, i)
    print("Plain now : ", plain)
    fernet_key = Fernet_process()

    x = 16 - len(plain)
    aes_plain = plain + (" "*x)

    NTRU_var.append(u_encrypt(plain, i))

    print("NTRU EES401 Plain Length : ",len(plain))
    print(" ")


    NTRU593_var.append(u_encrypt593(plain, i))

    print("NTRU EES593 Plain Length : ",len(plain))
    print(" ")

    time_s = time.perf_counter()
    send_data = Fernet(fernet_key).encrypt(plain.encode())
    rsa_file_handler(base_name_Fernet, i, send_data)
    time_e = time.perf_counter()
    end = time_e - time_s
    Fernet_var.append(time_e - time_s)

    print("Fernet Plain Length : ",len(plain))
    print("Fernet Key Length : ",len(fernet_key))
    print("Fernet Encryption Time : ",end)
    print(" ")

    start=time.perf_counter()
    AES128=aes.aes(aes_key, aes_plain)
    end=time.perf_counter()-start
    rsa_file_handler(base_name_aes, i, AES128)
    AES_var.append(end)

    print("AES128 Plain Length : ",len(aes_plain))
    print("AES128 Key Length : ",len(aes_key))
    print("AES128 Encryption Time : ",end)
    print(" ")

    rsa3072 = open("RSA3072_pub.pem", "r").read()
    start=time.perf_counter()
    rsa_pub = RSA.importKey(rsa3072)
    RSA_enc = rsa_pub.encrypt(plain.encode('utf8'), 32)
    #RSA_enc = rsa_pub.encrypt(plain.encode('utf8'), 32)
    end=time.perf_counter()-start
    cipher = RSA_enc[0]
    #cipher = str(cipher).encode('utf-8')
    rsa_file_handler(base_name_RSA3, i, cipher)
    RSA_var.append(end)
    print("RSA Plain length", len(plain))
    print("RSA Key Length : ",len(str(rsa_pub.n)))
    print("RSA Encryption Time : ", end)
    print(" ")

    rsa7680 = open("RSA7680_pub.pem", "r").read()
    start=time.perf_counter()
    RSA2_pub = RSA.importKey(rsa7680)
    #RSA2_pub = RSA2_key.publickey().exportKey("PEM")
    RSA2_enc = RSA2_pub.encrypt(plain.encode('utf8'), 32)
    #RSA2_enc = PKCS1_OAEP.new(RSA2_pub).encrypt(plain.encode('utf8'), 32)
    end=time.perf_counter()-start
    #print("RSA Stufff ", RSA2_enc[0])
    cipher = RSA2_enc[0]
    #cipher = str(cipher).encode('utf-8')
    rsa_file_handler(base_name_RSA7, i, cipher )
    RSA2_var.append(end)
    print("RSA2 Plain length", len(plain))
    print("RSA2 Key Length : ",len(str(RSA2_pub.n.bit_length)))
    print("RSA2 Encryption Time : ", end)
    print(" ")

NTRU_var    = [ round(i * 1000, 3) for i in  NTRU_var       ]
NTRU593_var = [ round(i * 1000, 3) for i in  NTRU593_var    ]
RSA_var     = [ round(i * 1000, 3) for i in  RSA_var        ]
RSA2_var    = [ round(i * 1000, 3) for i in  RSA2_var       ]
AES_var     = [ round(i * 1000, 3) for i in  AES_var        ]
Fernet_var  = [ round(i * 1000, 3) for i in  Fernet_var     ]

plain_length = range(1, n+1)

header = ["NTRU 401","NTRU 593", "RSA_2048", "RSA_7680", "AES128", "Fernet128"]
all_time = zip(NTRU_var, NTRU593_var, RSA_var, RSA2_var, AES_var, Fernet_var)
with open("Benchmark_Encryption_Time.csv", "w") as f:
    writer = csv.writer(f)
    writer.writerow(header)
    for row in all_time:
        writer.writerow(row)

v = np.linspace(0, 0.5, 5)
y =np.array(plain_length)
print(NTRU_var)
print(AES_var)
plt.figure(figsize=(11,5))
rcParams['font.family'] = 'sans-serif'
rcParams['font.size'] = '16'
plt.scatter(y ,AES_var      ,marker="x" ,label='AES 128'   )
plt.scatter(y ,Fernet_var   ,marker="^" ,label='Fernet 128')
plt.scatter(y ,RSA_var      ,marker="P" ,label='RSA 2048'  )
plt.scatter(y ,RSA2_var     ,marker="d" ,label='RSA 7680'  )
plt.scatter(y ,NTRU_var     ,marker="v" ,label='NTRU 401'  )
plt.scatter(y ,NTRU593_var  ,marker="o" ,label='NTRU 593'  )
plt.plot(y, AES_var,     linestyle=':')
plt.plot(y, Fernet_var,  linestyle=':')
plt.plot(y, RSA_var,     linestyle=':')
plt.plot(y, RSA2_var,    linestyle=':')
plt.plot(y, NTRU_var,    linestyle=':')
plt.plot(y, NTRU593_var, linestyle=':')
plt.xticks(list(range(1,17)))
plt.xlabel("Length Plain text")
plt.ylabel("Time (ms)")
plt.yscale('log')
plt.gca().set_yticks([0.03, 1.90], minor=True)
plt.legend(bbox_to_anchor=(1.04,1), loc="upper left", fontsize='14')
plt.subplots_adjust(right=0.8)
plt.tight_layout()
plt.savefig('time_comparsion_ENC.eps', format='eps', dpi=900)
plt.show()