using Base.Test  

importall AES

ctx = aes_context(Array(Uint32,64),Array(Uint32,64),0)
encrypted = Array(Uint8,16)
decrypted = Array(Uint8,16)


println("AES - NIST FIPS PUB 197 test cases:")
plaintext = b"\x00\x11\x22\x33\x44\x55\x66\x77\x88\x99\xaa\xbb\xcc\xdd\xee\xff"

print(" C.1 AES-128 ")
key128 = b"\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f"
aes_set_key(ctx, key128, 128)
aes_encrypt(ctx, plaintext, encrypted); @test encrypted==b"\x69\xc4\xe0\xd8\x6a\x7b\x04\x30\xd8\xcd\xb7\x80\x70\xb4\xc5\x5a"
aes_decrypt(ctx, encrypted, decrypted); @test decrypted==b"\x00\x11\x22\x33\x44\x55\x66\x77\x88\x99\xaa\xbb\xcc\xdd\xee\xff"
println("- Passed")

print(" C.2 AES-192 ")
key192 = b"\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f\x10\x11\x12\x13\x14\x15\x16\x17"
aes_set_key(ctx, key192, 192)
aes_encrypt(ctx, plaintext, encrypted); @test encrypted==b"\xdd\xa9\x7c\xa4\x86\x4c\xdf\xe0\x6e\xaf\x70\xa0\xec\x0d\x71\x91"
aes_decrypt(ctx, encrypted, decrypted); @test decrypted==b"\x00\x11\x22\x33\x44\x55\x66\x77\x88\x99\xaa\xbb\xcc\xdd\xee\xff"
println("- Passed")

print(" C.3 AES-256 ")
key256 = b"\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\x1a\x1b\x1c\x1d\x1e\x1f"
aes_set_key(ctx, key256, 256)
aes_encrypt(ctx, plaintext, encrypted); @test encrypted==b"\x8e\xa2\xb7\xca\x51\x67\x45\xbf\xea\xfc\x49\x90\x4b\x49\x60\x89"
aes_decrypt(ctx, encrypted, decrypted); @test decrypted==b"\x00\x11\x22\x33\x44\x55\x66\x77\x88\x99\xaa\xbb\xcc\xdd\xee\xff"
println("- Passed")



#==============================================================================
For fun you could scramble your data the other way: first with aes_decrypt, 
then aes_encrypt will get it back when provided with the same key.
==============================================================================#


println("NIST - Spec. Publ. 800-38A 2001 ED test cases:")
TxtBlock1 = b"\x6b\xc1\xbe\xe2\x2e\x40\x9f\x96\xe9\x3d\x7e\x11\x73\x93\x17\x2a"
TxtBlock2 = b"\xae\x2d\x8a\x57\x1e\x03\xac\x9c\x9e\xb7\x6f\xac\x45\xaf\x8e\x51"
TxtBlock3 = b"\x30\xc8\x1c\x46\xa3\x5c\xe4\x11\xe5\xfb\xc1\x19\x1a\x0a\x52\xef"
TxtBlock4 = b"\xf6\x9f\x24\x45\xdf\x4f\x9b\x17\xad\x2b\x41\x7b\xe6\x6c\x37\x10"
 

print(" F.1.1 ECB-AES128.Encrypt; F.1.2 ECB-AES128.Decrypt ")
Key128 = b"\x2b\x7e\x15\x16\x28\xae\xd2\xa6\xab\xf7\x15\x88\x09\xcf\x4f\x3c"
aes_set_key(ctx, Key128, 128)
aes_encrypt(ctx, TxtBlock1, encrypted); @test encrypted==b"\x3a\xd7\x7b\xb4\x0d\x7a\x36\x60\xa8\x9e\xca\xf3\x24\x66\xef\x97"
aes_decrypt(ctx, encrypted, decrypted); @test decrypted==b"\x6b\xc1\xbe\xe2\x2e\x40\x9f\x96\xe9\x3d\x7e\x11\x73\x93\x17\x2a"

aes_encrypt(ctx, TxtBlock2, encrypted); @test encrypted==b"\xf5\xd3\xd5\x85\x03\xb9\x69\x9d\xe7\x85\x89\x5a\x96\xfd\xba\xaf"
aes_decrypt(ctx, encrypted, decrypted); @test decrypted==b"\xae\x2d\x8a\x57\x1e\x03\xac\x9c\x9e\xb7\x6f\xac\x45\xaf\x8e\x51"

aes_encrypt(ctx, TxtBlock3, encrypted); @test encrypted==b"\x43\xb1\xcd\x7f\x59\x8e\xce\x23\x88\x1b\x00\xe3\xed\x03\x06\x88"
aes_decrypt(ctx, encrypted, decrypted); @test decrypted==b"\x30\xc8\x1c\x46\xa3\x5c\xe4\x11\xe5\xfb\xc1\x19\x1a\x0a\x52\xef"

aes_encrypt(ctx, TxtBlock4, encrypted); @test encrypted==b"\x7b\x0c\x78\x5e\x27\xe8\xad\x3f\x82\x23\x20\x71\x04\x72\x5d\xd4"
aes_decrypt(ctx, encrypted, decrypted); @test decrypted==b"\xf6\x9f\x24\x45\xdf\x4f\x9b\x17\xad\x2b\x41\x7b\xe6\x6c\x37\x10"
println("- Passed")


