NTRU_IoT/infidel-code/py_wrappers/Publish_rciot_AES.py

import paho.mqtt.client as paho
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
import paho.mqtt.publish as pahopub
#from simplecrypt import encrypt, decrypt
from cryptography.fernet import Fernet
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes
import base64
from ctypes import *
import _ctypes
from textwrap import wrap
import time
import csv
from Crypto.Cipher import AES
from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_OAEP as Cipher


# Konfigurasi MQTT server
broker = "nnag.xyz"
client = paho.Client()
port = 1883
aes_key = "1234123412341234"
IV = 16*"\x00"

def fernet_keygen():
    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 fernet_encrypt(plain):
    time_s = time.perf_counter()
    fernet_key = fernet_keygen()
    send_data = Fernet(fernet_key).encrypt(plain.encode())
    time_e = time.perf_counter()
    end = time_e - time_s

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

def aes_encrypt(plain):
    start = time.perf_counter()
    key = AES.new(aes_key.encode("utf-8"), AES.MODE_CBC, IV=IV.encode("utf-8"))
    # cipher = key.encrypt(plain.encode("utf-8"))
    # plain = "x"*16
    print(len(plain)) 
    x = 16 - (len(plain) % 16)
    plain = plain+" "*x
    print(x)
    # input()
    cipher = key.encrypt(plain.encode("utf-8"))
    end   = time.perf_counter()-start

    return end, cipher

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_aes(plain):

    dec_time, cipher = aes_encrypt(plain)
    send_data = base64.b64encode(cipher)
    return send_data, dec_time

def payload_fernet(plain):
    dec_time, cipher = fernet_encrypt(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 = [] 
    
    # ******************************8
    # 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])
    #     data_now, dec_time = payload_fernet(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("cipher len : ",len(xxx))
    # print("base64 cipher len : ",len(data_now))

    # msg_data = b''.join(cipher_arr)
    # ******************************8

    # send_data = my_ip+identifier+msg_data.decode('ascii')
    # msg_data, dec_time = payload_fernet(plain)
    msg_data, dec_time = payload_aes(plain)
    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("Total sent data ", len(send_data));
    # client.connect("nnag.xyz", 1883)
    # pahopub.single("device01/msg", "xxxx")
    client.publish("device01/msg", msg_data)
    # input()

    return dec_time

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