import streamlit as st
from cryptography.hazmat.primitives import padding
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
import base64
import os

# Generate a random 256-bit AES key (you can use different key sizes if needed)
key = b'Sixteen byte key'

def aes_encrypt(plaintext):
    backend = default_backend()

    # Generate a random IV (Initialization Vector)
    iv = os.urandom(16)

    cipher = Cipher(algorithms.AES(key), modes.CFB(iv), backend=backend)
    encryptor = cipher.encryptor()

    padder = padding.PKCS7(128).padder()
    padded_data = padder.update(plaintext) + padder.finalize()

    ciphertext = encryptor.update(padded_data) + encryptor.finalize()
    return base64.b64encode(iv + ciphertext)

def aes_decrypt(ciphertext):
    backend = default_backend()

    decoded_ciphertext = base64.b64decode(ciphertext)
    iv = decoded_ciphertext[:16]
    ciphertext = decoded_ciphertext[16:]

    cipher = Cipher(algorithms.AES(key), modes.CFB(iv), backend=backend)
    decryptor = cipher.decryptor()

    decrypted_data = decryptor.update(ciphertext) + decryptor.finalize()

    unpadder = padding.PKCS7(128).unpadder()
    unpadded_data = unpadder.update(decrypted_data) + unpadder.finalize()

    return unpadded_data

def main():
    st.title("AES Encryption and Decryption")

    operation = st.radio("Select operation:", ("Encrypt", "Decrypt"))

    if operation == "Encrypt":
        plaintext = st.text_area("Enter plaintext:", "")
        if st.button("Encrypt"):
            encrypted_data = aes_encrypt(plaintext.encode('utf-8'))
            st.success("Encryption successful!")
            st.text("Encrypted data:")
            st.write(encrypted_data.decode('utf-8'))
    else:
        ciphertext = st.text_area("Enter ciphertext:", "")
        if st.button("Decrypt"):
            decrypted_data = aes_decrypt(ciphertext.encode('utf-8'))
            st.success("Decryption successful!")
            st.text("Decrypted data:")
            st.write(decrypted_data.decode('utf-8'))

if __name__ == "__main__":
    main()