test.py

import os, sys
from aes import *

plain_text_1 = "Lorem ipsum dolor sit amet, consectetur tincidunt.".encode()
plain_text_2 = "Neque porro quisquam est qui dolorem ipsum quia dolor sit amet, consectetur, adipisci velit.".encode()
plain_text_3 = "Lorem ipsum dolor sit amet, consectetur adipiscing elit. In dui.".encode()

iv = os.urandom(16)
key_128 = os.urandom(16)
key_192 = os.urandom(24)
key_256 = os.urandom(32)

object_1 = AES(key_128)
object_2 = AES(key_192)
object_3 = AES(key_256)

def test_ecb():
  print("#####################")
  print("# Tests in ECB mode #")
  print("#####################\n")

  print("---> Plain text:", plain_text_1, "\n")
  temp = object_1.encrypt_ecb(plain_text_1)
  print("-Encrypted text (128-bit key):", temp)
  temp = object_1.decrypt_ecb(temp)
  print("-Decrypted text (128-bit key):", temp, "\n")
  temp = object_2.encrypt_ecb(plain_text_1)
  print("-Encrypted text (192-bit key):", temp)
  temp = object_2.decrypt_ecb(temp)
  print("-Decrypted text (192-bit key):", temp, "\n")
  temp = object_3.encrypt_ecb(plain_text_1)
  print("-Encrypted text (256-bit key):", temp)
  temp = object_3.decrypt_ecb(temp)
  print("-Decrypted text (256-bit key):", temp, "\n")

  print("-------------------------\n")

  print("---> Plain text:", plain_text_2, "\n")
  temp = object_1.encrypt_ecb(plain_text_2)
  print("-Encrypted text (128-bit key):", temp)
  temp = object_1.decrypt_ecb(temp)
  print("-Decrypted text (128-bit key):", temp, "\n")
  temp = object_2.encrypt_ecb(plain_text_2)
  print("-Encrypted text (192-bit key):", temp)
  temp = object_2.decrypt_ecb(temp)
  print("-Decrypted text (192-bit key):", temp, "\n")
  temp = object_3.encrypt_ecb(plain_text_2)
  print("-Encrypted text (256-bit key):", temp)
  temp = object_3.decrypt_ecb(temp)
  print("-Decrypted text (256-bit key):", temp, "\n")

  print("-------------------------\n")

  print("---> Plain text:", plain_text_3, "\n")
  temp = object_1.encrypt_ecb(plain_text_3)
  print("-Encrypted text (128-bit key):", temp)
  temp = object_1.decrypt_ecb(temp)
  print("-Decrypted text (128-bit key):", temp, "\n")
  temp = object_2.encrypt_ecb(plain_text_3)
  print("-Encrypted text (192-bit key):", temp)
  temp = object_2.decrypt_ecb(temp)
  print("-Decrypted text (192-bit key):", temp, "\n")
  temp = object_3.encrypt_ecb(plain_text_3)
  print("-Encrypted text (256-bit key):", temp)
  temp = object_3.decrypt_ecb(temp)
  print("-Decrypted text (256-bit key):", temp, "\n")

  print("--------------------------------------------------\n")

