import gradio as gr
from Crypto.Cipher import AES
from Crypto.Random import get_random_bytes
import base64
import io

def pad(data):
    pad_len = AES.block_size - len(data) % AES.block_size
    return data + bytes([pad_len] * pad_len)

def unpad(data):
    return data[:-data[-1]]

def encrypt_image(image, password):
    key = password.encode('utf-8').ljust(32, b'\0')[:32]
    img_bytes = image.read()
    img_padded = pad(img_bytes)
    
    cipher = AES.new(key, AES.MODE_CBC)
    ciphertext = cipher.encrypt(img_padded)
    
    encrypted_data = base64.b64encode(cipher.iv + ciphertext).decode('utf-8')
    return encrypted_data

def decrypt_image(encrypted_text, password):
    try:
        key = password.encode('utf-8').ljust(32, b'\0')[:32]
        encrypted_data = base64.b64decode(encrypted_text)
        iv = encrypted_data[:16]
        ciphertext = encrypted_data[16:]
        
        cipher = AES.new(key, AES.MODE_CBC, iv=iv)
        decrypted = unpad(cipher.decrypt(ciphertext))
        
        return decrypted
    except Exception as e:
        return f"Decryption failed: {str(e)}"

def encrypt_ui(image, password):
    if image is None or password.strip() == "":
        return "Please upload image and enter a password."
    return encrypt_image(image, password)

def decrypt_ui(encrypted_text, password):
    result = decrypt_image(encrypted_text, password)
    if isinstance(result, bytes):
        return gr.Image.update(value=io.BytesIO(result))
    return result

footer_note = "🔒 Made by Kratu Pandey, B.Tech CSE"

encrypt_tab = gr.Interface(
    fn=encrypt_ui,
    inputs=[
        gr.File(label="Upload Image"),
        gr.Text(label="Password", type="password")
    ],
    outputs="text",
    title="🔐 Encrypt Image",
    description=footer_note
)

decrypt_tab = gr.Interface(
    fn=decrypt_ui,
    inputs=[
        gr.Text(label="Encrypted Text"),
        gr.Text(label="Password", type="password")
    ],
    outputs="image",
    title="🔓 Decrypt Image",
    description=footer_note
)

gr.TabbedInterface(
    [encrypt_tab, decrypt_tab],
    tab_names=["Encrypt Image", "Decrypt Image"]
).launch()