/* SensorBoard Secure Library
 *  Subject: AES implementation from  FIPS 197
 *    (http://csrc.nist.gov/publications/fips/fips197/fips-197.pdf)
 *  Copyright (C) 2015, T.Kurosawa
 */
#include <string.h>
#include "aes128.h"

#include "aes128_sbox.h"

static int w[60];    // FIPS 197 P.19 5.2 Key Expansion

const int Nb = 4;

static const int nround         = 10; //aes_keylen_i + 6;

// local functions ...

// FIPS 197  P.20 Figure 11
// FIPS 197  P.19  5.2
static int SubWord(uint in)
{
  return (((Sbox[(in      ) & 0xff])      ) |
          ((Sbox[(in >>  8) & 0xff]) <<  8) |
          ((Sbox[(in >> 16) & 0xff]) << 16) |
          ((Sbox[(in >> 24) & 0xff]) << 24) );
}
// FIPS 197  P.16 Figure 6
static void SubBytes(uint data[])
{
  int i;
  for (i = 0; i< Nb; i++)
    data[i] = SubWord(data[i]);
}
// FIPS 197  P.20 Figure 11
// FIPS 197  P.19  5.2
static int invSubWord(uint in)
{
  return (((invSbox[(in      ) & 0xff])      ) |
          ((invSbox[(in >>  8) & 0xff]) <<  8) |
          ((invSbox[(in >> 16) & 0xff]) << 16) |
          ((invSbox[(in >> 24) & 0xff]) << 24) );
}
// FIPS 197  P.22 5.3.2
static void invSubBytes(uint data[])
{
  int i;
  for (i = 0; i< Nb; i++)
    data[i] = invSubWord(data[i]);
}
// FIPS 197  P.20 Figure 11
// FIPS 197  P.19  5.2
static int RotWord(int in)
{
  return (((in & 0xff000000)>> 8) |
          ((in & 0x00ff0000)>> 8) |
          ((in & 0x0000ff00)>> 8) |
          ((in & 0x000000ff)<<24) );
}
// FIPS 197  P.20 Figure 11
// FIPS 197  P.19  5.2
static void KeyExpansion(void *key)
{
  // FIPS 197  P.27 Appendix A.1 
  int Rcon[10] = {0x01,0x02,0x04,0x08,0x10,0x20,0x40,0x80,0x1b,0x36};
  int i, temp;

  memcpy(w, key, aes_keylen);
  for (i = aes_keylen_i; i< 4*(nround+1); i++) {
    temp = w[i-1];
    if ((i % aes_keylen_i) == 0)
      temp = SubWord(RotWord(temp)) ^ Rcon[(i/aes_keylen_i)-1];
    w[i] = w[i-aes_keylen_i] ^ temp;
  }
}
// FIPS 197 P.10 4.2 poly-multiply
static int mul(int dt, int n)
{
  int i, x = 0;
  for (i = 8; i >0; i>>=1) {
    x <<= 1;
    if (x & 0x100) x = (x ^ 0x1b) & 0xff;
    if (n & i) x ^= dt;
  }
  return x;
}
// FIPS 197  P.18 Figure 9
static void MixColumns(uint data[])
{
  int i;
  uint8 *p = (uint8 *) data;
  for (i = 0; i< Nb; i++) {
    int idx = i<<2, x;
    x  =  mul(p[idx+0], 2) ^ mul(p[idx+1], 3) ^ 
          mul(p[idx+2], 1) ^ mul(p[idx+3], 1);
    x |= (mul(p[idx+1], 2) ^ mul(p[idx+2], 3) ^
          mul(p[idx+3], 1) ^ mul(p[idx+0], 1)) << 8;
    x |= (mul(p[idx+2], 2) ^ mul(p[idx+3], 3) ^
          mul(p[idx+0], 1) ^ mul(p[idx+1], 1)) << 16;
    x |= (mul(p[idx+3], 2) ^ mul(p[idx+0], 3) ^
          mul(p[idx+1], 1) ^ mul(p[idx+2], 1)) << 24;
    data[i] = x;
  } 
}
// FIPS 197  P.23 5.3.3
static void invMixColumns(uint data[], int n)
{
  int i;
  uint8 *p = (uint8 *) data;
  for (i = 0; i< Nb; i++) {
    int idx = i<<2, x;
    x  =  mul(p[idx+0], 14) ^ mul(p[idx+1], 11) ^
          mul(p[idx+2], 13) ^ mul(p[idx+3], 9);
    x |= (mul(p[idx+1], 14) ^ mul(p[idx+2], 11) ^
          mul(p[idx+3], 13) ^ mul(p[idx+0], 9)) << 8;
    x |= (mul(p[idx+2], 14) ^ mul(p[idx+3], 11) ^
          mul(p[idx+0], 13) ^ mul(p[idx+1], 9)) << 16;
    x |= (mul(p[idx+3], 14) ^ mul(p[idx+0], 11) ^
          mul(p[idx+1], 13) ^ mul(p[idx+2], 9)) << 24;
    data[i] = x;
  } 
}
// FIPS 197  P.17 Figure 8
static void ShiftRows(uint data[])
{
  int i;
  uint8  *p=(uint8 *)data;
  uint8  temp[aes_datalen];

  for (i = 0; i< aes_datalen; i++) {
    int r = (i/4), c = i & 3;
    int from = (r*4+(c*5)) & 15;
    temp[i] = p[from];
  }
  memcpy(p, temp, aes_datalen);
}
// FIPS 197  P.22 Figure 13
static void invShiftRows(uint data[])
{
  int i;
  uint8  *p=(uint8 *)data;
  uint8  temp[aes_datalen];

  for (i = 0; i< aes_datalen; i++) {
    int r = (i/4), c = i & 3;
    int to = (r*4+(c*5)) & 15;
    temp[to] = p[i];
  }
  memcpy(p, temp, aes_datalen);
}
// FIPS 197  P.19 Figure 10
static void AddRoundKey(uint data[], int n)
{
  int i;
  for (i = 0; i< aes_datalen_i; i++) {
    data[i] ^= w[i+n*Nb];
  }
}

////// exported function
// encrypt()
int AES128_crypt(uint data[], uint8  *key)
{
  int i;

  KeyExpansion(key);

  AddRoundKey(data, 0);

  for (i = 1; i< nround; i++) {
    SubBytes(data);
    ShiftRows(data);
    MixColumns(data);
    AddRoundKey(data, i);
  }

  SubBytes(data);
  ShiftRows(data);
  AddRoundKey(data, i);
  return i;
}

////// exported function
// decrypt()
int AES128_decrypt(uint data[], uint8  *key)
{
  int i;

  AddRoundKey(data, nround);

  for (i = nround-1; i >0; i--) {
    invShiftRows(data);
    invSubBytes(data);
    AddRoundKey(data, i);
    invMixColumns(data, i);
  }

  invShiftRows(data);
  invSubBytes(data);
  AddRoundKey(data, 0);
  return nround;
}

// O.M.A.J.I.N.A.I.
// vi:expandtab:foldmethod=syntax sw=2 ts=2