import gradio as gr
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
import os
import base64
import hashlib
from pathlib import Path


def create_key(passw):
    key = passw.encode()
    h = hashlib.new('sha256')
    h.update(key)
    key = h.hexdigest()
    key = key.encode()
    #salt = os.urandom(16)
    salt = key
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=480000,
    )
    key = base64.urlsafe_b64encode(kdf.derive(key))

    return key


def encrypt(passw=None,mes=None,img=None,doc=None):
    #key = Fernet.generate_key()
    key = create_key(passw)
    fernet = Fernet(key)
    enc_mes=None
    enc_file=None
    if mes != None and mes != "":
        bytes_m = mes.encode()
        enc_mes = fernet.encrypt(bytes_m)
        enc_mes = f'{enc_mes}+aaa+'
    if img != None:
        with open(f'{img}', "rb") as image_file:
            bytes_i = base64.b64encode(image_file.read()) 
        if mes != None and mes != "":
            im_bytes = fernet.encrypt(bytes_i)
            enc_mes = f'{enc_mes}{im_bytes}+bbb+'
        else:
            enc_mes = f'{im_bytes}+bbb+'

    if doc != None:
        
        print(dir(doc))
        print(doc.name)
        #print(doc.file)
        doc_name = doc.name
        doc_name = doc_name.split("/",4)[4]
        #bytes_d = doc.encode()
        
        #doc = Path(doc)
        with open(doc.name, "rb") as file:
            # read all file data
            file_data = file.read()
        enc_doc = fernet.encrypt(file_data)
        og_name = doc.name
        og_end = og_name.split(".",1)[1]
        og_front=og_name.split(".",1)[0]
        
        
        enc_file=f'{doc.name}.ocrpt'
        with open(enc_file, "wb") as file:
            file.write(enc_doc)
        
    return enc_mes,enc_file
    
def decrypt(passw=None,enc_in=None):
    key = create_key(passw)
    fernet = Fernet(key)
    dec_im = None
    mes_dec= None
    enc_in=enc_in.strip('"')
    print (f'enc_in :::: {enc_in}')
    
    if "+aaa+" in enc_in:
        mes1=enc_in.split("+aaa+",1)[0]
        mes1=mes1.strip("b'").strip("'")
        mes_bytes = bytes(mes1,'utf-8')
        mes_dec = fernet.decrypt(mes_bytes).decode()
        if "+bbb+" in enc_in:
            mes12=enc_in.split("+aaa+",1)[1]
            mes2=mes12.split("+bbb+",1)[0]
            mes2=mes2.strip("b'").strip("'")
            im_bytes = bytes(mes2,'utf-8')
            print(f'im_bytes::{im_bytes}')
            
            mes2 = fernet.decrypt(mes2).decode()
            #base = bytes(decMessage, 'utf-8')
            with open(f"finished_im.png", "wb") as fh:
                #fh.write(base64.decodebytes(im_bytes))
                fh.write(base64.decodebytes(bytes(mes2, 'utf-8')))
            fh.close     
            dec_im = "finished_im.png"        
   
    if not "+aaa+" in enc_in:
        if "+bbb+" in enc_in:
            mes2 = enc_in.split("+bbb+",1)[0]
            mes2=mes2.strip("b'").strip("'")
            im_bytes = bytes(mes2,'utf-8')
            print(f'im_bytes2::{im_bytes}')
            mes2 = fernet.decrypt(mes2).decode()

            #base = bytes(decMessage, 'utf-8')
            with open(f"finished_im.png", "wb") as fh:
                #fh.write(base64.decodebytes(im_bytes))
                fh.write(base64.decodebytes(bytes(mes2, 'utf-8')))
            fh.close     
            dec_im = "finished_im.png"


    return(dec_im,mes_dec)
def decode_doc(passw=None,doc=None):
    key = create_key(passw)
    fernet = Fernet(key)
    doc_name = doc.name
    doc_name = doc_name.split("/",4)[4]
    #bytes_d = doc.encode()
    
    #doc = Path(doc)
    with open(doc.name, "rb") as file:
        # read all file data
        file_data = file.read()
    dec_doc = fernet.decrypt(file_data)
    og_name = doc.name
    og_end = og_name.split(".",1)[1]
    og_front=og_name.split(".",1)[0]
    
    dec_file = doc.name.strip(".ocrpt")
    #enc_file=f'{doc.name}.ocrpt'
    with open(dec_file, "wb") as file:
        file.write(dec_doc)    

    return dec_file
    
with gr.Blocks() as app:
    with gr.Tab("Encrypt"):
        pass_in=gr.Textbox(label="Set Password")
        mes = gr.Textbox(label = "Message")
        with gr.Row():
            im = gr.Image(type="filepath")
            doc=gr.File()
        en_btn = gr.Button("Encrypt")
        enc_out = gr.Textbox(label="Encrypted Bytes")
        enc_doc_out=gr.File()
    with gr.Tab("Decrypt"):
        pass_out = gr.Textbox(label="Enter Password")
        with gr.Tab("String"):
            enc_in = gr.Textbox(label="Encrypted Bytes")
            d_btn = gr.Button("Decrypt")
            d_txt = gr.Textbox(label="Decrypted")
            d_im =gr.Image(label="Decrytped Image")
        with gr.Tab("File"):
            
            dec_doc_in = gr.File()
            dec_doc_btn = gr.Button("Decrypt")
            dec_doc_out=gr.File()
    dec_doc_btn.click(decode_doc,[pass_out,dec_doc_in],dec_doc_out)
    en_btn.click(encrypt,[pass_in,mes,im,doc],[enc_out,enc_doc_out])
    d_btn.click(decrypt,[pass_out,enc_in],[d_im,d_txt])
app.launch()