Omnibus/RSA-Encryption-dev

1import gradio as gr2from Crypto.PublicKey import RSA3from Crypto.Random import get_random_bytes4from Crypto.Cipher import AES, PKCS1_OAEP5from Crypto.Hash import RIPEMD160, SHA2566import base587import stegan8import stegan29import qr10 11 12def calculate_hash(data, hash_function: str = "sha256") -> str:13    if type(data) == str:14        data = bytearray(data, "utf-8")15    if hash_function == "sha256":16        h = SHA256.new()17        h.update(data)18        return h.hexdigest()19    if hash_function == "ripemd160":20        h = RIPEMD160.new()21        h.update(data)22        return h.hexdigest()23 24 25def generate_keys():26    key = RSA.generate(2048)27    private_key = key.export_key('PEM')28    print(f'private_key:: {private_key}')29    30    file_out_priv = open("private.pem", "wb")31    file_out_priv.write(private_key)32    file_out_priv.close()33    34    public_key = key.publickey().export_key('PEM')35    file_out_pub = open("receiver.pem", "wb")36    file_out_pub.write(public_key)37    file_out_pub.close()38 39    qr_link="test"40    priv_key = stegan.conv_im("private_key.png",data=private_key)41    pub_key = stegan.conv_im("public_key.png",data=public_key)42    hash_1 = calculate_hash(public_key, hash_function="sha256")43    hash_2 = calculate_hash(hash_1, hash_function="ripemd160")44    address = base58.b58encode(hash_2)45    address_im=qr.make_qr(txt=address)46    47    return public_key,private_key,address_im,address,priv_key,pub_key48 49def encrypt_text(data,in2,address):50    pub_key = stegan2.decode(in2)51    52    data = data.encode("utf-8")53    #data = "I met aliens in UFO. Here is the map.".encode("utf-8")54    55    #recipient_key = RSA.import_key(open("receiver.pem").read())56    recipient_key = RSA.import_key(pub_key)57 58    session_key = get_random_bytes(16)59    60    # Encrypt the session key with the public RSA key61    cipher_rsa = PKCS1_OAEP.new(recipient_key)62    enc_session_key = cipher_rsa.encrypt(session_key)63    64    # Encrypt the data with the AES session key65    cipher_aes = AES.new(session_key, AES.MODE_EAX)66    ciphertext, tag = cipher_aes.encrypt_and_digest(data)67 68    file_out = open("encrypted_data.bin", "wb")69    70    [ file_out.write(x) for x in (enc_session_key, cipher_aes.nonce, tag, ciphertext) ]71    file_out.close()72    doc_name = "encrypted_data.bin"73    with open(doc_name, "rb") as file:74        file_data =(file.read())75        print (f'file_data::{file_data}')76    qr_link="test"77    enc_qr = stegan2.conv_im(qr_link=qr_link,data=file_data)78    79    file.close()80    return str(file_data),enc_qr81 82def decrypt_text(im,in2):83    enc_in = stegan2.decode(im)84    85    #private_key = RSA.import_key(open("private.pem").read())86    priv_key = stegan2.decode(in2)87    print(f'priv_key:: {priv_key}')88    private_key = RSA.import_key(priv_key)89    print(f'private_key:: {private_key}')90    enc_session_key = enc_in[:private_key.size_in_bytes()]91    end1 = private_key.size_in_bytes()+1692    nonce = enc_in[private_key.size_in_bytes():end1]93    start1=end1+194    end2 = private_key.size_in_bytes()+3295    start2=end2+196    tag = enc_in[end1:end2]97    ciphertext = enc_in[end2:]98    print (f'enc_session_key::{enc_session_key}')99    print (f'nonce::{nonce}')100    print (f'tag::{tag}')101    print (f'ciphertext::{ciphertext}')    102 103    # Decrypt the session key with the private RSA key104    cipher_rsa = PKCS1_OAEP.new(private_key)105    session_key = cipher_rsa.decrypt(enc_session_key)106    107    # Decrypt the data with the AES session key108    cipher_aes = AES.new(session_key, AES.MODE_EAX, nonce)109    data = cipher_aes.decrypt_and_verify(ciphertext, tag)110    return(data.decode("utf-8"))111 112 113 114with gr.Blocks() as app:115 116    with gr.Row():117        gr.Column()118        with gr.Column():119            with gr.Row():120                with gr.Tab("Gen Wal"):121                    gen_wal_btn=gr.Button()122                    seed = gr.Textbox(label='Seed Phrase')123                    img1=gr.Pil(label='Private Key')124                    out1 = gr.Textbox(label='Private Key',max_lines=4)125                    img2=gr.Pil(label='Public Key')126                    out2 = gr.Textbox(label='Public Key',max_lines=4)127                    img3=gr.Pil(label='Address')128                    out3 = gr.Textbox(label='Address')129                with gr.Tab("Encrypt"):130                    rsa_to_enc = gr.Textbox(label="txt to encrypt")131                    pub_key_in = gr.Image(label="Public Key", type="filepath")132                    rsa_enc_btn = gr.Button("RSA Encrypt")133                    rsa_enc_mes = gr.Textbox(label="encoded", max_lines=4)134                    qr_enc_mes = gr.Image(type="filepath")135                with gr.Tab("Decrypt"):136                    mes_in = gr.Image(label="Message", type="filepath")137                    priv_key_in = gr.Image(label="Private Key", type="filepath")138                    rsa_dec_btn = gr.Button("RSA Decrypt")139                    rsa_dec_mes = gr.Textbox(label="decoded")140        gr.Column()141 142    gen_wal_btn.click(generate_keys,None,[out2,out1, img3,out3,img1,img2])143    rsa_enc_btn.click(encrypt_text,[rsa_to_enc,pub_key_in,out3],[rsa_enc_mes,qr_enc_mes])144    rsa_dec_btn.click(decrypt_text,[mes_in,priv_key_in],rsa_dec_mes)145app.launch()