import streamlit as st
import base64
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
from Crypto.Random import get_random_bytes
import hashlib

st.title("🔒 Text Encryptor/Decryptor")
st.markdown("Secure text encryption using multiple algorithms")

# Encryption functions
def caesar_encrypt(text, shift):
    result = ""
    for char in text:
        if char.isalpha():
            ascii_offset = 65 if char.isupper() else 97
            result += chr((ord(char) + shift - ascii_offset) % 26 + ascii_offset)
        else:
            result += char
    return result

def caesar_decrypt(text, shift):
    return caesar_encrypt(text, -shift)

def vigenere_encrypt(text, key):
    key = key.upper()
    key_len = len(key)
    encrypted = []
    for i, char in enumerate(text):
        if char.isalpha():
            offset = 65 if char.isupper() else 97
            key_char = key[i % key_len]
            shift = ord(key_char) - 65
            encrypted_char = chr((ord(char) + shift - offset) % 26 + offset)
            encrypted.append(encrypted_char)
        else:
            encrypted.append(char)
    return "".join(encrypted)

def vigenere_decrypt(text, key):
    key = key.upper()
    key_len = len(key)
    decrypted = []
    for i, char in enumerate(text):
        if char.isalpha():
            offset = 65 if char.isupper() else 97
            key_char = key[i % key_len]
            shift = ord(key_char) - 65
            decrypted_char = chr((ord(char) - shift - offset) % 26 + offset)
            decrypted.append(decrypted_char)
        else:
            decrypted.append(char)
    return "".join(decrypted)

def aes_encrypt(text, password):
    salt = get_random_bytes(16)
    key = hashlib.scrypt(password.encode(), salt=salt, n=16384, r=8, p=1, dklen=32)
    iv = get_random_bytes(16)
    cipher = AES.new(key, AES.MODE_CBC, iv)
    ciphertext = cipher.encrypt(pad(text.encode(), AES.block_size))
    return base64.b64encode(salt + iv + ciphertext).decode()

def aes_decrypt(encrypted, password):
    encrypted = base64.b64decode(encrypted)
    salt = encrypted[:16]
    iv = encrypted[16:32]
    ciphertext = encrypted[32:]
    key = hashlib.scrypt(password.encode(), salt=salt, n=16384, r=8, p=1, dklen=32)
    cipher = AES.new(key, AES.MODE_CBC, iv)
    decrypted = unpad(cipher.decrypt(ciphertext), AES.block_size)
    return decrypted.decode()

# Interface
operation = st.radio("Select Operation:", ["Encrypt", "Decrypt"])
method = st.selectbox("Select Encryption Method:", 
                     ["Caesar Cipher", "Vigenère Cipher", "AES-256"])

text = st.text_area("Enter Text:", height=200)
password = st.text_input("Enter Key/Password:", type="password")

if st.button(f"Perform {operation}"):
    if text and password:
        try:
            if method == "Caesar Cipher":
                if not password.isdigit():
                    raise ValueError("Caesar cipher requires numeric key")
                shift = int(password) % 26
                if operation == "Encrypt":
                    result = caesar_encrypt(text, shift)
                else:
                    result = caesar_decrypt(text, shift)
                    
            elif method == "Vigenère Cipher":
                if operation == "Encrypt":
                    result = vigenere_encrypt(text, password)
                else:
                    result = vigenere_decrypt(text, password)
                    
            elif method == "AES-256":
                if operation == "Encrypt":
                    result = aes_encrypt(text, password)
                else:
                    result = aes_decrypt(text, password)
                    
            st.success(f"Result:")
            st.code(result)
            
            st.download_button(
                label="Download Result",
                data=result,
                file_name="encrypted.txt" if operation == "Encrypt" else "decrypted.txt",
                mime="text/plain"
            )
            
        except Exception as e:
            st.error(f"Error: {str(e)}")
    else:
        st.warning("Please enter both text and key/password")

st.sidebar.markdown("## Instructions")
st.sidebar.markdown("""
1. Select encrypt/decrypt operation
2. Choose encryption method
3. Enter text and key/password
4. Click the action button
                    
**Key Requirements:**
- Caesar: Numeric key (1-25)
- Vigenère: Text key (letters only)
- AES-256: Strong password (any characters)
""")