import paho.mqtt.client as paho
import paho.mqtt.publish as pahopub
#from simplecrypt import encrypt, decrypt
import base64
from ctypes import *
import _ctypes
from textwrap import wrap
import time
import csv
from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_OAEP as Cipher
from  SABER_KEM.test.so_handler import main as sec_pub


# Konfigurasi MQTT server
broker = "192.168.1.17"
client = paho.Client()
port = 1883
def u_encrypt(arg):
    sce_pub(arg)

def u_encrypt(arg):
    so_file = 'EES401/URG_encrypt.so'
    u_enc = CDLL(so_file)
    u_enc.main.restype = c_double
    u_enc.main.argtype = c_char_p
    str_temp = "07"
    c_return = u_enc.main(arg.encode('utf-8'), str_temp.encode('utf-8'))
    print("Python Log, C Return         : ", c_return)
    # time.sleep(0.01)
    _ctypes.dlclose(u_enc._handle)

    return c_return


def rsa_encrypt(plain):

    rsa3072 = open("RSA7680_pub.pem", "r").read()
    start=time.perf_counter()
    rsa_pub = RSA.importKey(rsa3072)
    rsa_pub = Cipher.new(rsa_pub)
    RSA_enc = rsa_pub.encrypt(plain.encode('utf8'))
    # RSA_enc = rsa_pub.encrypt(plain)
    #RSA_enc = rsa_pub.encrypt(plain.encode('utf8'), 32)
    end=time.perf_counter()-start
    cipher = RSA_enc
    # cipher = str(cipher).encode('utf-8')
    # rsa_file_handler(base_name_RSA3, i, cipher)
    print("RSA Encryption Time : ", end)
    print("Cipher ", cipher)
    print("Length ", len(cipher))
    print(" ")

    return end, cipher


def file_handler(f_name, num):
    # num = str(num).zfill(2)
    # f = open("./cipher/"+f_name+"_"+num+".dat", "rb")
    f = open("/tmp/cipher_07.dat", "rb")
    ret = f.read()
    f.close()
    return ret

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

#Proses enkripsi
def payload_process(plain):
    dec_time = u_encrypt(plain)
    # cipher = u_encrypt(plain)
    # time.sleep(0.05)
    cipher = file_handler("cipher_EES401", 7)
    # print("Cipher Length                : ", len(plain))
    send_data = base64.b64encode(cipher)
    return send_data, dec_time

def payload_RSA(plain):
    dec_time, cipher = rsa_encrypt(plain)
    # cipher = u_encrypt(plain)
    # time.sleep(0.05)
    # cipher = rsa_file_handler("cipher_EES401", 7)
    # print("Cipher Length                : ", len(plain))
    send_data = base64.b64encode(cipher)
    return send_data, dec_time

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


# Add this to test the Blacklist feature
# blacklist_trigger = base64.b64encode(b"I'm gonna make this program crash!"*10)
# cipher_arr.append(blacklist_trigger)
# msg_data = blacklist_trigger

# ======================== || START FROM HERE BUDDY || ==================
def main(msg):
    # plain = input("Enter The Message : ")
    plain = msg
    plain_ori = plain
    # plain = plain_handler("plain", i)
    i = 0
    my_ip = "device01"
    identifier = "&"
    f_name = "cipher_EES401"
    plain_arr = []
    cipher_arr = []
    # plain = input("Enter plain  : ")
    # plain = "This message is longer than 4 char"
    x = 16 - (len(plain)%16)
    plain = plain+(" "*x)
    plain_arr = wrap(plain, 16, replace_whitespace=False, drop_whitespace=False)
    # print("Join List Test ", ''.join(plain_arr))
    # client.tls_set()
    # client.username_pw_set(username="aaa", password="pass")
    status = client.connect(broker,port)
    print("Connection status, ", status)
    # input()
    # plain_now = plain_arr[1]
    time_tmp = [] 
    for i in range(len(plain_arr)):
        # print(" ")
        # print("===== [ Sequence  "+str(i)+" ] "+"======================================================")
        data_now, dec_time = payload_process(plain_arr[i])
        # data_now, dec_time = payload_RSA(plain_arr[i])
        cipher_arr.append(data_now)
        time_tmp.append(dec_time)

    dec_acc = sum(map(float, time_tmp)) 
    xxx = (base64.b64decode(cipher_arr[0]))
    print(len(xxx))
    print(len(data_now))
    msg_data = b''.join(cipher_arr)
    # send_data = my_ip+identifier+msg_data.decode('ascii')
    send_data = msg_data.decode('ascii')
    # print("Python LOG Plain : ", plain)
    # print(" ")
    # print("===== Result ======================================================")
    # print("Python LOG Plain Length              : ", len(plain_ori))
    # print("Python LOG Plain with Padding Length : ", len(plain))
    # print("Sending Cipher with length           : ", len(send_data))
    # print("Python LOG Plain Array               : ", plain_arr)
    # print("Python LOG Cipher Length : ", len(send_data))
    # client.publish("device01/msg", "XXXXX")
    client.loop_start()
    print(len(send_data));
    # client.connect("nnag.xyz", 1883)
    # pahopub.single("device01/msg", "xxxx")
    # client.publish("device01/msg", msg)
    client.publish("device01/msg", msg_data)

    return dec_acc


# #main("AAA")
# test_str = "Test String"
# client.publish("device01/msg", base64.b64encode(test_str.encode("utf-8")) )