Update app.py

app.py

@@ -1,10 +1,8 @@
+import gradio as gr
 import os
-import tkinter as tk
-from tkinter import filedialog
 from PIL import Image
 from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
 from cryptography.hazmat.backends import default_backend
-import gradio as gr
 
 class CBCEncryption:
     def __init__(self, key, iv):
@@ -24,28 +22,28 @@ class CBCEncryption:
     def finalize_decrypt(self):
         return self.decryptor.finalize()
 
-def EncryptImage(encryption, image_path, output_path):
-    image = Image.open(image_path)
+def EncryptImage(encryption, image, key):
+    image = Image.open(image)
     image.save('temp.bmp')
     with open('temp.bmp', 'rb') as reader:
-        with open(output_path, 'wb') as writer:
+        with open('encrypted.bmp', 'wb') as writer:
             image_data = reader.read()
             header, body = image_data[:54], image_data[54:]
             body += b'\x35' * (16 - (len(body) % 16))
             body = encryption.encrypt(body) + encryption.finalize_encrypt()
             writer.write(header + body)
     os.remove('temp.bmp')
+    return 'encrypted.bmp'
 
-def DecryptImage(decryption, image_path, output_path):
-    with open(image_path, 'rb') as reader:
-        with open(output_path, 'wb') as writer:
+def DecryptImage(decryption, image, key):
+    with open(image, 'rb') as reader:
+        with open('decrypted.bmp', 'wb') as writer:
             image_data = reader.read()
             header, body = image_data[:54], image_data[54:]
             body = decryption.decrypt(body)
             unpadded_body = unpad_data(body)
             writer.write(header + unpadded_body)
-    decrypted_image = Image.open(output_path)
-    decrypted_image.show()
+    return 'decrypted.bmp'
 
 def unpad_data(data):
     padding_byte = data[-1]
@@ -54,14 +52,49 @@ def unpad_data(data):
     return unpadded_data
 
 def main():
-    app = gr.Interface(
-        fn=lambda key, image_path: EncryptImage(CBCEncryption(key.encode('utf-8'), key.encode('utf-8')[:16]), image_path, filedialog.asksaveasfilename()),
-        inputs=["text", "file"],
-        outputs="file",
-        title="AES CBC Encryption"
+    def action_encrypt(image, key):
+        if len(key) > 32:
+            return 'Key is too long. Maximum key length is 32 characters.'
+        key = key.encode('utf-8')
+        key = key.ljust(32, b'\x35')
+
+        iv = key[:16]
+        iv = bytearray(iv)
+        for i in range(16):
+            iv[i] = iv[i] ^ 0x35
+        iv = bytes(iv)
+
+        AesCbc = CBCEncryption(key, iv)
+        output_path = EncryptImage(AesCbc, image, key)
+        return output_path
+
+    def action_decrypt(image, key):
+        if len(key) > 32:
+            return 'Key is too long. Maximum key length is 32 characters.'
+        key = key.encode('utf-8')
+        key = key.ljust(32, b'\x35')
+
+        iv = key[:16]
+        iv = bytearray(iv)
+        for i in range(16):
+            iv[i] = iv[i] ^ 0x35
+        iv = bytes(iv)
+
+        AesCbc = CBCEncryption(key, iv)
+        output_path = DecryptImage(AesCbc, image, key)
+        return output_path
+
+    # Define the Gradio interface
+    iface = gr.Interface(
+        [action_encrypt, action_decrypt], # Functions to wrap
+        ["image", "text"], # Input types
+        ["image", "text"], # Output types
+        ["Encrypt", "Decrypt"], # Button labels
+        title="AES CBC Encryption/Decryption",
+        description="Upload an image and a key to encrypt or decrypt it."
     )
 
-    app.launch()
+    iface.launch()
 
 if __name__ == '__main__':
-    main()
+    main()