import Base
import ../../../src/crypto/aes/types.bend as T
import ./poly.bend as P

# Executable specification of AES (FIPS 197), transcribed from the standard.
# Only the neutral Quad/State representation is shared with the
# implementation. A byte is a U32; the finite-field operations read its low
# eight bits b7..b0 (section 3.2), the bytes of the API are all below 256.

# Section 3.2/4: the byte b7 b6 ... b0 is the polynomial b7 x^7 + ... + b0.
def bits(x: U32) -> List<&2, Bool>:
  match x:
    case U32{WCon{b0, WCon{b1, WCon{b2, WCon{b3, WCon{b4, WCon{b5, WCon{b6, WCon{b7, rest}}}}}}}}}:
      [b0, b1, b2, b3, b4, b5, b6, b7]

def byte(w: Word(8n)) -> U32:
  match w:
    case WCon{b0, WCon{b1, WCon{b2, WCon{b3, WCon{b4, WCon{b5, WCon{b6, WCon{b7, WNil{}}}}}}}}}:
      U32{WCon{b0, WCon{b1, WCon{b2, WCon{b3, WCon{b4, WCon{b5, WCon{b6, WCon{b7, Word.zero(24n)}}}}}}}}}

# Section 4.2: m(x) = x^8 + x^4 + x^3 + x + 1, so x^8 = x^4 + x^3 + x + 1.
def low() -> Word(8n):
  WCon{True{}, WCon{True{}, WCon{False{}, WCon{True{}, WCon{True{}, WCon{False{}, WCon{False{}, WCon{False{}, WNil{}}}}}}}}}

# Section 4.2: the product of two bytes is their polynomial product modulo m(x).
def mul(+a: U32, +b: U32) -> U32:
  byte(P.reduce(8n, low(), P.mul(bits(a), bits(b))))

# Section 4.4 / 5.1.1: the multiplicative inverse of a nonzero a in GF(2^8),
# {00} mapped to itself. The multiplicative group of GF(2^8) has 255
# elements, so a^255 = {01} and a^-1 = a^254 for a != {00} (and
# {00}^254 = {00}): computed by squaring, a^254 = a^2 a^4 a^8 ... a^128.
# (tools/check_aes.py checks this against the defining property a * b = {01}
# for all 256 bytes.)
def inverse(+a: U32) -> U32:
  +a2 = mul(a, a)
  +a4 = mul(a2, a2)
  +a8 = mul(a4, a4)
  +a16 = mul(a8, a8)
  +a32 = mul(a16, a16)
  +a64 = mul(a32, a32)
  +a128 = mul(a64, a64)
  mul(a2, mul(a4, mul(a8, mul(a16, mul(a32, mul(a64, a128))))))

def nth_bit(xs: List<&2, Bool>, n: Nat) -> Bool:
  match xs n:
    case Nil{} _:
      False{}
    case x <> rest 0n:
      x
    case x <> rest 1n+p:
      nth_bit(rest, p)

# Section 5.1.1, equation (5.1): b'_i = b_i + b_(i+4 mod 8) + b_(i+5 mod 8)
#   + b_(i+6 mod 8) + b_(i+7 mod 8) + c_i, with c = {63}.
def affine_bit(+b: List<&2, Bool>, +c: List<&2, Bool>, +i: Nat) -> Bool:
  Bool.xor(Bool.xor(Bool.xor(Bool.xor(Bool.xor(nth_bit(b, i),
    nth_bit(b, Nat.mod(Nat.add(i, 4n), 8n))), nth_bit(b, Nat.mod(Nat.add(i, 5n), 8n))),
    nth_bit(b, Nat.mod(Nat.add(i, 6n), 8n))), nth_bit(b, Nat.mod(Nat.add(i, 7n), 8n))),
    nth_bit(c, i))