print(" F.1.3 ECB-AES192.Encrypt; F.1.4 ECB-AES192.Decrypt ")
Key192 = b"\x8e\x73\xb0\xf7\xda\x0e\x64\x52\xc8\x10\xf3\x2b\x80\x90\x79\xe5\x62\xf8\xea\xd2\x52\x2c\x6b\x7b"
aes_set_key(ctx, Key192, 192)
aes_encrypt(ctx, TxtBlock1, encrypted); @test encrypted==b"\xbd\x33\x4f\x1d\x6e\x45\xf2\x5f\xf7\x12\xa2\x14\x57\x1f\xa5\xcc"
aes_decrypt(ctx, encrypted, decrypted); @test decrypted==b"\x6b\xc1\xbe\xe2\x2e\x40\x9f\x96\xe9\x3d\x7e\x11\x73\x93\x17\x2a"

aes_encrypt(ctx, TxtBlock2, encrypted); @test encrypted==b"\x97\x41\x04\x84\x6d\x0a\xd3\xad\x77\x34\xec\xb3\xec\xee\x4e\xef"
aes_decrypt(ctx, encrypted, decrypted); @test decrypted==b"\xae\x2d\x8a\x57\x1e\x03\xac\x9c\x9e\xb7\x6f\xac\x45\xaf\x8e\x51"

aes_encrypt(ctx, TxtBlock3, encrypted); @test encrypted==b"\xef\x7a\xfd\x22\x70\xe2\xe6\x0a\xdc\xe0\xba\x2f\xac\xe6\x44\x4e"
aes_decrypt(ctx, encrypted, decrypted); @test decrypted==b"\x30\xc8\x1c\x46\xa3\x5c\xe4\x11\xe5\xfb\xc1\x19\x1a\x0a\x52\xef"

aes_encrypt(ctx, TxtBlock4, encrypted); @test encrypted==b"\x9a\x4b\x41\xba\x73\x8d\x6c\x72\xfb\x16\x69\x16\x03\xc1\x8e\x0e"
aes_decrypt(ctx, encrypted, decrypted); @test decrypted==b"\xf6\x9f\x24\x45\xdf\x4f\x9b\x17\xad\x2b\x41\x7b\xe6\x6c\x37\x10"
println("- Passed")


print(" F.1.5 ECB-AES256.Encrypt; F.1.6 ECB-AES256.Decrypt ")
Key256 = b"\x60\x3d\xeb\x10\x15\xca\x71\xbe\x2b\x73\xae\xf0\x85\x7d\x77\x81\x1f\x35\x2c\x07\x3b\x61\x08\xd7\x2d\x98\x10\xa3\x09\x14\xdf\xf4"
aes_set_key(ctx, Key256, 256)
aes_encrypt(ctx, TxtBlock1, encrypted); @test encrypted==b"\xf3\xee\xd1\xbd\xb5\xd2\xa0\x3c\x06\x4b\x5a\x7e\x3d\xb1\x81\xf8"
aes_decrypt(ctx, encrypted, decrypted); @test decrypted==b"\x6b\xc1\xbe\xe2\x2e\x40\x9f\x96\xe9\x3d\x7e\x11\x73\x93\x17\x2a"

aes_encrypt(ctx, TxtBlock2, encrypted); @test encrypted==b"\x59\x1c\xcb\x10\xd4\x10\xed\x26\xdc\x5b\xa7\x4a\x31\x36\x28\x70"
aes_decrypt(ctx, encrypted, decrypted); @test decrypted==b"\xae\x2d\x8a\x57\x1e\x03\xac\x9c\x9e\xb7\x6f\xac\x45\xaf\x8e\x51"

aes_encrypt(ctx, TxtBlock3, encrypted); @test encrypted==b"\xb6\xed\x21\xb9\x9c\xa6\xf4\xf9\xf1\x53\xe7\xb1\xbe\xaf\xed\x1d"
aes_decrypt(ctx, encrypted, decrypted); @test decrypted==b"\x30\xc8\x1c\x46\xa3\x5c\xe4\x11\xe5\xfb\xc1\x19\x1a\x0a\x52\xef"

aes_encrypt(ctx, TxtBlock4, encrypted); @test encrypted==b"\x23\x30\x4b\x7a\x39\xf9\xf3\xff\x06\x7d\x8d\x8f\x9e\x24\xec\xc7"
aes_decrypt(ctx, encrypted, decrypted); @test decrypted==b"\xf6\x9f\x24\x45\xdf\x4f\x9b\x17\xad\x2b\x41\x7b\xe6\x6c\x37\x10"
println("- Passed")



#==============================================================================
File encryption/decryption utility examples:

file_encrypt("original_file.txt", "encrypted_128.aes", 128, "up_to_16char_pwd")
file_decrypt("encrypted_128.aes", "original_file.txt", 128, "up_to_16char_pwd")

file_encrypt("original_file.txt", "encrypted_192.aes", 192, "upto24character_password")
file_decrypt("encrypted_192.aes", "original_file.txt", 192, "upto24character_password")

file_encrypt("original_file.txt", "encrypted_256.aes", 256, "longest_upto32character_password")
file_decrypt("encrypted_256.aes", "original_file.txt", 256, "longest_upto32character_password")

==============================================================================#
