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
import qr
import stegan

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,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:
            im_bytes = fernet.encrypt(bytes_i)
            enc_mes = f'{im_bytes}+bbb+'
            
    #enc_qr = qr.make_qr(txt=(enc_mes.strip('"')))
    qr_link = "https://omnibus-fernet-encryption-qr.hf.space"
    enc_qr = stegan.conv_im(qr_link=qr_link,data=enc_mes)
    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,enc_qr
    
def decrypt(passw,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,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

def decode_qr(im,passw):
    with open(f'{im}', "rb") as image_file:
        bytes_i = base64.b64encode(image_file.read())     
    decode_qr = stegan.decode(im)
    dec_im = decrypt(passw, enc_in=decode_qr)
    '''
    print(f'decoded QR::{decode_qr}')
    key = create_key(passw)
    fernet = Fernet(key)
    dec_im = None
    mes_dec= None
    enc_in=decode_qr.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)
    
    
with gr.Blocks() as app:
    with gr.Tab("Encrypt"):
        pass_in=gr.Textbox(label="Set Password")
        with gr.Tab("String"):
            mes = gr.Textbox(label = "Message")
            with gr.Row():
                im = gr.Image(type="filepath")
        with gr.Tab("File"):
            doc=gr.File()
            enc_doc_out=gr.File()
        en_btn = gr.Button("Encrypt")
        enc_out = gr.Textbox(label="Encrypted Bytes")
        enc_qr_out = gr.Image(label = "Encrypted QR")

    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()
        with gr.Tab("QR"):
            dec_qr_im = gr.Image(type="filepath")
            dec_qr_txt = gr.Textbox(label="Decrypted")
            dec_qr_im_out =gr.Image(label="Decrytped Image")
            dec_qr_btn = gr.Button("Decrypt")

    dec_qr_btn.click(decode_qr,[dec_qr_im,pass_out],[dec_qr_im_out,dec_qr_txt])
    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,enc_qr_out])
    d_btn.click(decrypt,[pass_out,enc_in],[d_im,d_txt])
app.launch()