def affine(+x: U32) -> U32:
  +b = bits(x)
  +c = bits(99)
  byte(WCon{affine_bit(b, c, 0n), WCon{affine_bit(b, c, 1n), WCon{affine_bit(b, c, 2n),
    WCon{affine_bit(b, c, 3n), WCon{affine_bit(b, c, 4n), WCon{affine_bit(b, c, 5n),
    WCon{affine_bit(b, c, 6n), WCon{affine_bit(b, c, 7n), WNil{}}}}}}}}})

# Section 5.1.1: the S-box is the inverse followed by the affine transformation.
def sbox_byte(+x: U32) -> U32:
  affine(inverse(x))

# The byte c7..c0 (the S-box reads only the low eight bits of a U32).
def byte8(c0: Bool, c1: Bool, c2: Bool, c3: Bool, c4: Bool, c5: Bool, c6: Bool, c7: Bool) -> U32:
  U32{WCon{c0, WCon{c1, WCon{c2, WCon{c3, WCon{c4, WCon{c5, WCon{c6, WCon{c7, Word.zero(24n)}}}}}}}}}

# Both arms are the S-box of the byte: the match on bit 7 only keeps a
# proof about a byte whose bits are unknown from expanding the field
# arithmetic (the specification is not run in constant time).
def sbox_at(+c0: Bool, +c1: Bool, +c2: Bool, +c3: Bool, +c4: Bool, +c5: Bool, +c6: Bool, c7: Bool) -> U32:
  match c7:
    case False{}: sbox_byte(byte8(c0, c1, c2, c3, c4, c5, c6, False{}))
    case True{}: sbox_byte(byte8(c0, c1, c2, c3, c4, c5, c6, True{}))

def sbox(x: U32) -> U32:
  match x:
    case U32{WCon{c0, WCon{c1, WCon{c2, WCon{c3, WCon{c4, WCon{c5, WCon{c6, WCon{c7, rest}}}}}}}}}:
      sbox_at(c0, c1, c2, c3, c4, c5, c6, c7)

# s[r,c]: byte r of column c.
def byte_at(w: T.Quad, r: Nat) -> U32:
  match w r:
    case T.W{a, b, c, d} 0n:
      a
    case T.W{a, b, c, d} 1n:
      b
    case T.W{a, b, c, d} 2n:
      c
    case T.W{a, b, c, d} 3n+p:
      d

def column(s: T.State, c: Nat) -> T.Quad:
  match s c:
    case T.S{c0, c1, c2, c3} 0n:
      c0
    case T.S{c0, c1, c2, c3} 1n:
      c1
    case T.S{c0, c1, c2, c3} 2n:
      c2
    case T.S{c0, c1, c2, c3} 3n+p:
      c3

def at(+s: T.State, +r: Nat, +c: Nat) -> U32:
  byte_at(column(s, c), r)

# Section 5.1.1: SubBytes applies the S-box to every byte of the state.
def sub_word(w: T.Quad) -> T.Quad:
  match w:
    case T.W{a, b, c, d}:
      T.W{sbox(a), sbox(b), sbox(c), sbox(d)}

def sub_bytes(s: T.State) -> T.State:
  match s:
    case T.S{c0, c1, c2, c3}:
      T.S{sub_word(c0), sub_word(c1), sub_word(c2), sub_word(c3)}

# Section 5.1.2: s'[r,c] = s[r, (c + shift(r, 4)) mod 4], shift(r, 4) = r.
def shifted(+s: T.State, +c: Nat) -> T.Quad:
  T.W{at(s, 0n, c), at(s, 1n, Nat.mod(Nat.add(c, 1n), 4n)),
    at(s, 2n, Nat.mod(Nat.add(c, 2n), 4n)), at(s, 3n, Nat.mod(Nat.add(c, 3n), 4n))}

def shift_rows(+s: T.State) -> T.State:
  T.S{shifted(s, 0n), shifted(s, 1n), shifted(s, 2n), shifted(s, 3n)}

