Update app.py

app.py

@@ -52,7 +52,7 @@ def unpad_data(data):
     return unpadded_data
 
 def main():
-    def action_encrypt(image, key):
+    def action(image, key, action_type):
         if len(key) > 32:
             return 'Key is too long. Maximum key length is 32 characters.'
         key = key.encode('utf-8')
@@ -65,33 +65,21 @@ def main():
         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)
+        if action_type == 'Encrypt':
+            output_path = EncryptImage(AesCbc, image, key)
+        elif action_type == 'Decrypt':
+            output_path = DecryptImage(AesCbc, image, key)
+        else:
+            return 'Invalid action type. Choose either Encrypt or Decrypt.'
         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
+        action, # Single function to wrap
+        ["image", "text", "text"], # Input types: image, key, action_type
+        "image", # Output type
         title="AES CBC Encryption/Decryption",
-        description="Upload an image and a key to encrypt or decrypt it."
+        description="Upload an image, a key, and choose Encrypt or Decrypt to process the image."
     )
 
     iface.launch()