import secrets
import string
from Crypto.Cipher import AES
from Crypto.Random import get_random_bytes
from Crypto.Util.Padding import pad, unpad
from Crypto.Hash import SHA256
import base64
import gradio as gr

def generate_random_key(length=32):
    """Generate a random key with digits, uppercase, lowercase, and special characters."""
    characters = string.ascii_lowercase + string.ascii_uppercase + string.digits + string.punctuation
    key = (
        secrets.choice(string.ascii_lowercase) +
        secrets.choice(string.ascii_uppercase) +
        secrets.choice(string.digits) +
        secrets.choice(string.punctuation)
    )
    key += ''.join(secrets.choice(characters) for _ in range(length - 4))
    key_list = list(key)
    secrets.SystemRandom().shuffle(key_list)
    return ''.join(key_list)

def derive_key(custom_key, key_size):
    """Derive a key of specified size (128, 192, or 256 bits) from a custom string using SHA-256."""
    hasher = SHA256.new()
    hasher.update(custom_key.encode('utf-8'))
    full_hash = hasher.digest()
    # Convert key_size to integer before division
    return full_hash[:int(key_size) // 8]

def encrypt_string(plain_text, key, key_size):
    """Encrypt a string using AES in CBC mode with specified key size."""
    try:
        if not plain_text:
            return "Error: Please provide text to encrypt."
        if not key:
            return "Error: Please provide a key."
        derived_key = derive_key(key, key_size)
        plain_text_bytes = plain_text.encode('utf-8')
        iv = get_random_bytes(AES.block_size)
        cipher = AES.new(derived_key, AES.MODE_CBC, iv)
        padded_text = pad(plain_text_bytes, AES.block_size)
        cipher_text = cipher.encrypt(padded_text)
        return base64.b64encode(iv + cipher_text).decode('utf-8')
    except Exception as e:
        return f"Encryption error: {str(e)}"

def decrypt_string(encrypted_text, key, key_size):
    """Decrypt a string using AES in CBC mode with specified key size."""
    try:
        if not encrypted_text:
            return "Error: Please provide encrypted text."
        if not key:
            return "Error: Please provide a key."
        derived_key = derive_key(key, key_size)
        raw = base64.b64decode(encrypted_text)
        iv = raw[:AES.block_size]
        cipher_text = raw[AES.block_size:]
        cipher = AES.new(derived_key, AES.MODE_CBC, iv)
        padded_text = cipher.decrypt(cipher_text)
        return unpad(padded_text, AES.block_size).decode('utf-8')
    except Exception as e:
        return f"Decryption error: {str(e)}"

def generate_and_display_key():
    """Generate and return a random key."""
    return generate_random_key()

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# AES Encryption/Decryption App")
    gr.Markdown("Encrypt or decrypt text using AES (128, 192, or 256-bit keys).")

    with gr.Row():
        key_size = gr.Radio(
            choices=["128", "192", "256"],
            label="AES Key Size (bits)",
            value="256"
        )
    
    with gr.Row():
        custom_key = gr.Textbox(
            label="Custom Key",
            placeholder="Enter your key or generate one",
            lines=1
        )
        generate_key_btn = gr.Button("Generate Random Key")

    with gr.Row():
        input_text = gr.Textbox(
            label="Input Text",
            placeholder="Enter text to encrypt",
            lines=3
        )
    
    with gr.Row():
        encrypt_btn = gr.Button("Encrypt")
        decrypt_btn = gr.Button("Decrypt")
    
    output_text = gr.Textbox(
        label="Output",
        placeholder="Result will appear here",
        lines=3,
        interactive=False
    )

    # Event handlers
    generate_key_btn.click(
        fn=generate_and_display_key,
        outputs=custom_key
    )
    
    encrypt_btn.click(
        fn=encrypt_string,
        inputs=[input_text, custom_key, key_size],
        outputs=output_text
    )
    
    decrypt_btn.click(
        fn=decrypt_string,
        inputs=[input_text, custom_key, key_size],
        outputs=output_text
    )

# Launch the app
demo.launch()