app.py

import gradio as gr
import os
import numpy as np
from PIL import Image
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
import cv2

class CBCEncryption:
    def __init__(self, key, iv):
        self.cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
        self.encryptor = self.cipher.encryptor()
        self.decryptor = self.cipher.decryptor()

    def encrypt(self, image):
        return self.encryptor.update(image)

    def decrypt(self, image):
        return self.decryptor.update(image)

    def finalize_encrypt(self):
        return self.encryptor.finalize()

    def finalize_decrypt(self):
        return self.decryptor.finalize()

def add_gaussian_noise(image, mean=0, var=0.01):
    noise = np.random.normal(mean, var**0.5, image.shape)
    noisy_image = image + noise
    return noisy_image

def remove_gaussian_noise(image, mean=0, var=0.01):
    # This is a simplified approach to remove noise.
    # In practice, you might need a more sophisticated method to estimate the noise statistics.
    return image - np.random.normal(mean, var**0.5, image.shape)

def EncryptImage(encryption, image, key):
    image = Image.open(image)
    image_np = np.array(image)
    noisy_image_np = add_gaussian_noise(image_np)
    noisy_image = Image.fromarray(noisy_image_np.astype(np.uint8))
    noisy_image.save('temp.bmp')
    with open('temp.bmp', 'rb') as reader:
        with open('encrypted.bmp', 'wb') as writer:
            image_data = reader.read()
            header, body = image_data[:54], image_data[54:]
            body += b'\x35' * (16 - (len(body) % 16))
            body = encryption.encrypt(body) + encryption.finalize_encrypt()
            writer.write(header + body)
    os.remove('temp.bmp')
    return 'encrypted.bmp'

def DecryptImage(decryption, image, key):
    with open(image, 'rb') as reader:
        with open('decrypted.bmp', 'wb') as writer:
            image_data = reader.read()
            header, body = image_data[:54], image_data[54:]
            body = decryption.decrypt(body)
            unpadded_body = unpad_data(body)
            decrypted_image_np = np.frombuffer(unpadded_body, dtype=np.uint8).reshape(image.shape)
            original_image_np = remove_gaussian_noise(decrypted_image_np)
            original_image = Image.fromarray(original_image_np.astype(np.uint8))
            original_image.save('decrypted.bmp')
    return 'decrypted.bmp'

def unpad_data(data):
    padding_byte = data[-1]
    padding_length = padding_byte if padding_byte < 16 else 16
    unpadded_data = data[:-padding_length]
    return unpadded_data

def main():
    def action(image, key, action_type):
        if len(key) > 32:
            return 'Key is too long. Maximum key length is 32 characters.'
        key = key.encode('utf-8')
        key = key.ljust(32, b'\x35')

        iv = key[:16]
        iv = bytearray(iv)
        for i in range(16):
            iv[i] = iv[i] ^ 0x35
        iv = bytes(iv)

        AesCbc = CBCEncryption(key, iv)
        if action_type == 'Encrypt':
            output_path = EncryptImage(AesCbc, image, key)
        elif action_type == 'Decrypt':
            output_path = DecryptImage(AesCbc, image, key)
        else:
            return 'Invalid action type. Choose either Encrypt or Decrypt.'
        return output_path

    # Define the Gradio interface
    iface = gr.Interface(
        action, # Single function to wrap
        ["image", "text", "text"], # Input types: image, key, action_type
        "image", # Output type
        title="AES CBC Encryption/Decryption with Gaussian Noise",
        description="Upload an image, a key, and choose Encrypt or Decrypt to process the image."
    )

    iface.launch()

if __name__ == '__main__':
    main()