3 | module Crypto.AES.Small

  5 | import Utils.ConstantTable

  6 | import Crypto.AES.Common

  7 | import Data.Bits

  8 | import Data.DPair

  9 | import Data.Fin

 10 | import Data.List

 11 | import Data.Nat

 12 | import Data.Stream

 13 | import Data.Vect

 14 | import Utils.Misc

 15 | import Utils.Bytes

 17 | gbox_9 : ConstantTable 256 Bits8

 18 | gbox_9 = from_vect

 19 |   [ 0x00, 0x09, 0x12, 0x1b, 0x24, 0x2d, 0x36, 0x3f, 0x48, 0x41, 0x5a, 0x53, 0x6c, 0x65, 0x7e, 0x77

 20 |   , 0x90, 0x99, 0x82, 0x8b, 0xb4, 0xbd, 0xa6, 0xaf, 0xd8, 0xd1, 0xca, 0xc3, 0xfc, 0xf5, 0xee, 0xe7

 21 |   , 0x3b, 0x32, 0x29, 0x20, 0x1f, 0x16, 0x0d, 0x04, 0x73, 0x7a, 0x61, 0x68, 0x57, 0x5e, 0x45, 0x4c

 22 |   , 0xab, 0xa2, 0xb9, 0xb0, 0x8f, 0x86, 0x9d, 0x94, 0xe3, 0xea, 0xf1, 0xf8, 0xc7, 0xce, 0xd5, 0xdc

 23 |   , 0x76, 0x7f, 0x64, 0x6d, 0x52, 0x5b, 0x40, 0x49, 0x3e, 0x37, 0x2c, 0x25, 0x1a, 0x13, 0x08, 0x01

 24 |   , 0xe6, 0xef, 0xf4, 0xfd, 0xc2, 0xcb, 0xd0, 0xd9, 0xae, 0xa7, 0xbc, 0xb5, 0x8a, 0x83, 0x98, 0x91

 25 |   , 0x4d, 0x44, 0x5f, 0x56, 0x69, 0x60, 0x7b, 0x72, 0x05, 0x0c, 0x17, 0x1e, 0x21, 0x28, 0x33, 0x3a

 26 |   , 0xdd, 0xd4, 0xcf, 0xc6, 0xf9, 0xf0, 0xeb, 0xe2, 0x95, 0x9c, 0x87, 0x8e, 0xb1, 0xb8, 0xa3, 0xaa

 27 |   , 0xec, 0xe5, 0xfe, 0xf7, 0xc8, 0xc1, 0xda, 0xd3, 0xa4, 0xad, 0xb6, 0xbf, 0x80, 0x89, 0x92, 0x9b

 28 |   , 0x7c, 0x75, 0x6e, 0x67, 0x58, 0x51, 0x4a, 0x43, 0x34, 0x3d, 0x26, 0x2f, 0x10, 0x19, 0x02, 0x0b

 29 |   , 0xd7, 0xde, 0xc5, 0xcc, 0xf3, 0xfa, 0xe1, 0xe8, 0x9f, 0x96, 0x8d, 0x84, 0xbb, 0xb2, 0xa9, 0xa0

 30 |   , 0x47, 0x4e, 0x55, 0x5c, 0x63, 0x6a, 0x71, 0x78, 0x0f, 0x06, 0x1d, 0x14, 0x2b, 0x22, 0x39, 0x30

 31 |   , 0x9a, 0x93, 0x88, 0x81, 0xbe, 0xb7, 0xac, 0xa5, 0xd2, 0xdb, 0xc0, 0xc9, 0xf6, 0xff, 0xe4, 0xed

 32 |   , 0x0a, 0x03, 0x18, 0x11, 0x2e, 0x27, 0x3c, 0x35, 0x42, 0x4b, 0x50, 0x59, 0x66, 0x6f, 0x74, 0x7d

 33 |   , 0xa1, 0xa8, 0xb3, 0xba, 0x85, 0x8c, 0x97, 0x9e, 0xe9, 0xe0, 0xfb, 0xf2, 0xcd, 0xc4, 0xdf, 0xd6

 34 |   , 0x31, 0x38, 0x23, 0x2a, 0x15, 0x1c, 0x07, 0x0e, 0x79, 0x70, 0x6b, 0x62, 0x5d, 0x54, 0x4f, 0x46

 35 |   ]

 37 | gbox_11 : ConstantTable 256 Bits8

 38 | gbox_11 = from_vect

 39 |   [ 0x00, 0x0b, 0x16, 0x1d, 0x2c, 0x27, 0x3a, 0x31, 0x58, 0x53, 0x4e, 0x45, 0x74, 0x7f, 0x62, 0x69

 40 |   , 0xb0, 0xbb, 0xa6, 0xad, 0x9c, 0x97, 0x8a, 0x81, 0xe8, 0xe3, 0xfe, 0xf5, 0xc4, 0xcf, 0xd2, 0xd9

 41 |   , 0x7b, 0x70, 0x6d, 0x66, 0x57, 0x5c, 0x41, 0x4a, 0x23, 0x28, 0x35, 0x3e, 0x0f, 0x04, 0x19, 0x12

 42 |   , 0xcb, 0xc0, 0xdd, 0xd6, 0xe7, 0xec, 0xf1, 0xfa, 0x93, 0x98, 0x85, 0x8e, 0xbf, 0xb4, 0xa9, 0xa2

 43 |   , 0xf6, 0xfd, 0xe0, 0xeb, 0xda, 0xd1, 0xcc, 0xc7, 0xae, 0xa5, 0xb8, 0xb3, 0x82, 0x89, 0x94, 0x9f

 44 |   , 0x46, 0x4d, 0x50, 0x5b, 0x6a, 0x61, 0x7c, 0x77, 0x1e, 0x15, 0x08, 0x03, 0x32, 0x39, 0x24, 0x2f

 45 |   , 0x8d, 0x86, 0x9b, 0x90, 0xa1, 0xaa, 0xb7, 0xbc, 0xd5, 0xde, 0xc3, 0xc8, 0xf9, 0xf2, 0xef, 0xe4

 46 |   , 0x3d, 0x36, 0x2b, 0x20, 0x11, 0x1a, 0x07, 0x0c, 0x65, 0x6e, 0x73, 0x78, 0x49, 0x42, 0x5f, 0x54

 47 |   , 0xf7, 0xfc, 0xe1, 0xea, 0xdb, 0xd0, 0xcd, 0xc6, 0xaf, 0xa4, 0xb9, 0xb2, 0x83, 0x88, 0x95, 0x9e

 48 |   , 0x47, 0x4c, 0x51, 0x5a, 0x6b, 0x60, 0x7d, 0x76, 0x1f, 0x14, 0x09, 0x02, 0x33, 0x38, 0x25, 0x2e

 49 |   , 0x8c, 0x87, 0x9a, 0x91, 0xa0, 0xab, 0xb6, 0xbd, 0xd4, 0xdf, 0xc2, 0xc9, 0xf8, 0xf3, 0xee, 0xe5

 50 |   , 0x3c, 0x37, 0x2a, 0x21, 0x10, 0x1b, 0x06, 0x0d, 0x64, 0x6f, 0x72, 0x79, 0x48, 0x43, 0x5e, 0x55

 51 |   , 0x01, 0x0a, 0x17, 0x1c, 0x2d, 0x26, 0x3b, 0x30, 0x59, 0x52, 0x4f, 0x44, 0x75, 0x7e, 0x63, 0x68

 52 |   , 0xb1, 0xba, 0xa7, 0xac, 0x9d, 0x96, 0x8b, 0x80, 0xe9, 0xe2, 0xff, 0xf4, 0xc5, 0xce, 0xd3, 0xd8

 53 |   , 0x7a, 0x71, 0x6c, 0x67, 0x56, 0x5d, 0x40, 0x4b, 0x22, 0x29, 0x34, 0x3f, 0x0e, 0x05, 0x18, 0x13

 54 |   , 0xca, 0xc1, 0xdc, 0xd7, 0xe6, 0xed, 0xf0, 0xfb, 0x92, 0x99, 0x84, 0x8f, 0xbe, 0xb5, 0xa8, 0xa3

 55 |   ]

 57 | gbox_13 : ConstantTable 256 Bits8

 58 | gbox_13 = from_vect

 59 |   [ 0x00, 0x0d, 0x1a, 0x17, 0x34, 0x39, 0x2e, 0x23, 0x68, 0x65, 0x72, 0x7f, 0x5c, 0x51, 0x46, 0x4b

 60 |   , 0xd0, 0xdd, 0xca, 0xc7, 0xe4, 0xe9, 0xfe, 0xf3, 0xb8, 0xb5, 0xa2, 0xaf, 0x8c, 0x81, 0x96, 0x9b

 61 |   , 0xbb, 0xb6, 0xa1, 0xac, 0x8f, 0x82, 0x95, 0x98, 0xd3, 0xde, 0xc9, 0xc4, 0xe7, 0xea, 0xfd, 0xf0

 62 |   , 0x6b, 0x66, 0x71, 0x7c, 0x5f, 0x52, 0x45, 0x48, 0x03, 0x0e, 0x19, 0x14, 0x37, 0x3a, 0x2d, 0x20

 63 |   , 0x6d, 0x60, 0x77, 0x7a, 0x59, 0x54, 0x43, 0x4e, 0x05, 0x08, 0x1f, 0x12, 0x31, 0x3c, 0x2b, 0x26

 64 |   , 0xbd, 0xb0, 0xa7, 0xaa, 0x89, 0x84, 0x93, 0x9e, 0xd5, 0xd8, 0xcf, 0xc2, 0xe1, 0xec, 0xfb, 0xf6

 65 |   , 0xd6, 0xdb, 0xcc, 0xc1, 0xe2, 0xef, 0xf8, 0xf5, 0xbe, 0xb3, 0xa4, 0xa9, 0x8a, 0x87, 0x90, 0x9d

 66 |   , 0x06, 0x0b, 0x1c, 0x11, 0x32, 0x3f, 0x28, 0x25, 0x6e, 0x63, 0x74, 0x79, 0x5a, 0x57, 0x40, 0x4d

 67 |   , 0xda, 0xd7, 0xc0, 0xcd, 0xee, 0xe3, 0xf4, 0xf9, 0xb2, 0xbf, 0xa8, 0xa5, 0x86, 0x8b, 0x9c, 0x91

 68 |   , 0x0a, 0x07, 0x10, 0x1d, 0x3e, 0x33, 0x24, 0x29, 0x62, 0x6f, 0x78, 0x75, 0x56, 0x5b, 0x4c, 0x41

 69 |   , 0x61, 0x6c, 0x7b, 0x76, 0x55, 0x58, 0x4f, 0x42, 0x09, 0x04, 0x13, 0x1e, 0x3d, 0x30, 0x27, 0x2a

 70 |   , 0xb1, 0xbc, 0xab, 0xa6, 0x85, 0x88, 0x9f, 0x92, 0xd9, 0xd4, 0xc3, 0xce, 0xed, 0xe0, 0xf7, 0xfa

 71 |   , 0xb7, 0xba, 0xad, 0xa0, 0x83, 0x8e, 0x99, 0x94, 0xdf, 0xd2, 0xc5, 0xc8, 0xeb, 0xe6, 0xf1, 0xfc

 72 |   , 0x67, 0x6a, 0x7d, 0x70, 0x53, 0x5e, 0x49, 0x44, 0x0f, 0x02, 0x15, 0x18, 0x3b, 0x36, 0x21, 0x2c

 73 |   , 0x0c, 0x01, 0x16, 0x1b, 0x38, 0x35, 0x22, 0x2f, 0x64, 0x69, 0x7e, 0x73, 0x50, 0x5d, 0x4a, 0x47

 74 |   , 0xdc, 0xd1, 0xc6, 0xcb, 0xe8, 0xe5, 0xf2, 0xff, 0xb4, 0xb9, 0xae, 0xa3, 0x80, 0x8d, 0x9a, 0x97

 75 |   ]

 77 | gbox_14 : ConstantTable 256 Bits8

 78 | gbox_14 = from_vect

 79 |   [ 0x00, 0x0e, 0x1c, 0x12, 0x38, 0x36, 0x24, 0x2a, 0x70, 0x7e, 0x6c, 0x62, 0x48, 0x46, 0x54, 0x5a

 80 |   , 0xe0, 0xee, 0xfc, 0xf2, 0xd8, 0xd6, 0xc4, 0xca, 0x90, 0x9e, 0x8c, 0x82, 0xa8, 0xa6, 0xb4, 0xba

 81 |   , 0xdb, 0xd5, 0xc7, 0xc9, 0xe3, 0xed, 0xff, 0xf1, 0xab, 0xa5, 0xb7, 0xb9, 0x93, 0x9d, 0x8f, 0x81

 82 |   , 0x3b, 0x35, 0x27, 0x29, 0x03, 0x0d, 0x1f, 0x11, 0x4b, 0x45, 0x57, 0x59, 0x73, 0x7d, 0x6f, 0x61

 83 |   , 0xad, 0xa3, 0xb1, 0xbf, 0x95, 0x9b, 0x89, 0x87, 0xdd, 0xd3, 0xc1, 0xcf, 0xe5, 0xeb, 0xf9, 0xf7

 84 |   , 0x4d, 0x43, 0x51, 0x5f, 0x75, 0x7b, 0x69, 0x67, 0x3d, 0x33, 0x21, 0x2f, 0x05, 0x0b, 0x19, 0x17

 85 |   , 0x76, 0x78, 0x6a, 0x64, 0x4e, 0x40, 0x52, 0x5c, 0x06, 0x08, 0x1a, 0x14, 0x3e, 0x30, 0x22, 0x2c

 86 |   , 0x96, 0x98, 0x8a, 0x84, 0xae, 0xa0, 0xb2, 0xbc, 0xe6, 0xe8, 0xfa, 0xf4, 0xde, 0xd0, 0xc2, 0xcc

 87 |   , 0x41, 0x4f, 0x5d, 0x53, 0x79, 0x77, 0x65, 0x6b, 0x31, 0x3f, 0x2d, 0x23, 0x09, 0x07, 0x15, 0x1b

 88 |   , 0xa1, 0xaf, 0xbd, 0xb3, 0x99, 0x97, 0x85, 0x8b, 0xd1, 0xdf, 0xcd, 0xc3, 0xe9, 0xe7, 0xf5, 0xfb

 89 |   , 0x9a, 0x94, 0x86, 0x88, 0xa2, 0xac, 0xbe, 0xb0, 0xea, 0xe4, 0xf6, 0xf8, 0xd2, 0xdc, 0xce, 0xc0

 90 |   , 0x7a, 0x74, 0x66, 0x68, 0x42, 0x4c, 0x5e, 0x50, 0x0a, 0x04, 0x16, 0x18, 0x32, 0x3c, 0x2e, 0x20

 91 |   , 0xec, 0xe2, 0xf0, 0xfe, 0xd4, 0xda, 0xc8, 0xc6, 0x9c, 0x92, 0x80, 0x8e, 0xa4, 0xaa, 0xb8, 0xb6

 92 |   , 0x0c, 0x02, 0x10, 0x1e, 0x34, 0x3a, 0x28, 0x26, 0x7c, 0x72, 0x60, 0x6e, 0x44, 0x4a, 0x58, 0x56

 93 |   , 0x37, 0x39, 0x2b, 0x25, 0x0f, 0x01, 0x13, 0x1d, 0x47, 0x49, 0x5b, 0x55, 0x7f, 0x71, 0x63, 0x6d

 94 |   , 0xd7, 0xd9, 0xcb, 0xc5, 0xef, 0xe1, 0xf3, 0xfd, 0xa7, 0xa9, 0xbb, 0xb5, 0x9f, 0x91, 0x83, 0x8d

 95 |   ]

 97 | data G : Type where

 98 |   G1 : G

 99 |   G2 : G

