Update app.py

app.py

@@ -24,76 +24,17 @@ def decrypt_image(decryption, image_path, output_path):
 def unpad_data(data):
     # Your unpad_data function code here...
 
-def main():
-    # Your main function code here...
-
-import subprocess
-subprocess.run(["pip", "install", "altair==4.1.0"])
-import os
-import streamlit as st
-from PIL import Image
-from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
-from cryptography.hazmat.backends import default_backend
-
-class CBCEncryption:
-    def __init__(self, key, iv):
-        self.cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
-        self.encryptor = self.cipher.encryptor()
-        self.decryptor = self.cipher.decryptor()
-
-    def encrypt(self, image):
-        return self.encryptor.update(image)
-
-    def decrypt(self, image):
-        return self.decryptor.update(image)
-
-    def finalize_encrypt(self):
-        return self.encryptor.finalize()
-
-    def finalize_decrypt(self):
-        return self.decryptor.finalize()
-
-def encrypt_image(encryption, image_path, output_path):
-    image = Image.open(image_path)
-    image.save('temp.bmp')
-    with open('temp.bmp', 'rb') as reader:
-        with open(output_path, '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')
-
-def decrypt_image(decryption, image_path, output_path):
-    with open(image_path, 'rb') as reader:
-        with open(output_path, '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)
-    st.image(decrypted_image, caption="Decrypted Image")
-
-def unpad_data(data):
-    padding_byte = data[-1]
-    padding_length = padding_byte if padding_byte < 16 else 16
-    unpadded_data = data[:-padding_length]
-    return unpadded_data
-
 def main():
     st.title("AES CBC Encryption/Decryption")
 
     uploaded_file = st.file_uploader("Choose an image...", type=["bmp", "png", "jpg", "jpeg"])
     key = st.text_input("Enter Key:", type="password")
-    action = st.radio("Choose Action:", ("Encrypt", "Decrypt"))
 
     if uploaded_file is not None and key:
-        if len(key) > 32:
-            st.error("Key is too long. Maximum key length is 32 characters.")
-            return
-        
+        st.write("")  # Add space between encryption and decryption options
+        encrypt_button = st.button("Encrypt Image")
+        decrypt_button = st.button("Decrypt Image")
+
         key = key.encode('utf-8')
         key = key.ljust(32, b'\x35')
         iv = key[:16]
@@ -110,16 +51,17 @@ def main():
         with open(image_path, 'wb') as f:
             f.write(uploaded_file.getbuffer())
         
-        if action == "Encrypt":
+        if encrypt_button:
             encrypt_image(encryption=AesCbc, image_path=image_path, output_path=output_path)
             st.success("Encryption done!")
             with open(output_path, 'rb') as f:
                 st.download_button(label="Download Encrypted Image", data=f, file_name="encrypted_image.bmp")
         
-        elif action == "Decrypt":
+        elif decrypt_button:
             decrypt_image(decryption=AesCbc, image_path=image_path, output_path=output_path)
             st.success("Decryption done!")
+            decrypted_image = Image.open(output_path)
+            st.image(decrypted_image, caption="Decrypted Image")
 
 if __name__ == '__main__':
     main()
-