{-# LANGUAGE OverloadedStrings, CPP #-}
module KAT (katTests) where

-- unfortunately due to a bug in some version of cabal
-- there's no way to have a condition cpp-options in the cabal file
-- for test suite. to run test with AESni, uncomment the following
-- #define HAVE_AESNI

import Test.Framework.Providers.QuickCheck2 (testProperty)
import qualified Data.ByteString as B
import qualified Data.ByteString.Char8 as BC

import Data.Word

import qualified Crypto.Cipher.AES as AES
import qualified Crypto.Cipher.Blowfish as Blowfish
import qualified Crypto.Cipher.Camellia as Camellia
import qualified Crypto.Cipher.RC4 as RC4

encryptStream fi fc key plaintext = B.unpack $ snd $ fc (fi $ B.pack key) plaintext

encryptBlock fi fc key plaintext = B.unpack $ fc (fi $ B.pack key) plaintext

unright (Right r) = r
unright (Left e) = error e

wordify :: [Char] -> [Word8]
wordify = map (toEnum . fromEnum)

vectors_aes128_enc =
	[
	  ( [0x10, 0xa5, 0x88, 0x69, 0xd7, 0x4b, 0xe5, 0xa3,0x74,0xcf,0x86,0x7c,0xfb,0x47,0x38,0x59]
	  , B.replicate 16 0
	  , [0x6d,0x25,0x1e,0x69,0x44,0xb0,0x51,0xe0,0x4e,0xaa,0x6f,0xb4,0xdb,0xf7,0x84,0x65]
	  )
	, ( replicate 16 0
	  , B.replicate 16 0
	  , [0x66,0xe9,0x4b,0xd4,0xef,0x8a,0x2c,0x3b,0x88,0x4c,0xfa,0x59,0xca,0x34,0x2b,0x2e]
	  )
	, ( replicate 16 0
	  , B.replicate 16 1
	  , [0xe1,0x4d,0x5d,0x0e,0xe2,0x77,0x15,0xdf,0x08,0xb4,0x15,0x2b,0xa2,0x3d,0xa8,0xe0]
	  )
	, ( replicate 16 1
	  , B.replicate 16 2
	  , [0x17,0xd6,0x14,0xf3,0x79,0xa9,0x35,0x90,0x77,0xe9,0x55,0x77,0xfd,0x31,0xc2,0x0a]
	  )
	, ( replicate 16 2
	  , B.replicate 16 1
	  , [0x8f,0x42,0xc2,0x4b,0xee,0x6e,0x63,0x47,0x2b,0x16,0x5a,0xa9,0x41,0x31,0x2f,0x7c]
	  )
	, ( replicate 16 3
	  , B.replicate 16 2
	  , [0x90,0x98,0x85,0xe4,0x77,0xbc,0x20,0xf5,0x8a,0x66,0x97,0x1d,0xa0,0xbc,0x75,0xe3]
	  )
	]

vectors_aes192_enc =
	[
	  ( replicate 24 0
	  , B.replicate 16 0
	  , [0xaa,0xe0,0x69,0x92,0xac,0xbf,0x52,0xa3,0xe8,0xf4,0xa9,0x6e,0xc9,0x30,0x0b,0xd7]
	  )
	, ( replicate 24 0
	  , B.replicate 16 1
	  , [0xcf,0x1e,0xce,0x3c,0x44,0xb0,0x78,0xfb,0x27,0xcb,0x0a,0x3e,0x07,0x1b,0x08,0x20]
	  )
	, ( replicate 24 1
	  , B.replicate 16 2
	  , [0xeb,0x8c,0x17,0x30,0x90,0xc7,0x5b,0x77,0xd6,0x72,0xb4,0x57,0xa7,0x78,0xd9,0xd0]
	  )
	, ( replicate 24 2
	  , B.replicate 16 1
	  , [0xf2,0xf0,0xae,0xd8,0xcd,0xc9,0x21,0xca,0x4b,0x55,0x84,0x5d,0xa4,0x15,0x21,0xc2]
	  )
	, ( replicate 24 3
	  , B.replicate 16 2
	  , [0xca,0xcc,0x30,0x79,0xe4,0xb7,0x95,0x27,0x63,0xd2,0x55,0xd6,0x34,0x10,0x46,0x14]
	  )
	]

vectors_aes256_enc =
	[ ( replicate 32 0
	  , B.replicate 16 0
	  , [0xdc,0x95,0xc0,0x78,0xa2,0x40,0x89,0x89,0xad,0x48,0xa2,0x14,0x92,0x84,0x20,0x87]
	  )
	, ( replicate 32 0
	  , B.replicate 16 1
	  , [0x7b,0xc3,0x02,0x6c,0xd7,0x37,0x10,0x3e,0x62,0x90,0x2b,0xcd,0x18,0xfb,0x01,0x63]
	  )
	, ( replicate 32 1
	  , B.replicate 16 2
	  , [0x62,0xae,0x12,0xf3,0x24,0xbf,0xea,0x08,0xd5,0xf6,0x75,0xb5,0x13,0x02,0x6b,0xbf]
	  )
	, ( replicate 32 2
	  , B.replicate 16 1
	  , [0x00,0xf9,0xc7,0x44,0x4b,0xb0,0xcc,0x80,0x6c,0x7c,0x39,0xee,0x22,0x11,0xf1,0x46]
	  )
	, ( replicate 32 3
	  , B.replicate 16 2
	  , [0xb4,0x05,0x87,0x3e,0xa0,0x76,0x1b,0x9c,0xa9,0x9f,0x70,0xb0,0x16,0x16,0xce,0xb1]
	  )
	]

vectors_aes128_dec =
	[ ( replicate 16 0
	  , B.replicate 16 0
	  , [0x14,0x0f,0x0f,0x10,0x11,0xb5,0x22,0x3d,0x79,0x58,0x77,0x17,0xff,0xd9,0xec,0x3a]
	  )
	, ( replicate 16 0
	  , B.replicate 16 1
	  , [0x15,0x6d,0x0f,0x85,0x75,0xd5,0x33,0x07,0x52,0xf8,0x4a,0xf2,0x72,0xff,0x30,0x50]
	  )
	, ( replicate 16 1
	  , B.replicate 16 2
	  , [0x34,0x37,0xd6,0xe2,0x31,0xd7,0x02,0x41,0x9b,0x51,0xb4,0x94,0x72,0x71,0xb6,0x11]
	  )
	, ( replicate 16 2
	  , B.replicate 16 1
	  , [0xe3,0xcd,0xe2,0x37,0xc8,0xf2,0xd9,0x7b,0x8d,0x79,0xf9,0x17,0x1d,0x4b,0xda,0xc1]
	  )
	, ( replicate 16 3
	  , B.replicate 16 2
	  , [0x5b,0x94,0xaa,0xed,0xd7,0x83,0x99,0x8c,0xd5,0x15,0x35,0x35,0x18,0xcc,0x45,0xe2]
	  )
	]

vectors_aes192_dec =
	[
	  ( replicate 24 0
	  , B.replicate 16 0
	  , [0x13,0x46,0x0e,0x87,0xa8,0xfc,0x02,0x3e,0xf2,0x50,0x1a,0xfe,0x7f,0xf5,0x1c,0x51]
	  )
	, ( replicate 24 0
	  , B.replicate 16 1
	  , [0x92,0x17,0x07,0xc3,0x3d,0x1c,0xc5,0x96,0x7d,0xa5,0x1d,0xbb,0xb0,0x66,0xb2,0x6c]
	  )
	, ( replicate 24 1
	  , B.replicate 16 2
	  , [0xee,0x92,0x97,0xc6,0xba,0xe8,0x26,0x4d,0xff,0x08,0x0e,0xbb,0x1e,0x74,0x11,0xc1]
	  )
	, ( replicate 24 2
	  , B.replicate 16 1
	  , [0x49,0x67,0xdf,0x70,0xd2,0x9e,0x9a,0x7f,0x5d,0x7c,0xb9,0xc1,0x20,0xc3,0x8a,0x71]
	  )
	, ( replicate 24 3
	  , B.replicate 16 2
	  , [0x74,0x38,0x62,0x42,0x6b,0x56,0x7f,0xd5,0xf0,0x1d,0x1b,0x59,0x56,0x01,0x26,0x29]
	  )
	]

vectors_aes256_dec =
	[ ( replicate 32 0
	  , B.replicate 16 0
	  , [0x67,0x67,0x1c,0xe1,0xfa,0x91,0xdd,0xeb,0x0f,0x8f,0xbb,0xb3,0x66,0xb5,0x31,0xb4]
	  )
	, ( replicate 32 0
	  , B.replicate 16 1
	  , [0xcc,0x09,0x21,0xa3,0xc5,0xca,0x17,0xf7,0x48,0xb7,0xc2,0x7b,0x73,0xba,0x87,0xa2]
	  )
	, ( replicate 32 1
	  , B.replicate 16 2
	  , [0xc0,0x4b,0x27,0x90,0x1a,0x50,0xcf,0xfa,0xf1,0xbb,0x88,0x9f,0xc0,0x92,0x5e,0x14]
	  )
	, ( replicate 32 2
	  , B.replicate 16 1
	  , [0x24,0x61,0x53,0x5d,0x16,0x1c,0x15,0x39,0x88,0x32,0x77,0x29,0xc5,0x8c,0xc0,0x3a]
	  )
	, ( replicate 32 3
	  , B.replicate 16 2
	  , [0x30,0xc9,0x1c,0xce,0xfe,0x89,0x30,0xcf,0xff,0x31,0xdb,0xcc,0xfc,0x11,0xc5,0x23]
	  )
	]

aes128InitKey = AES.initKey
aes192InitKey = AES.initKey
aes256InitKey = AES.initKey

vectors_rc4 =
	[ (wordify "Key", "Plaintext", [ 0xBB,0xF3,0x16,0xE8,0xD9,0x40,0xAF,0x0A,0xD3 ])
	, (wordify "Wiki", "pedia", [ 0x10,0x21,0xBF,0x04,0x20 ])
	, (wordify "Secret", "Attack at dawn", [ 0x45,0xA0,0x1F,0x64,0x5F,0xC3,0x5B,0x38,0x35,0x52,0x54,0x4B,0x9B,0xF5 ])
	]

vectors_camellia128 =
	[ 
	  ( replicate 16 0
	  , B.replicate 16 0
	  , [0x3d,0x02,0x80,0x25,0xb1,0x56,0x32,0x7c,0x17,0xf7,0x62,0xc1,0xf2,0xcb,0xca,0x71]
	  )
	, ( [0x01,0x23,0x45,0x67,0x89,0xab,0xcd,0xef,0xfe,0xdc,0xba,0x98,0x76,0x54,0x32,0x10]
	  , B.pack [0x01,0x23,0x45,0x67,0x89,0xab,0xcd,0xef,0xfe,0xdc,0xba,0x98,0x76,0x54,0x32,0x10]
	  , [0x67,0x67,0x31,0x38,0x54,0x96,0x69,0x73,0x08,0x57,0x06,0x56,0x48,0xea,0xbe,0x43]
	  )
	]

vectors_camellia192 =
	[
	  ( [0x01,0x23,0x45,0x67,0x89,0xab,0xcd,0xef,0xfe,0xdc,0xba,0x98,0x76,0x54,0x32,0x10,0x00,0x11,0x22,0x33,0x44,0x55,0x66,0x77]
	  , B.pack [0x01,0x23,0x45,0x67,0x89,0xab,0xcd,0xef,0xfe,0xdc,0xba,0x98,0x76,0x54,0x32,0x10]
	  ,[0xb4,0x99,0x34,0x01,0xb3,0xe9,0x96,0xf8,0x4e,0xe5,0xce,0xe7,0xd7,0x9b,0x09,0xb9]
	  )
	]

vectors_camellia256 =
	[
	  ( [0x01,0x23,0x45,0x67,0x89,0xab,0xcd,0xef,0xfe,0xdc,0xba,0x98,0x76,0x54,0x32,0x10
	    ,0x00,0x11,0x22,0x33,0x44,0x55,0x66,0x77,0x88,0x99,0xaa,0xbb,0xcc,0xdd,0xee,0xff]
	  , B.pack [0x01,0x23,0x45,0x67,0x89,0xab,0xcd,0xef,0xfe,0xdc,0xba,0x98,0x76,0x54,0x32,0x10]
	  , [0x9a,0xcc,0x23,0x7d,0xff,0x16,0xd7,0x6c,0x20,0xef,0x7c,0x91,0x9e,0x3a,0x75,0x09]
	  )
	]

vectors_blowfish =
    [
      ( replicate 8 0
      , B.replicate 8 0
      , [0x4e,0xf9,0x97,0x45,0x61,0x98,0xDD,0x78]
      )
    , ( replicate 8 255
      , B.replicate 8 255
      , [0x51,0x86,0x6F,0xD5,0xB8,0x5E,0xCB,0x8A]
      )
    , ( [0x7C,0xA1,0x10,0x45,0x4A,0x1A,0x6E,0x57]
      , B.pack [0x01,0xA1,0xD6,0xD0,0x39,0x77,0x67,0x42]
      , [0x59,0xC6,0x82,0x45,0xEB,0x05,0x28,0x2B]
      )
    ]

vectors =
	[ ("RC4",        vectors_rc4,         encryptStream RC4.initCtx RC4.encrypt)
	-- AES haskell implementation
	, ("AES 128 Enc", vectors_aes128_enc,  encryptBlock aes128InitKey AES.encryptECB)
	, ("AES 192 Enc", vectors_aes192_enc,  encryptBlock aes192InitKey AES.encryptECB)
	, ("AES 256 Enc", vectors_aes256_enc,  encryptBlock aes256InitKey AES.encryptECB)
	, ("AES 128 Dec", vectors_aes128_dec,  encryptBlock aes128InitKey AES.decryptECB)
	, ("AES 192 Dec", vectors_aes192_dec,  encryptBlock aes192InitKey AES.decryptECB)
	, ("AES 256 Dec", vectors_aes256_dec,  encryptBlock aes256InitKey AES.decryptECB)
    -- Blowfish implementation
    , ("Blowfish",   vectors_blowfish,    encryptBlock (unright . Blowfish.initKey) Blowfish.encrypt)
	-- Camellia implementation
	, ("Camellia",   vectors_camellia128, encryptBlock (unright . Camellia.initKey128) Camellia.encrypt)
	]

katTests = map makeTests vectors
	where makeTests (name, v, f) = testProperty name (and $ map makeTest v)
		where makeTest (key,plaintext,expected) = assertEq expected $ f key plaintext

assertEq expected got
	| expected == got = True
	| otherwise       = error ("expected: " ++ show expected ++ " got: " ++ show got)