from Crypto.Cipher import DES, AES
from Crypto.Util.Padding import pad, unpad
import time
from bokeh.plotting import figure, show
from bokeh.models import ColumnDataSource
from bokeh.io import output_file


def benchmark_encryption_decryption(cipher_func, data, key):
    # Measure encryption time
    start_time = time.time()
    encrypted = cipher_func(data, key)
    encrypt_time = time.time() - start_time

    # Measure decryption time
    start_time = time.time()
    decrypted = cipher_func(encrypted, key, decrypt=True)
    decrypt_time = time.time() - start_time

    return encrypt_time, decrypt_time


def des_cipher(data, key, decrypt=False):
    cipher = DES.new(key.encode('utf-8'), DES.MODE_ECB)
    if decrypt:
        return unpad(cipher.decrypt(data), DES.block_size).decode('utf-8')
    else:
        padded_data = pad(data.encode('utf-8'), DES.block_size)
        return cipher.encrypt(padded_data)


def aes_cipher(data, key, decrypt=False):
    cipher = AES.new(key.encode('utf-8'), AES.MODE_ECB)
    if decrypt:
        return unpad(cipher.decrypt(data), AES.block_size).decode('utf-8')
    else:
        padded_data = pad(data.encode('utf-8'), AES.block_size)
        return cipher.encrypt(padded_data)


def main():
    message = "Performance Testing of Encryption Algorithms"

    # Test with different key sizes
    des_key = "12345678"  # 8 bytes for DES
    aes_key = "12345678901234567890123456789012"  # 32 bytes for AES-256

    # Benchmark DES
    des_encrypt_time, des_decrypt_time = benchmark_encryption_decryption(des_cipher, message, des_key)

    # Benchmark AES-256
    aes_encrypt_time, aes_decrypt_time = benchmark_encryption_decryption(aes_cipher, message, aes_key)

    # Print results
    print("Performance Results:")
    print("-" * 40)
    print(f"DES Encryption: {des_encrypt_time:.9f} seconds")
    print(f"DES Decryption: {des_decrypt_time:.9f} seconds")
    print(f"AES-256 Encryption: {aes_encrypt_time:.9f} seconds")
    print(f"AES-256 Decryption: {aes_decrypt_time:.9f} seconds")

    # Create Bokeh visualization
    algorithms = ['DES', 'AES-256']
    encrypt_times = [des_encrypt_time, aes_encrypt_time]
    decrypt_times = [des_decrypt_time, aes_decrypt_time]

    # Create the plot
    p = figure(x_range=algorithms, height=400, title="Encryption/Decryption Performance Comparison",
               toolbar_location=None, responsive=True)

    p.vbar(x='algorithm', top='encrypt_time', width=0.4, color='blue', alpha=0.7, legend_label="Encryption",
           source=ColumnDataSource(data=dict(algorithm=algorithms, encrypt_time=encrypt_times)))
    p.vbar(x='algorithm', top='decrypt_time', width=0.4, color='red', alpha=0.7, legend_label="Decryption",
           source=ColumnDataSource(data=dict(algorithm=algorithms, decrypt_time=decrypt_times)))

    p.legend.location = "top_right"
    p.xaxis.axis_label = "Algorithm"
    p.yaxis.axis_label = "Time (seconds)"
    p.title.text_font_size = "14px"

    # Save and show the plot
    output_file("encryption_performance.html")
    show(p)


if __name__ == '__main__':
    main()