def test_cbc():
  print("#####################")
  print("# Tests in CBC mode #")
  print("#####################\n")

  print("---> Plain text:", plain_text_1, "\n")
  temp = object_1.encrypt_cbc(plain_text_1, iv)
  print("-Encrypted text (128-bit key):", temp)
  temp = object_1.decrypt_cbc(temp, iv)
  print("-Decrypted text (128-bit key):", temp, "\n")
  temp = object_2.encrypt_cbc(plain_text_1, iv)
  print("-Encrypted text (192-bit key):", temp)
  temp = object_2.decrypt_cbc(temp, iv)
  print("-Decrypted text (192-bit key):", temp, "\n")
  temp = object_3.encrypt_cbc(plain_text_1, iv)
  print("-Encrypted text (256-bit key):", temp)
  temp = object_3.decrypt_cbc(temp, iv)
  print("-Decrypted text (256-bit key):", temp, "\n")

  print("-------------------------\n")

  print("---> Plain text:", plain_text_2, "\n")
  temp = object_1.encrypt_cbc(plain_text_2, iv)
  print("-Encrypted text (128-bit key):", temp)
  temp = object_1.decrypt_cbc(temp, iv)
  print("-Decrypted text (128-bit key):", temp, "\n")
  temp = object_2.encrypt_cbc(plain_text_2, iv)
  print("-Encrypted text (192-bit key):", temp)
  temp = object_2.decrypt_cbc(temp, iv)
  print("-Decrypted text (192-bit key):", temp, "\n")
  temp = object_3.encrypt_cbc(plain_text_2, iv)
  print("-Encrypted text (256-bit key):", temp)
  temp = object_3.decrypt_cbc(temp, iv)
  print("-Decrypted text (256-bit key):", temp, "\n")

  print("-------------------------\n")

  print("---> Plain text:", plain_text_3, "\n")
  temp = object_1.encrypt_cbc(plain_text_3, iv)
  print("-Encrypted text (128-bit key):", temp)
  temp = object_1.decrypt_cbc(temp, iv)
  print("-Decrypted text (128-bit key):", temp, "\n")
  temp = object_2.encrypt_cbc(plain_text_3, iv)
  print("-Encrypted text (192-bit key):", temp)
  temp = object_2.decrypt_cbc(temp, iv)
  print("-Decrypted text (192-bit key):", temp, "\n")
  temp = object_3.encrypt_cbc(plain_text_3, iv)
  print("-Encrypted text (256-bit key):", temp)
  temp = object_3.decrypt_cbc(temp, iv)
  print("-Decrypted text (256-bit key):", temp, "\n")

  print("--------------------------------------------------\n")

def test_ofb():
  print("#####################")
  print("# Tests in OFB mode #")
  print("#####################\n")

  print("---> Plain text:", plain_text_1, "\n")
  temp = object_1.encrypt_ofb(plain_text_1, iv)
  print("-Encrypted text (128-bit key):", temp)
  temp = object_1.decrypt_ofb(temp, iv)
  print("-Decrypted text (128-bit key):", temp, "\n")
  temp = object_2.encrypt_ofb(plain_text_1, iv)
  print("-Encrypted text (192-bit key):", temp)
  temp = object_2.decrypt_ofb(temp, iv)
  print("-Decrypted text (192-bit key):", temp, "\n")
  temp = object_3.encrypt_ofb(plain_text_1, iv)
  print("-Encrypted text (256-bit key):", temp)
  temp = object_3.decrypt_ofb(temp, iv)
  print("-Decrypted text (256-bit key):", temp, "\n")

  print("-------------------------\n")

  print("---> Plain text:", plain_text_2, "\n")
  temp = object_1.encrypt_ofb(plain_text_2, iv)
  print("-Encrypted text (128-bit key):", temp)
  temp = object_1.decrypt_ofb(temp, iv)
  print("-Decrypted text (128-bit key):", temp, "\n")
  temp = object_2.encrypt_ofb(plain_text_2, iv)
  print("-Encrypted text (192-bit key):", temp)
  temp = object_2.decrypt_ofb(temp, iv)
  print("-Decrypted text (192-bit key):", temp, "\n")
  temp = object_3.encrypt_ofb(plain_text_2, iv)
  print("-Encrypted text (256-bit key):", temp)
  temp = object_3.decrypt_ofb(temp, iv)
  print("-Decrypted text (256-bit key):", temp, "\n")

  print("-------------------------\n")

  print("---> Plain text:", plain_text_3, "\n")
  temp = object_1.encrypt_ofb(plain_text_3, iv)
  print("-Encrypted text (128-bit key):", temp)
  temp = object_1.decrypt_ofb(temp, iv)
  print("-Decrypted text (128-bit key):", temp, "\n")
  temp = object_2.encrypt_ofb(plain_text_3, iv)
  print("-Encrypted text (192-bit key):", temp)
  temp = object_2.decrypt_ofb(temp, iv)
  print("-Decrypted text (192-bit key):", temp, "\n")
  temp = object_3.encrypt_ofb(plain_text_3, iv)
  print("-Encrypted text (256-bit key):", temp)
  temp = object_3.decrypt_ofb(temp, iv)
  print("-Decrypted text (256-bit key):", temp, "\n")

  print("--------------------------------------------------\n")

if __name__ == "__main__":
  test_ecb()
  test_cbc()
  test_ofb()