Update app.py

app.py

@@ -7,19 +7,10 @@ import os
 import base64
 import hashlib
 
-def encrypt(mes=None,img=None,document=None):
 
-    if mes != None and mes != "":
-        bytes_e = mes.encode()
-        print (bytes_e)
-    elif mes == None or mes=="" and img != None:
-        with open(f'{img}', "rb") as image_file:
-            bytes_e = base64.b64encode(image_file.read()) 
-        print (bytes_e)
-    #key = Fernet.generate_key()
-    key = "potato"
-  
-    key = key.encode()
+
+def create_key(passw):
+    key = passw.encode()
     h = hashlib.new('sha256')
     h.update(key)
     key = h.hexdigest()
@@ -27,9 +18,7 @@ def encrypt(mes=None,img=None,document=None):
     print (key)
     print (h)
     #key = h
-    
-    
-    password=key
+    #password=key
     #salt = os.urandom(16)
     salt = key
     kdf = PBKDF2HMAC(
@@ -38,8 +27,21 @@ def encrypt(mes=None,img=None,document=None):
         salt=salt,
         iterations=480000,
     )
-    key = base64.urlsafe_b64encode(kdf.derive(password))
-    
+    key = base64.urlsafe_b64encode(kdf.derive(key))
+    return key
+
+
+def encrypt(passw=None,mes=None,img=None,document=None):
+
+    if mes != None and mes != "":
+        bytes_e = mes.encode()
+        print (bytes_e)
+    elif mes == None or mes=="" and img != None:
+        with open(f'{img}', "rb") as image_file:
+            bytes_e = base64.b64encode(image_file.read()) 
+        print (bytes_e)
+    #key = Fernet.generate_key()
+    key = create_key(passw)  
     fernet = Fernet(key)
     
     # then use the Fernet class instance
@@ -49,24 +51,37 @@ def encrypt(mes=None,img=None,document=None):
     print(f'Length in Bytes: {len(encMessage)}')
     #print("original string: ", mes)
     print("encrypted string: ", encMessage)
-     
-    if mes != None and mes != "":
+    return encMessage
+def decrypt(passw=None,enc_in=None):
+    key = create_key(passw)
 
-        decMessage = fernet.decrypt(encMessage).decode()
-        print("decrypted string: ", decMessage)
-        
-    if mes == None or mes=="" and img != None:
-        decMessage = fernet.decrypt(encMessage).decode()
+    
+    decMessage = fernet.decrypt(enc_in).decode()
+    print("decrypted string: ", decMessage)
+    try:        
         base = bytes(decMessage, 'utf-8')
         with open(f"finished_im.png", "wb") as fh:
             fh.write(base64.decodebytes(bytes(decMessage, 'utf-8')))
         fh.close     
-        return("finished_im.png")
+        dec_im = "finished_im.png"
+    except Exception:
+        dec_im = None
+    return(dec_im,decMessage)
 with gr.Blocks() as app:
-    mes = gr.Textbox()
-    im = gr.Image(type="filepath")
-    d_im =gr.Image()
-    btn = gr.Button()
-    btn.click(encrypt,[mes,im],d_im)
+    with gr.Tab("Encrypt"):
+        pass_in=gr.Textbox(label="Set Password")
+        mes = gr.Textbox(label = "Message")
+        im = gr.Image(type="filepath")
+        en_btn = gr.Button("Encrypt")
+        enc_out = gr.Textbox("Encoded Bytes")
+    with gr.Tab("Decrypt"):
+        pass_out = gr.Textbox("Enter Password")
+        enc_in = gr.Textbox("Encrypted Text")
+        d_btn = gr.Button("Decrypt")
+        d_txt = gr.Textbox()
+        d_im =gr.Image()
+        d_txt=gr.Textbox()
+    en_btn.click(encrypt,[pass_in,mes,im],enc_out)
+    d_btn.click(decrypt,[pass_out,enc_in],[d_im,d_txt])
 app.launch()