# Section 5.1.3, equation (5.6).
def mix_column(w: T.Quad) -> T.Quad:
  match w:
    case T.W{+a, +b, +c, +d}:
      T.W{U32.xor(U32.xor(U32.xor(mul(2, a), mul(3, b)), c), d),
        U32.xor(U32.xor(U32.xor(a, mul(2, b)), mul(3, c)), d),
        U32.xor(U32.xor(U32.xor(a, b), mul(2, c)), mul(3, d)),
        U32.xor(U32.xor(U32.xor(mul(3, a), b), c), mul(2, d))}

def mix_columns(s: T.State) -> T.State:
  match s:
    case T.S{c0, c1, c2, c3}:
      T.S{mix_column(c0), mix_column(c1), mix_column(c2), mix_column(c3)}

def xor_word(a: T.Quad, b: T.Quad) -> T.Quad:
  match a b:
    case T.W{a0, a1, a2, a3} T.W{b0, b1, b2, b3}:
      T.W{U32.xor(a0, b0), U32.xor(a1, b1), U32.xor(a2, b2), U32.xor(a3, b3)}

# Section 5.1.4: column c is added to the key word w[round*4 + c].
def add_round_key(s: T.State, w0: T.Quad, w1: T.Quad, w2: T.Quad, w3: T.Quad) -> T.State:
  match s:
    case T.S{c0, c1, c2, c3}:
      T.S{xor_word(c0, w0), xor_word(c1, w1), xor_word(c2, w2), xor_word(c3, w3)}

# Section 5.1, Figure 5. After the initial AddRoundKey, n more full rounds
# (SubBytes, ShiftRows, MixColumns, AddRoundKey) and the final round (no
# MixColumns), each taking the next four words of the key schedule.
def rounds(n: Nat, s: T.State, w: List<&2, T.Quad>) -> T.State:
  match n w:
    case 0n w0 <> w1 <> w2 <> w3 <> rest:
      add_round_key(shift_rows(sub_bytes(s)), w0, w1, w2, w3)
    case 1n+p w0 <> w1 <> w2 <> w3 <> rest:
      rounds(p, add_round_key(mix_columns(shift_rows(sub_bytes(s))), w0, w1, w2, w3), rest)
    case _ _:
      s

def cipher(+nr: Nat, s: T.State, w: List<&2, T.Quad>) -> T.State:
  match w:
    case w0 <> w1 <> w2 <> w3 <> rest:
      rounds(Nat.sub(nr, 1n), add_round_key(s, w0, w1, w2, w3), rest)
    case _:
      s

# Section 5.2: RotWord, SubWord (sub_word above), and Rcon[j] = [x^(j-1), 0, 0, 0].
def rot_word(w: T.Quad) -> T.Quad:
  match w:
    case T.W{a, b, c, d}:
      T.W{b, c, d, a}

def xpow(n: Nat) -> U32:
  match n:
    case 0n:
      1
    case 1n+p:
      mul(2, xpow(p))

def rcon(j: Nat) -> T.Quad:
  T.W{xpow(Nat.sub(j, 1n)), 0, 0, 0}

# w[i] = [key[4i], key[4i+1], key[4i+2], key[4i+3]] for i < Nk.
def key_words(key: List<&2, U32>) -> List<&2, T.Quad>:
  match key:
    case a <> b <> c <> d <> rest:
      T.W{a, b, c, d} <> key_words(rest)
    case _:
      Nil{}

def nth_word(ws: List<&2, T.Quad>, n: Nat) -> T.Quad:
  match ws n:
    case Nil{} _:
      T.W{0, 0, 0, 0}
    case w <> rest 0n:
      w
    case w <> rest 1n+p:
      nth_word(rest, p)

