Update app.py

app.py

@@ -28,40 +28,60 @@ def create_key(passw):
     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)
+def encrypt(passw=None,mes=None,img=None,doc=None):
     #key = Fernet.generate_key()
     key = create_key(passw)
     fernet = Fernet(key)
-    encMessage = fernet.encrypt(bytes_e)
-  
-    return encMessage
+    
+    if mes != None and mes != "":
+        bytes_m = mes.encode()
+        print (bytes_m)
+        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()) 
+        print (bytes_i)
+        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+'
+    
+    return enc_mes
 
 def decrypt(passw=None,enc_in=None):
-    enc_in=enc_in.strip('"').strip("b'").strip("'")
-    enc_in = bytes(enc_in,'utf-8')
-
     key = create_key(passw)
-    
     fernet = Fernet(key)
-    decMessage = fernet.decrypt(enc_in).decode()
-    print("decrypted string: ", decMessage)
-    try:        
-        base = bytes(decMessage, 'utf-8')
+
+    enc_in=enc_in.strip('"').strip("b'").strip("'")
+    if "+aaa+" in enc_in:
+        mes1=enc_in.split("+aaa+",1)[1]
+        if "+bbb+" in mes1:
+            mes1=mes1.split("+bbb+",1)[0]
+            mes2=mes1.split("+bbb+",1)[1]
+            im_bytes = bytes(mes2,'utf-8')
+            #base = bytes(decMessage, 'utf-8')
+            with open(f"finished_im.png", "wb") as fh:
+                fh.write(base64.decodebytes(bytes(im_bytes, 'utf-8')))
+            fh.close     
+            dec_im = "finished_im.png"        
+        else:
+            dec_im = None
+        mes_bytes = bytes(mes1,'utf-8')
+        mes_dec = fernet.decrypt(mes_bytes).decode()
+   
+    if not "+aaa+" in enc_in and "+bbb+" in enc_in:
+        mes_dec=None
+        mes2 = enc_in.split("+bbb+",1)[1]
+        im_bytes = bytes(mes2,'utf-8')
+        #base = bytes(decMessage, 'utf-8')
         with open(f"finished_im.png", "wb") as fh:
-            fh.write(base64.decodebytes(bytes(decMessage, 'utf-8')))
+            fh.write(base64.decodebytes(bytes(im_bytes, 'utf-8')))
         fh.close     
         dec_im = "finished_im.png"
-    except Exception:
-        dec_im = None
-    return(dec_im,decMessage)
+
+    return(dec_im,mes_dec)