100 |   G3 : G

102 | data G' : Type where

103 |   G9 : G'

104 |   G11 : G'

105 |   G13 : G'

106 |   G14 : G'

108 | mix_columns_mat : Vect 4 (Vect 4 G)

109 | mix_columns_mat = [[G2, G1, G1, G3], [G3, G2, G1, G1], [G1, G3, G2, G1], [G1, G1, G3, G2]]

111 | inv_mix_columns_mat : Vect 4 (Vect 4 G')

112 | inv_mix_columns_mat = [[G14, G9, G13, G11], [G11, G14, G9, G13], [G13, G11, G14, G9], [G9, G13, G11, G14]]

114 | gmod : Bits8 -> G -> Bits8

115 | gmod x G1 = x

116 | gmod x G2 = (if x < 128 then id else (xor 0x1b)) $ shiftL x 1

117 | gmod x G3 = xor (gmod x G2) x

119 | inv_gmod : Bits8 -> G' -> Bits8

120 | inv_gmod x g =

121 |   index_bits8 x $ case g of

122 |     G9 => gbox_9

123 |     G11 => gbox_11

124 |     G13 => gbox_13

125 |     G14 => gbox_14

127 | mix_columns : Vect n (Vect 4 Bits8) -> Vect n (Vect 4 Bits8)

128 | mix_columns xs = matmul (\a, b => foldl xor 0 $ zipWith gmod a b) xs mix_columns_mat

130 | inv_mix_columns : Vect n (Vect 4 Bits8) -> Vect n (Vect 4 Bits8)

131 | inv_mix_columns xs = matmul (\a, b => foldl xor 0 $ zipWith inv_gmod a b) xs inv_mix_columns_mat

133 | mk_round_keys : {n_k, n_b : _} -> {auto 0 ok : NonZero n_b} -> {auto 0 ok2 : NonZero n_k} -> (n_main_rounds : Nat) -> (key : Vect n_k (Vect n_b Bits8)) -> (Vect 4 (Vect n_b Bits8), Vect n_main_rounds (Vect 4 (Vect n_b Bits8)), Vect 4 (Vect n_b Bits8))

134 | mk_round_keys n_main_rounds k =

135 |   let

136 |     (rk_init :: rks) = Stream.take (S (S n_main_rounds)) $ chunk _ $ prepend (toList k) $ expand_key k

137 |   in

138 |     (rk_init, unsnoc rks)

140 | shift_rows' : {m : _} -> Nat -> Vect n (Vect (S m) Bits8) -> Vect n (Vect (S m) Bits8)

141 | shift_rows' n [] = []

142 | shift_rows' n (x :: xs) = rot_word n x :: shift_rows' (S n) xs

144 | shift_rows : {n, m : Nat} -> Vect (S n) (Vect m Bits8) -> Vect (S n) (Vect m Bits8)

145 | shift_rows = transpose . shift_rows' Z . transpose

147 | inv_shift_rows' : {m : _} -> Nat -> Vect n (Vect (S m) Bits8) -> Vect n (Vect (S m) Bits8)

148 | inv_shift_rows' n [] = []

149 | inv_shift_rows' n (x :: xs) = inv_rot_word n x :: inv_shift_rows' (S n) xs

151 | inv_shift_rows : {n, m : Nat} -> Vect (S n) (Vect m Bits8) -> Vect (S n) (Vect m Bits8)

152 | inv_shift_rows = transpose . inv_shift_rows' Z . transpose

154 | encrypt_block' : {n_k : _} -> {auto 0 ok : NonZero n_k} -> (n_main_rounds : Nat) -> (key : Vect n_k (Vect 4 Bits8)) -> (state : Vect 4 (Vect 4 Bits8)) -> Vect 4 (Vect 4 Bits8)

155 | encrypt_block' n_main_rounds k x =

156 |   let

157 |     (init_rk, main_rks, final_rk) = mk_round_keys n_main_rounds k

158 |     -- initial round

159 |     x = matxor x init_rk

160 |     -- main rounds

161 |     x = foldl (flip $ \rk => matxor rk . mix_columns . shift_rows . sub_bytes) x main_rks

162 |     -- final round

163 |     x = matxor final_rk $ shift_rows $ sub_bytes x

164 |   in

165 |     x

168 | encrypt_block : (mode : Mode) -> (key : Vect ((get_n_k mode) * 4) Bits8) -> (state : Vect 16 Bits8) -> Vect 16 Bits8

169 | encrypt_block mode k s = concat $ encrypt_block' {ok = n_k_never_Z mode} (get_main_rounds mode) (group _ _ k) (group _ _ s)

171 | decrypt_block' : {n_k : _} -> {auto 0 ok : NonZero n_k} -> (n_main_rounds : Nat) -> (key : Vect n_k (Vect 4 Bits8)) -> (state : Vect 4 (Vect 4 Bits8)) -> Vect 4 (Vect 4 Bits8)

172 | decrypt_block' n_main_rounds k x =

173 |   let

174 |     (init_rk, main_rks, final_rk) = mk_round_keys n_main_rounds k

175 |     -- initial round

176 |     x = matxor x init_rk

177 |     -- main rounds

178 |     x = foldl (flip $ \rk => inv_mix_columns . matxor rk . inv_sub_bytes . inv_shift_rows) x main_rks

179 |     -- final round

180 |     x = matxor final_rk $ inv_sub_bytes $ inv_shift_rows x

181 |   in

182 |     x

185 | decrypt_block : (mode : Mode) -> (key : Vect ((get_n_k mode) * 4) Bits8) -> (state : Vect 16 Bits8) -> Vect 16 Bits8

186 | decrypt_block mode k x = concat $ decrypt_block' {ok = n_k_never_Z mode} (get_main_rounds mode) (group _ _ k) (group _ _ x)