# The transformation of temp = w[i-1] in Figure 11, by k = i mod Nk,
# big = (Nk > 6) and q = i / Nk.
def transform(k: Nat, big: Bool, +q: Nat, temp: T.Quad) -> T.Quad:
  match k:
    case 0n:
      xor_word(sub_word(rot_word(temp)), rcon(q))
    case 1n:
      temp
    case 2n:
      temp
    case 3n:
      temp
    case 4n:
      match big:
        case True{}:
          sub_word(temp)
        case False{}:
          temp
    case 5n+p:
      temp

# w[i] = w[i-Nk] xor transform(w[i-1]); history holds w[i-1], w[i-2], ...
def next_word(+nk: Nat, +i: Nat, +history: List<&2, T.Quad>) -> T.Quad:
  xor_word(nth_word(history, Nat.sub(nk, 1n)),
    transform(Nat.mod(i, nk), Nat.is_lt(6n, nk), Nat.div(i, nk), nth_word(history, 0n)))

def expand(n: Nat, +nk: Nat, +i: Nat, +history: List<&2, T.Quad>) -> List<&2, T.Quad>:
  match n:
    case 0n:
      history
    case 1n+p:
      expand(p, nk, 1n+i, next_word(nk, i, history) <> history)

# Figure 11: Nb*(Nr+1) = 4*(Nr+1) words, the first Nk from the key.
def key_expansion(+nk: Nat, +nr: Nat, key: List<&2, U32>) -> List<&2, T.Quad>:
  List.reverse(&2, T.Quad, expand(Nat.sub(Nat.mul(4n, 1n+nr), nk), nk, nk,
    List.reverse(&2, T.Quad, key_words(key))))

def encrypt(+nk: Nat, +nr: Nat, key: List<&2, U32>, block: T.State) -> T.State:
  cipher(nr, block, key_expansion(nk, nr, key))

# Figure 4: Nk is the key length in 32-bit words and Nr = Nk + 6
# (AES-128: 4, 10; AES-192: 6, 12; AES-256: 8, 14).
def nk(+key: List<&2, U32>) -> Nat:
  Nat.div(List.length(&2, U32, key), 4n)

def nr(+key: List<&2, U32>) -> Nat:
  Nat.add(nk(key), 6n)

# AES with the key's own Nk and Nr.
def aes(+key: List<&2, U32>, block: T.State) -> T.State:
  encrypt(nk(key), nr(key), key, block)

def aes128(key: List<&2, U32>, block: T.State) -> T.State:
  encrypt(4n, 10n, key, block)

def aes256(key: List<&2, U32>, block: T.State) -> T.State:
  encrypt(8n, 14n, key, block)

# Section 3.4: s[r,c] = in[r + 4c], out[r + 4c] = s[r,c]. Missing input
# bytes read as {00}.
def nth_byte(xs: List<&2, U32>, n: Nat) -> U32:
  match xs n:
    case Nil{} _:
      0
    case x <> rest 0n:
      x
    case x <> rest 1n+p:
      nth_byte(rest, p)

def state_of(+input: List<&2, U32>) -> T.State:
  T.S{T.W{nth_byte(input, 0n), nth_byte(input, 1n), nth_byte(input, 2n), nth_byte(input, 3n)},
    T.W{nth_byte(input, 4n), nth_byte(input, 5n), nth_byte(input, 6n), nth_byte(input, 7n)},
    T.W{nth_byte(input, 8n), nth_byte(input, 9n), nth_byte(input, 10n), nth_byte(input, 11n)},
    T.W{nth_byte(input, 12n), nth_byte(input, 13n), nth_byte(input, 14n), nth_byte(input, 15n)}}

def bytes_of(s: T.State) -> List<&2, U32>:
  match s:
    case T.S{T.W{a0, a1, a2, a3}, T.W{b0, b1, b2, b3}, T.W{c0, c1, c2, c3}, T.W{d0, d1, d2, d3}}:
      [a0, a1, a2, a3, b0, b1, b2, b3, c0, c1, c2, c3, d0, d1, d2, d3]