import gradio as gr
from Crypto.PublicKey import RSA
from Crypto.Random import get_random_bytes
from Crypto.Cipher import AES, PKCS1_OAEP
from Crypto.Hash import RIPEMD160, SHA256
import base58

import stegan2
import qr


def calculate_hash(data, hash_function: str = "sha256") -> str:
    if type(data) == str:
        data = bytearray(data, "utf-8")
    if hash_function == "sha256":
        h = SHA256.new()
        h.update(data)
        return h.hexdigest()
    if hash_function == "ripemd160":
        h = RIPEMD160.new()
        h.update(data)
        return h.hexdigest()


def generate_keys():
    key = RSA.generate(2048)
    private_key = key.export_key('PEM')
    file_out = open("private.pem", "wb")
    file_out.write(private_key)
    file_out.close()
    
    public_key = key.publickey().export_key('PEM')
    file_out = open("receiver.pem", "wb")
    file_out.write(public_key)
    file_out.close()

    hash_1 = calculate_hash(public_key, hash_function="sha256")
    hash_2 = calculate_hash(hash_1, hash_function="ripemd160")
    address = base58.b58encode(hash_2)
    address_im=qr.make_qr(txt=address)
    
    return public_key,private_key,address_im,address

def encrypt_text(data,in2,address):
    data = data.encode("utf-8")
    #data = "I met aliens in UFO. Here is the map.".encode("utf-8")
    file_out = open("encrypted_data.bin", "wb")
    
    recipient_key = RSA.import_key(open("receiver.pem").read())
    session_key = get_random_bytes(16)
    
    # Encrypt the session key with the public RSA key
    cipher_rsa = PKCS1_OAEP.new(recipient_key)
    enc_session_key = cipher_rsa.encrypt(session_key)
    
    # Encrypt the data with the AES session key
    cipher_aes = AES.new(session_key, AES.MODE_EAX)
    ciphertext, tag = cipher_aes.encrypt_and_digest(data)
    [ file_out.write(x) for x in (enc_session_key, cipher_aes.nonce, tag, ciphertext) ]
    file_out.close()
    doc_name = "encrypted_data.bin"
    with open(doc_name, "rb") as file:
        file_data =(file.read())
        print (f'file_data::{file_data}')
    qr_link="test"
    enc_qr = stegan2.conv_im(qr_link=qr_link,data=file_data)
    
    file.close()
    return str(file_data),enc_qr

def decrypt_text(im,in2):
    enc_in = stegan2.decode(im)
    private_key = RSA.import_key(open("private.pem").read())
    enc_session_key = enc_in[:private_key.size_in_bytes()]
    end1 = private_key.size_in_bytes()+16
    nonce = enc_in[private_key.size_in_bytes():end1]
    start1=end1+1
    end2 = private_key.size_in_bytes()+32
    start2=end2+1
    tag = enc_in[end1:end2]
    ciphertext = enc_in[end2:]
    print (f'enc_session_key::{enc_session_key}')
    print (f'nonce::{nonce}')
    print (f'tag::{tag}')
    print (f'ciphertext::{ciphertext}')    

    # Decrypt the session key with the private RSA key
    cipher_rsa = PKCS1_OAEP.new(private_key)
    session_key = cipher_rsa.decrypt(enc_session_key)
    
    # Decrypt the data with the AES session key
    cipher_aes = AES.new(session_key, AES.MODE_EAX, nonce)
    data = cipher_aes.decrypt_and_verify(ciphertext, tag)
    return(data.decode("utf-8"))



with gr.Blocks() as app:

    with gr.Row():
        gr.Column()
        with gr.Column():
            with gr.Box():
                with gr.Row():
                    with gr.Column():
                        gen_wal_btn=gr.Button()
                        seed = gr.Textbox(label='Seed Phrase')
                        out1 = gr.Textbox(label='Private Key')
                        out2 = gr.Textbox(label='Public Key')
                        img3=gr.Pil()
                        out3 = gr.Textbox(label='Address')
            rsa_to_enc = gr.Textbox(label="txt to encrypt")
            rsa_enc_btn = gr.Button("RSA Encrypt")
            rsa_enc_mes = gr.Textbox(label="encoded")
            qr_enc_mes = gr.Image(type="filepath")
            rsa_dec_btn = gr.Button("RSA Decrypt")
            rsa_dec_mes = gr.Textbox(label="decoded")
        gr.Column()

    gen_wal_btn.click(generate_keys,None,[out2,out1, img3,out3])
    rsa_enc_btn.click(encrypt_text,[rsa_to_enc,out2,out3],[rsa_enc_mes,qr_enc_mes])
    rsa_dec_btn.click(decrypt_text,[qr_enc_mes,out1],rsa_dec_mes)
app.launch()