# The only difference between these two files is this one does _not_ generate a new key for every run.
# The other file (speed-test.py) does generate a new key for every run and it makes the RSA encryption test take _much_ longer.

from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_OAEP, AES
from Crypto.Util.Padding import pad, unpad
from Crypto.Random import get_random_bytes

import time

plaintext = input("Please input the test message: ").encode()
IV = b'0' * 16

def encrypt_aes(key: bytes, plaintext: bytes) -> bytes:
    plaintext = pad(plaintext, 16)

    cipher = AES.new(key, AES.MODE_CBC, iv=IV)

    return cipher.encrypt(plaintext)

def decrypt_aes(key: bytes, ciphertext: bytes) -> bytes:

    cipher = AES.new(key, AES.MODE_CBC, iv=IV)

    return unpad(cipher.decrypt(ciphertext), 16)



for size in (128, 192, 256):

    e_times = []
    d_times = []

    for n in range(100):
        aes_key = get_random_bytes(size // 8)

        start = time.time()
        ciphertext = encrypt_aes(aes_key, plaintext)
        end = time.time()

        e_times.append(end - start)

        start = time.time()
        decrypt_aes(aes_key, ciphertext)
        end = time.time()

        d_times.append(end - start)

    e_avg = sum(e_times) / len(e_times)
    d_avg = sum(d_times) / len(d_times)

    print(f'AES {size} key encrpytion average: {e_avg * 1000} ms')
    print(f'AES {size} key decryption average: {d_avg * 1000} ms\n')

def encrypt_rsa(public_key: bytes, plaintext: bytes) -> bytes:
    cipher = PKCS1_OAEP.new(public_key)
    return cipher.encrypt(plaintext)

def decrypt_rsa(private_key: bytes, ciphertext: bytes) -> bytes:
    cipher = PKCS1_OAEP.new(private_key)
    return cipher.decrypt(ciphertext)

for size in (1024, 2048, 4096):

    e_times = []
    d_times = []

    key = RSA.generate(size)
    private_key = key
    public_key = key.publickey()

    for n in range(100):

        start = time.time()
        ciphertext = encrypt_rsa(public_key, plaintext)
        end = time.time()
        e_times.append(end - start)

        start = time.time()
        decrypt_rsa(private_key, ciphertext)
        end = time.time()
        d_times.append(end - start)

    e_avg = sum(e_times) / len(e_times)
    d_avg = sum(d_times) / len(d_times)

    print(f'RSA {size} key encrpytion average: {e_avg * 1000} ms')
    print(f'RSA {size} key decryption average: {d_avg * 1000} ms\n')