#!/usr/bin/env python3

import matplotlib.pyplot as plt
from matplotlib.pyplot import rcParams
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
from Crypto.Cipher import AES
import base64
import csv
import struct
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 os
#from multiprocessing import Pool

def u_decrypt(arg):
    so_file2 = 'EES401/URG_decrypt.so'
    arg_len = len(arg)
    print("Length ", arg_len)
    u_dec = CDLL(so_file2)
    u_dec.main.restype = c_char_p
    u_dec.main.argtype = c_char_p
    c_return = u_dec.main(arg)
    _ctypes.dlclose(u_dec._handle)
    print("C Return ",c_return)
    return c_return

def u_decrypt593(arg):
    so_file2 = 'EES593/URG_decrypt.so'
    arg_len = len(arg)
    print("Length ", arg_len)
    u_dec = CDLL(so_file2)
    u_dec.main.restype = c_char_p
    u_dec.main.argtype = c_char_p
    c_return = u_dec.main(arg)
    _ctypes.dlclose(u_dec._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 file_handler(f_name, num):
    num = str(num).zfill(2)
    f = open("./cipher_bak/"+f_name+"_"+num+".dat", "rb")
    #file_val = list(f.read())
    #file_val = bytearray(file_val)
    file_val = f.read()
    return file_val

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

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

aes_key = "1234123412341234"
IV      = 16 * '\x00'

base_name           = "cipher_EES401"
base_name_593       = "cipher_EES593"
base_name_RSA3      = "RSA_3072"
base_name_RSA7      = "RSA_7680"
base_name_aes       = "AES_128"
base_name_Fernet    = "Fernet_128"

i = 1
while i <= 16:
    plain_length.append(i)
    i = i+1

for i in plain_length:

    plain = rsa_file_handler(base_name, i)
    plain2 = rsa_file_handler(base_name_593, i)
    fernet_key = Fernet_process()
    #print("Cip now : ", plain)

    start = time.perf_counter()
    print("NTRU PLAIN ",u_decrypt(plain))
    end=time.perf_counter()-start
    NTRU_var.append(end)

    start=time.perf_counter()
    print("NTRU PLAIN 593 ",u_decrypt593(plain2))
    end=time.perf_counter()-start
    NTRU593_var.append(end)

    fernet_cipher = rsa_file_handler(base_name_Fernet, i)
    time_s = time.perf_counter()
    Fernet_data = Fernet(fernet_key).decrypt(fernet_cipher)
    time_e = time.perf_counter()
    end = time_e - time_s
    Fernet_var.append(time_e - time_s)

    print("Fernet Result : ", Fernet_data)
    print("Fernet Cipher Length : ",len(fernet_cipher))
    print("Fernet Decryption Time : ",end)
    print(" ")

    aes_cipher = rsa_file_handler(base_name_aes, i)
    start=time.perf_counter()
    AES128= AES.new(aes_key, AES.MODE_CBC, IV).decrypt(aes_cipher)
    end=time.perf_counter()-start
    AES_var.append(end)

    print("AES Result : ", AES128.decode('ascii'))
    print("AES128 Cipher Length : ",len(aes_cipher))
    print("AES128 Decryption Time : ",end)
    print(" ")

    rsa3072 = open("RSA3072_priv.pem", "r").read()
    start=time.perf_counter()
    rsa_priv = RSA.importKey(rsa3072)
    RSA_dec = rsa_priv.decrypt(rsa_file_handler(base_name_RSA3, i))
    end=time.perf_counter()-start
    RSA_var.append(end)
    print("3072 Result : ", RSA_dec)
    print("RSA Cipher length", len(plain))
    print("RSA Decryption Time : ", end)
    print(" ")

    rsa7680 = open("RSA7680_priv.pem", "r").read()
    start=time.perf_counter()
    RSA2_priv = RSA.importKey(rsa7680)
    RSA2_dec = RSA2_priv.decrypt(rsa_file_handler(base_name_RSA3, i))
    end=time.perf_counter()-start
    RSA2_var.append(end)
    print("7680 Result : ", RSA_dec)
    print("RSA2 Plain length", len(plain))
    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     ]


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_Decryption_Time.csv", "w") as f:
    writer = csv.writer(f)
    writer.writerow(header)
    for row in all_time:
        writer.writerow(row)

print(NTRU_var)
print(AES_var)
y =np.array(plain_length)
plt.figure(figsize=(11,5))
rcParams['legend.fontsize'] = '50'
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_DEC.eps', format='eps', dpi=900)