extra/crypto/aes/aes.factor

   1 ! Copyright (C) 2013 Fred Alger
   2 ! Some parts Copyright (C) 2008 Doug Coleman.
   3 ! See http://factorcode.org/license.txt for BSD license.
   4 USING: accessors arrays assocs combinators crypto.aes.utils
   5 generalizations grouping kernel locals math math.bitwise
   6 math.ranges memoize namespaces sequences sequences.private
   7 sequences.unrolled ;
   8 IN: crypto.aes
   9
  10 CONSTANT: AES_BLOCK_SIZE 16
  11 #! FIPS-197 AES
  12 #! input block, state, output block -- 4 32-bit words
  13 CONSTANT: FIPS-197 {
  14     { 128 10 } ! aes-128 -- Key(4) Block(4) Rounds(10)
  15     { 192 12 } ! aes-192 -- Key(6) Block(4) Rounds(12)
  16     { 256 14 } ! aes-256 -- Key(8) Block(4) Rounds(14)
  17 }
  18
  19 <PRIVATE
  20
  21 : (nrounds) ( byte-array -- rounds )
  22     length 8 * FIPS-197 at ;
  23
  24 : sbox ( -- array )
  25 {
  26     0x63 0x7c 0x77 0x7b 0xf2 0x6b 0x6f 0xc5
  27     0x30 0x01 0x67 0x2b 0xfe 0xd7 0xab 0x76
  28     0xca 0x82 0xc9 0x7d 0xfa 0x59 0x47 0xf0
  29     0xad 0xd4 0xa2 0xaf 0x9c 0xa4 0x72 0xc0
  30     0xb7 0xfd 0x93 0x26 0x36 0x3f 0xf7 0xcc
  31     0x34 0xa5 0xe5 0xf1 0x71 0xd8 0x31 0x15
  32     0x04 0xc7 0x23 0xc3 0x18 0x96 0x05 0x9a
  33     0x07 0x12 0x80 0xe2 0xeb 0x27 0xb2 0x75
  34     0x09 0x83 0x2c 0x1a 0x1b 0x6e 0x5a 0xa0
  35     0x52 0x3b 0xd6 0xb3 0x29 0xe3 0x2f 0x84
  36     0x53 0xd1 0x00 0xed 0x20 0xfc 0xb1 0x5b
  37     0x6a 0xcb 0xbe 0x39 0x4a 0x4c 0x58 0xcf
  38     0xd0 0xef 0xaa 0xfb 0x43 0x4d 0x33 0x85
  39     0x45 0xf9 0x02 0x7f 0x50 0x3c 0x9f 0xa8
  40     0x51 0xa3 0x40 0x8f 0x92 0x9d 0x38 0xf5
  41     0xbc 0xb6 0xda 0x21 0x10 0xff 0xf3 0xd2
  42     0xcd 0x0c 0x13 0xec 0x5f 0x97 0x44 0x17
  43     0xc4 0xa7 0x7e 0x3d 0x64 0x5d 0x19 0x73
  44     0x60 0x81 0x4f 0xdc 0x22 0x2a 0x90 0x88
  45     0x46 0xee 0xb8 0x14 0xde 0x5e 0x0b 0xdb
  46     0xe0 0x32 0x3a 0x0a 0x49 0x06 0x24 0x5c
  47     0xc2 0xd3 0xac 0x62 0x91 0x95 0xe4 0x79
  48     0xe7 0xc8 0x37 0x6d 0x8d 0xd5 0x4e 0xa9
  49     0x6c 0x56 0xf4 0xea 0x65 0x7a 0xae 0x08
  50     0xba 0x78 0x25 0x2e 0x1c 0xa6 0xb4 0xc6
  51     0xe8 0xdd 0x74 0x1f 0x4b 0xbd 0x8b 0x8a
  52     0x70 0x3e 0xb5 0x66 0x48 0x03 0xf6 0x0e
  53     0x61 0x35 0x57 0xb9 0x86 0xc1 0x1d 0x9e
  54     0xe1 0xf8 0x98 0x11 0x69 0xd9 0x8e 0x94
  55     0x9b 0x1e 0x87 0xe9 0xce 0x55 0x28 0xdf
  56     0x8c 0xa1 0x89 0x0d 0xbf 0xe6 0x42 0x68
  57     0x41 0x99 0x2d 0x0f 0xb0 0x54 0xbb 0x16
  58 } ;
  59
  60 : inv-sbox ( -- array )
  61     256 0 <array>
  62     dup 256 [ dup sbox nth rot set-nth ] with each-integer ;
  63
  64 #! applies sbox to each byte of word
  65 : subword ( word -- word' )
  66     [ gb0 sbox nth ] keep [ gb1 sbox nth ] keep
  67     [ gb2 sbox nth ] keep gb3 sbox nth >ui32 ;
  68
  69 #! applies inverse sbox to each byte of word
  70 : inv-subword ( word -- word' )
  71     [ gb0 inv-sbox nth ] keep [ gb1 inv-sbox nth ] keep
  72     [ gb2 inv-sbox nth ] keep gb3 inv-sbox nth >ui32 ;
  73
  74 : rotword ( n -- n ) 8 bitroll-32 ;
  75
  76 #! round constants, 2^n over GF(2^8)
  77 : rcon ( -- array )
  78     {
  79         0x00 0x01 0x02 0x04 0x08 0x10
  80         0x20 0x40 0x80 0x1b 0x36
  81     } ;
  82
  83 : (rcon-nth) ( n -- rcon[n] ) rcon nth 24 shift ;
  84
  85 #! Galois field product related
  86 : xtime ( x -- x' )
  87     [ 1 shift ]
  88     [ 0x80 bitand 0 = 0 0x1b ? ] bi bitxor 8 bits ;
  89
  90 #! generate t-box
  91 :: set-t ( T i -- )
  92     i sbox nth :> a1
  93     a1 xtime :> a2
  94     a1 a2 bitxor :> a3
  95
  96     a3 a1 a1 a2 >ui32 i T set-nth
  97     a1 a1 a2 a3 >ui32 i 0x100 + T set-nth
  98     a1 a2 a3 a1 >ui32 i 0x200 + T set-nth
  99     a2 a3 a1 a1 >ui32 i 0x300 + T set-nth ;
 100
 101 MEMO:: t-table ( -- array )
 102     1024 0 <array>
 103     dup 256 [ set-t ] with each-integer ;
 104
 105 #! generate inverse t-box
 106 :: set-d ( D i -- )
 107     i inv-sbox nth :> a1
 108     a1 xtime :> a2
 109     a2 xtime :> a4
 110     a4 xtime :> a8
 111     a8 a1 bitxor :> a9
 112     a9 a2 bitxor :> ab
 113     a9 a4 bitxor :> ad
 114     a8 a4 a2 bitxor bitxor :> ae
 115
 116     ab ad a9 ae >ui32 i D set-nth
 117     ad a9 ae ab >ui32 i 0x100 + D set-nth
 118     a9 ae ab ad >ui32 i 0x200 + D set-nth
 119     ae ab ad a9 >ui32 i 0x300 + D set-nth ;
 120
 121 MEMO:: d-table ( -- array )
 122     1024 0 <array>
 123     dup 256 [ set-d ] with each-integer ;
 124
 125
 126 :: (transform) ( a0 a1 a2 a3 table -- word' )
 127   a0 a1 a2 a3
 128   [ 0x100 + ] [ 0x200 + ] [ 0x300 + ] tri*
 129   [ table nth ] 4 napply
 130   bitxor bitxor bitxor ; inline
 131
 132 : t-transform ( a0 a1 a2 a3 -- word' ) t-table (transform) ;
 133 : d-transform ( a0 a1 a2 a3 -- word' ) d-table (transform) ;
 134
 135 #! key schedule
 136 #! expands an 128/192/256 bit key into an 176/208/240 byte schedule
 137
 138 SYMBOL: aes-expand-inner
 139 HOOK: key-expand-round aes-expand-inner  ( temp i -- temp' )
 140
 141 SINGLETON: aes-128-key
 142 SINGLETON: aes-256-key
 143
 144 : (add-rcon) ( word rcon-ndx -- word' )
 145     (rcon-nth) [ rotword subword ] dip bitxor ;
 146
 147 M: aes-128-key key-expand-round ( temp i -- temp' )
 148     4 /mod 0 = swap and [ (add-rcon) ] when* ;
 149
 150 ERROR: aes-192-256-not-implemented ;
 151 M: aes-256-key key-expand-round ( temp i -- temp' )
 152     aes-192-256-not-implemented ;
 153
 154 : (key-sched-round) ( output temp i -- output' )
 155     key-expand-round
 156     [ dup 4th-from-end ] dip bitxor suffix! ; inline
 157
 158 : (sched-interval) ( K Nr -- seq )
 159     [ length ] dip 1 + 4 * [a,b) ;    ! over the interval Nk...Nb(Nr + 1)
 160
 161 : (init-round) ( out -- out temp quot )
 162     [ ]
 163     [ last ]
 164     [
 165         length
 166         6 > [ aes-256-key ] [ aes-128-key ] if
 167     ] tri ;
 168
 169 #! K -- input key (byte-array), Nr -- number of rounds
 170 #! output: sched, Nb(Nr+1) byte key schedule
 171 : (expand-enc-key) ( K Nr -- sched )
 172     [ bytes>words ] dip
 173     [ drop (init-round) ]
 174     [ (sched-interval) ] 2bi
 175     [
 176         [ aes-expand-inner set ] dip
 177         [ (key-sched-round) dup last ] each
 178     ] with-scope
 179     drop ;
 180
 181 TUPLE: aes-state nrounds key state ;
 182
 183 : <aes> ( nrounds key state -- aes-state ) \ aes-state boa ;
 184
 185 #! grabs the 4n...4(n+1) words of the key
 186 : (key-at-nth-round) ( nth aes -- seq )
 187     [ 4 * dup 4 + ] [ key>> ] bi* <slice> ;
 188
 189 SYMBOL: aes-strategy
 190 HOOK: (expand-key) aes-strategy ( K Nr -- sched )
 191 HOOK: (first-round) aes-strategy ( aes -- aes' )
 192 HOOK: (counter) aes-strategy ( nrounds -- seq )
 193 HOOK: (round) aes-strategy ( state -- )
 194 HOOK: (add-key) aes-strategy ( aes -- aes' )
 195 HOOK: (final-round) aes-strategy ( aes -- aes' )
 196
 197 SINGLETON: aes-decrypt
 198 SINGLETON: aes-encrypt
 199
 200
 201 #! rotates the 2nd row left by one element
 202 #! rotates the 3rd row left by two elements
 203 #! rotates the 4th row left by three elements
 204 #!
 205 #! Kind of ugly because the algorithm is specified and
 206 #! implemented in terms of columns. This approach is very
 207 #! efficient in terms of execution and only requires one new
 208 #! word to implement.
 209 #!
 210 #! The alternative is to split into arrays of bytes, transpose,
 211 #! rotate each row n times, transpose again, and then
 212 #! smash them back into 4-byte words.
 213 :: (shift-rows) ( c0 c1 c2 c3 -- c0' c1' c2' c3' )
 214     c3 gb0   c2 gb1   c1 gb2   c0 gb3   >ui32   ! c0'
 215     c0 gb0   c3 gb1   c2 gb2   c1 gb3   >ui32   ! c1'
 216     c1 gb0   c0 gb1   c3 gb2   c2 gb3   >ui32   ! c2'
 217     c2 gb0   c1 gb1   c0 gb2   c3 gb3   >ui32 ; ! c3'
 218
 219 :: (unshift-rows) ( c0 c1 c2 c3 -- c0' c1' c2' c3' )
 220     c1 gb0   c2 gb1   c3 gb2   c0 gb3   >ui32   ! c0'
 221     c2 gb0   c3 gb1   c0 gb2   c1 gb3   >ui32   ! c1'
 222     c3 gb0   c0 gb1   c1 gb2   c2 gb3   >ui32   ! c2'
 223     c0 gb0   c1 gb1   c2 gb2   c3 gb3   >ui32 ; ! c3'
 224
 225 : (add-round-key) ( key state -- state' )
 226    4 [ bitxor ] unrolled-2map ;
 227
 228 : add-round-key ( aes n -- aes' )
 229     over (key-at-nth-round) swap
 230     [ (add-round-key) ] change-state ;
 231
 232 : add-final-round-key ( aes -- aes' )
 233     dup nrounds>> add-round-key ;
 234
 235 : add-first-round-key ( aes -- aes' )
 236     0 add-round-key ;
 237
 238 : aes-round ( state -- )
 239     dup first4-unsafe
 240     { [ first-diag t-transform ]
 241       [ second-diag t-transform ]
 242       [ third-diag t-transform ]
 243       [ fourth-diag t-transform ] } 4 ncleave
 244       set-first4-unsafe ;
 245
 246
 247 : shift-rows ( state -- state' )
 248     first4 (shift-rows) 4array ;
 249
 250 : unshift-rows ( state -- state' )
 251     first4 (unshift-rows) 4array ;
 252
 253 : final-round ( state -- state' )
 254     4 [ subword ] unrolled-map shift-rows ;
 255
 256 : (do-round) ( aes -- aes' )
 257     [ state>> (round) ] keep ;
 258
 259 M: aes-encrypt (expand-key) (expand-enc-key) ;
 260 M: aes-encrypt (first-round) add-first-round-key ;
 261 M: aes-encrypt (counter) 0 swap (a,b) ;
 262 M: aes-encrypt (round) aes-round ;
 263 M: aes-encrypt (final-round) [ final-round ] change-state add-final-round-key ;
 264
 265 M:: aes-decrypt (expand-key) ( K Nr -- sched )
 266     K Nr (expand-enc-key) dup length :> key-length
 267     [
 268         [ 4 >= ] [ key-length 4 - < ] bi and
 269         [ subword ui32-rev> d-transform ]
 270         [ ] if
 271     ] map-index ;
 272
 273 M: aes-decrypt (first-round) ( aes -- aes' )
 274     add-final-round-key ;
 275
 276 M: aes-decrypt (counter) ( nrounds -- seq ) 0 swap (a,b) <reversed> ;
 277 M: aes-decrypt (final-round) ( aes -- aes' )
 278     [ [ inv-subword ] map unshift-rows  ] change-state
 279     add-first-round-key ;
 280
 281 M: aes-decrypt (round) ( state -- )
 282     dup first4-unsafe
 283     { [ -first-diag d-transform ]
 284       [ -fourth-diag d-transform ]
 285       [ -third-diag d-transform ]
 286       [ -second-diag d-transform ] } 4 ncleave
 287       set-first4-unsafe ;
 288
 289
 290 : (aes-crypt) ( aes -- aes' )
 291     (first-round) [
 292         dup nrounds>> (counter)
 293         [ [ (do-round) ] dip add-round-key drop ] with each
 294     ] keep
 295     (final-round) ;
 296
 297 : (aes-expand-key) ( key -- nrounds expanded-key )
 298     [ (nrounds) ] keep over (expand-key) ;
 299
 300 : (aes-crypt-block-inner) ( nrounds key block -- crypted-block )
 301     <aes> (aes-crypt) state>> ;
 302
 303 : (aes-crypt-block) ( key block -- output-block )
 304     [ (aes-expand-key) ] dip bytes>words (aes-crypt-block-inner) ;
 305
 306 PRIVATE>
 307
 308 : aes-encrypt-block ( key block -- output )
 309     [ aes-encrypt aes-strategy set (aes-crypt-block) ] with-scope
 310     [ ui32> 4array reverse ] map concat ;
 311
 312 : aes-decrypt-block ( key block -- output )
 313     [ aes-decrypt aes-strategy set (aes-crypt-block) ] with-scope
 314     [ ui32> 4array reverse ] map concat ;