import gradio as gr
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import serialization
import base64
import os

# Caesar Cipher functions
def caesar_cipher_encrypt(text, shift):
    result = ""
    for i in range(len(text)):
        char = text[i]
        if char.isupper():
            result += chr((ord(char) + shift - 65) % 26 + 65)
        else:
            result += chr((ord(char) + shift - 97) % 26 + 97)
    return result

def caesar_cipher_decrypt(text, shift):
    return caesar_cipher_encrypt(text, -shift)

# AES functions
def aes_encrypt(key, plaintext):
    salt = os.urandom(16)
    kdf = PBKDF2HMAC(algorithm=hashes.SHA256(), length=32, salt=salt, iterations=100000)
    key = kdf.derive(key.encode())
    iv = os.urandom(16)
    cipher = Cipher(algorithms.AES(key), modes.CFB(iv))
    encryptor = cipher.encryptor()
    ciphertext = encryptor.update(plaintext.encode()) + encryptor.finalize()
    return base64.b64encode(salt + iv + ciphertext).decode('utf-8')

def aes_decrypt(key, ciphertext):
    raw = base64.b64decode(ciphertext)
    salt, iv, ciphertext = raw[:16], raw[16:32], raw[32:]
    kdf = PBKDF2HMAC(algorithm=hashes.SHA256(), length=32, salt=salt, iterations=100000)
    key = kdf.derive(key.encode())
    cipher = Cipher(algorithms.AES(key), modes.CFB(iv))
    decryptor = cipher.decryptor()
    return decryptor.update(ciphertext) + decryptor.finalize()

# RSA functions
def rsa_generate_keys():
    private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    public_key = private_key.public_key()
    pem_private = private_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption())
    pem_public = public_key.public_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PublicFormat.SubjectPublicKeyInfo)
    return pem_private.decode('utf-8'), pem_public.decode('utf-8')

def rsa_encrypt(public_key_pem, plaintext):
    public_key = serialization.load_pem_public_key(public_key_pem.encode('utf-8'))
    ciphertext = public_key.encrypt(
        plaintext.encode(),
        padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA256()), algorithm=hashes.SHA256(), label=None)
    )
    return base64.b64encode(ciphertext).decode('utf-8')

def rsa_decrypt(private_key_pem, ciphertext):
    private_key = serialization.load_pem_private_key(private_key_pem.encode('utf-8'), password=None)
    decrypted_text = private_key.decrypt(
        base64.b64decode(ciphertext),
        padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA256()), algorithm=hashes.SHA256(), label=None)
    )
    return decrypted_text.decode('utf-8')

def caesar_encrypt_decrypt(text, shift, mode):
    if mode == "Encrypt":
        return caesar_cipher_encrypt(text, int(shift))
    else:
        return caesar_cipher_decrypt(text, int(shift))

def aes_encrypt_decrypt(key, text, mode):
    if mode == "Encrypt":
        return aes_encrypt(key, text)
    else:
        return aes_decrypt(key, text)

def rsa_encrypt_decrypt(key, text, mode, key_type):
    if mode == "Encrypt":
        return rsa_encrypt(key, text)
    else:
        return rsa_decrypt(key, text)

def rsa_keys():
    private_key, public_key = rsa_generate_keys()
    return private_key, public_key

# Gradio Blocks interface
with gr.Blocks() as demo:
    gr.Markdown(
        """
        ### Simple Gradio Demo of encryption and decryption process, using Caesar Chiper, Advanced Encryption Standard (AES), and Rivest–Shamir–Adleman (RSA) method.
        
        """
    )
    with gr.Tab("Caesar Cipher"):
        with gr.Row():
            caesar_text = gr.Textbox(label="Text")
            caesar_shift = gr.Number(label="Shift")
            caesar_mode = gr.Radio(["Encrypt", "Decrypt"], label="Mode")
        caesar_output = gr.Textbox(label="Output")
        gr.Button("Submit").click(caesar_encrypt_decrypt, [caesar_text, caesar_shift, caesar_mode], caesar_output)

    with gr.Tab("AES"):
        with gr.Row():
            aes_key = gr.Textbox(label="Key")
            aes_text = gr.Textbox(label="Text")
            aes_mode = gr.Radio(["Encrypt", "Decrypt"], label="Mode")
        aes_output = gr.Textbox(label="Output")
        gr.Button("Submit").click(aes_encrypt_decrypt, [aes_key, aes_text, aes_mode], aes_output)

    with gr.Tab("RSA"):
        with gr.Row():
            rsa_key = gr.Textbox(label="Key (Public/Private)")
            rsa_text = gr.Textbox(label="Text")
            rsa_mode = gr.Radio(["Encrypt", "Decrypt"], label="Mode")
            rsa_key_type = gr.Radio(["Public Key", "Private Key"], label="Key Type")
        rsa_output = gr.Textbox(label="Output")
        gr.Button("Submit").click(rsa_encrypt_decrypt, [rsa_key, rsa_text, rsa_mode, rsa_key_type], rsa_output)

    with gr.Tab("RSA Key Generation"):
        rsa_private_output = gr.Textbox(label="Private Key")
        rsa_public_output = gr.Textbox(label="Public Key")
        gr.Button("Generate Keys").click(rsa_keys, [], [rsa_private_output, rsa_public_output])

demo.launch()