! Copyright (C) 2013 Gabriel Kerneis
! See http://factorcode.org/license.txt for BSD license.
USING: kernel sequences grouping math math.order math.parser math.bitwise
math.bits strings locals arrays make crypto.aes tools.test ;
IN: fips197

! Words are represented as arrays of 4 bytes.
! A representation based on uint32 would certainly be more
! efficient, but [shift-rows] would need significant changes.

: sub-word ( word -- word )
    [ sbox nth ] map ;

: rot-word ( word n -- word )
     cut-slice prepend ;

: xor-word ( word word -- word )
    [ bitxor ] 2map ;

: key-step ( key rcon -- next-key )
    { 0 0 0 } swap prefix
    over last 1 rot-word sub-word
    xor-word
    [ xor-word ] accumulate swap suffix 1 tail ;

: expand-key ( key -- round-keys )
    ! XXX rcon is broken in crypto.aes, missing last round;
    ! let's build directly the values we need here
    11 0 <array> 0x01 [ drop xtime ] accumulate nip swap
    [ key-step ] accumulate nip ;

: sub-bytes ( state -- state )
    [ sub-word ] map ;

: shift-rows ( state -- state )
    flip 4 iota [ rot-word ] 2map flip ;

! See FIPS 4.2.1 for the details -- this is very inefficient,
! the same operations are performed many times.  Fast exponentiation
! should be doable with accumulate or produce, instead of map-index.

: nxtime ( x y -- x' )
    make-bits
    [ [ xtime ] times 0 ? ] with map-index
    0 [ bitxor ] reduce ;

: word-product ( word word -- byte )
    [ nxtime ] 2map
    0 [ bitxor ] reduce ;

: matrix-product ( word matrix -- word )
    [ word-product ] with map ;

: mix-column ( word -- word )
    { { 2 3 1 1 }
      { 1 2 3 1 }
      { 1 1 2 3 }
      { 3 1 1 2 } } matrix-product ;

: mix-columns ( state -- state )
    [ mix-column ] map ;

: add-round-key ( round-key state -- state )
    [ xor-word ] 2map ;

: aes-round ( round-key state -- state )
   sub-bytes shift-rows mix-columns add-round-key ;


: aes-128-encrypt ( expanded-key block -- block )
    [ unclip ] dip add-round-key
    [ unclip-last swap ] dip
    [ swap aes-round ] reduce
    sub-bytes shift-rows add-round-key ;

! Inverse transformations for decrypt

: inv-shift-rows ( state -- state )
    flip { 0 3 2 1 } [ rot-word ] 2map flip ;

: inv-sub-bytes ( state -- state )
    [ [ inv-sbox nth ] map ] map ;

: inv-mix-column ( word -- word )
    { { 0xe 0xb 0xd 0x9 }
      { 0x9 0xe 0xb 0xd }
      { 0xd 0x9 0xe 0xb }
      { 0xb 0xd 0x9 0xe } } matrix-product ;

: inv-mix-columns ( state -- state )
    [ inv-mix-column ] map ;

: inv-aes-round ( round-key state -- state )
   inv-shift-rows inv-sub-bytes add-round-key inv-mix-columns ;

: aes-128-decrypt ( expanded-key block -- block )
    [ reverse ] dip
    [ unclip ] dip add-round-key
    [ unclip-last swap ] dip
    [ swap inv-aes-round ] reduce
    inv-shift-rows inv-sub-bytes add-round-key ;

! Example use: decrypt some AES ECB string
! XXX Do not use for real-world code; ECB mode is NOT secure.

: cut-padding ( seq -- seq )
    dup last head* ;

: aes-128-ecb-decrypt ( key cipher -- cleartext )
    [ 4 group expand-key ] dip
    2 [ 4 group ] times
    [ aes-128-decrypt ] with map
    concat concat
    cut-padding ;

! I also have some unit tests, not shown here.