AES.cpp

/*
 * AES.cpp - AES implementation optimized for ATTiny85.
 *
 * The MIT License (MIT)
 * Copyright © 2022 Matt Borja (https://mattborja.dev/)
 * 
 * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
 * 
 * The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
 * 
 * THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 
 *
 * Created by Matt Borja, December 27, 2016.
 * 
 * Tested OK:
 * - getColumn
 * - getRow
 * - rotWord
 * - subBytes
 * - keyExpansion
 * - shiftRows
 * 
 * Notes:
 * - Unrolling certain for() loops seems to alleviate memory corruption behavior observed (i.e. addRoundKey)
 * 
 * References:
 * - FIPS-197 (http://csrc.nist.gov/publications/fips/fips197/fips-197.pdf)
 * - Rijndael Flash Animation (http://www.formaestudio.com/rijndaelinspector/archivos/Rijndael_Animation_v4_eng.swf)
 */
#include "Arduino.h"
#include "AES.h"

AES::AES(byte *key)
{
  keyExpansion(key);
}

/*
 * getColumn
 * @param byte *result - Pointer to byte[SIZEOF_BLOCK] array containing fetched column
 * @param byte *state - Pointer to byte[N] state array
 * @param uint8_t row - 0-index of column
 */
void AES::getColumn(byte *result, byte *state, uint8_t column) {
  uint8_t offset = column * SIZEOF_BLOCK;

  result[0] = state[offset];
  result[1] = state[++offset];
  result[2] = state[++offset];
  result[3] = state[++offset];
}

/*
 * getRow
 * @param byte *result - Pointer to byte[SIZEOF_BLOCK] array containing fetched row
 * @param byte *state - Pointer to byte[N] state array
 * @param uint8_t row - 0-index of row
 */
void AES::getRow(byte *result, byte *state, uint8_t row) {
  result[0] = state[row];
  result[1] = state[row + SIZEOF_BLOCK];
  result[2] = state[row + SIZEOF_BLOCK + SIZEOF_BLOCK];
  result[3] = state[row + SIZEOF_BLOCK + SIZEOF_BLOCK + SIZEOF_BLOCK];
}

/*
 * rotWord
 * @param byte *result - Pointer to byte[SIZEOF_BLOCK] array to be rotated
 */
void AES::rotWord(byte *block) {
  byte first = block[0];
  block[0] = block[1];
  block[1] = block[2];
  block[2] = block[3];
  block[3] = first;
}

/*
 * subBytes
 * @param byte *result - Pointer to byte[len] array containing bytes to be substituted
 * @param uint8_t len - Length of byte array to be substituted
 */
void AES::subBytes(byte *result, uint8_t len)
{
  uint8_t i = 0;
  byte r, c;

  r = result[i] >> 4; c = result[i] & 0x0f; result[i++] = SBox[r * 0x10 + c];
  r = result[i] >> 4; c = result[i] & 0x0f; result[i++] = SBox[r * 0x10 + c];
  r = result[i] >> 4; c = result[i] & 0x0f; result[i++] = SBox[r * 0x10 + c];
  r = result[i] >> 4; c = result[i] & 0x0f; result[i++] = SBox[r * 0x10 + c];

  if (len == 4)
    return;

  r = result[i] >> 4; c = result[i] & 0x0f; result[i++] = SBox[r * 0x10 + c];
  r = result[i] >> 4; c = result[i] & 0x0f; result[i++] = SBox[r * 0x10 + c];
  r = result[i] >> 4; c = result[i] & 0x0f; result[i++] = SBox[r * 0x10 + c];
  r = result[i] >> 4; c = result[i] & 0x0f; result[i++] = SBox[r * 0x10 + c];
  
  r = result[i] >> 4; c = result[i] & 0x0f; result[i++] = SBox[r * 0x10 + c];
  r = result[i] >> 4; c = result[i] & 0x0f; result[i++] = SBox[r * 0x10 + c];
  r = result[i] >> 4; c = result[i] & 0x0f; result[i++] = SBox[r * 0x10 + c];
  r = result[i] >> 4; c = result[i] & 0x0f; result[i++] = SBox[r * 0x10 + c];

  r = result[i] >> 4; c = result[i] & 0x0f; result[i++] = SBox[r * 0x10 + c];
  r = result[i] >> 4; c = result[i] & 0x0f; result[i++] = SBox[r * 0x10 + c];
  r = result[i] >> 4; c = result[i] & 0x0f; result[i++] = SBox[r * 0x10 + c];
  r = result[i] >> 4; c = result[i] & 0x0f; result[i++] = SBox[r * 0x10 + c];
}

void AES::buildKeyRound(uint8_t w, uint8_t r) {
  byte block[SIZEOF_BLOCK];

  getColumn(block, _expandedKey, w - 1);      
  rotWord(block);                             
  subBytes(block, sizeof(block));             
  memcpy(_expandedKey+(w*4), block, SIZEOF_BLOCK);
  
  getColumn(block, _expandedKey, w - 4);      
  _expandedKey[w*4] ^= block[0] ^ RCon[r];  
  _expandedKey[w*4+1] ^= block[1] ^ 0x00;     
  _expandedKey[w*4+2] ^= block[2] ^ 0x00;     
  _expandedKey[w*4+3] ^= block[3] ^ 0x00;     

  w++; memcpy(_expandedKey+(w*4), _expandedKey+((w-1) * 4), SIZEOF_BLOCK);
  getColumn(block, _expandedKey, w - 4);
  _expandedKey[w*4] ^= block[0]; _expandedKey[w*4+1] ^= block[1]; _expandedKey[w*4+2] ^= block[2]; _expandedKey[w*4+3] ^= block[3];
  
  w++; memcpy(_expandedKey+(w*4), _expandedKey+((w-1) * 4), SIZEOF_BLOCK);
  getColumn(block, _expandedKey, w - 4);
  _expandedKey[w*4] ^= block[0]; _expandedKey[w*4+1] ^= block[1]; _expandedKey[w*4+2] ^= block[2]; _expandedKey[w*4+3] ^= block[3];
  
  w++; memcpy(_expandedKey+(w*4), _expandedKey+((w-1) * 4), SIZEOF_BLOCK);
  getColumn(block, _expandedKey, w - 4);
  _expandedKey[w*4] ^= block[0]; _expandedKey[w*4+1] ^= block[1]; _expandedKey[w*4+2] ^= block[2]; _expandedKey[w*4+3] ^= block[3];
}

/*
 * keyExpansion
 * @param byte *key - Pointer to byte[SIZEOF_KEY] key array
 */
void AES::keyExpansion(byte *key) {
  if (_keyExpanded)
    return;
  
  // Cipher key (OK)
  memcpy(_expandedKey, key, SIZEOF_KEY);

  // Word position (iW)
  uint8_t iW = 4;
  uint8_t r = 0;

  // Round key (rK)
  buildKeyRound(iW, r++); iW += 4;
  buildKeyRound(iW, r++); iW += 4;
  buildKeyRound(iW, r++); iW += 4;
  buildKeyRound(iW, r++); iW += 4;
  buildKeyRound(iW, r++); iW += 4;
  buildKeyRound(iW, r++); iW += 4;
  buildKeyRound(iW, r++); iW += 4;
  buildKeyRound(iW, r++); iW += 4;
  buildKeyRound(iW, r++); iW += 4;
  buildKeyRound(iW, r++); iW += 4;

  _keyExpanded = true;
}

/*
 * shiftRows
 * @param byte *state - Pointer to byte[SIZEOF_STATE] state array containing rows to be shifted
 */
void AES::shiftRows(byte *state) {
  byte block[SIZEOF_BLOCK] = { 0 };
  uint8_t r = 0;

  getRow(block, state, ++r);
  state[r] = block[1];
  state[r + SIZEOF_BLOCK] = block[2];
  state[r + SIZEOF_BLOCK + SIZEOF_BLOCK] = block[3];
  state[r + SIZEOF_BLOCK + SIZEOF_BLOCK + SIZEOF_BLOCK] = block[0];

  getRow(block, state, ++r);
  state[r] = block[2];
  state[r + SIZEOF_BLOCK] = block[3];
  state[r + SIZEOF_BLOCK + SIZEOF_BLOCK] = block[0];
  state[r + SIZEOF_BLOCK + SIZEOF_BLOCK + SIZEOF_BLOCK] = block[1];

  getRow(block, state, ++r);
  state[r] = block[3];
  state[r + SIZEOF_BLOCK] = block[0];
  state[r + SIZEOF_BLOCK + SIZEOF_BLOCK] = block[1];
  state[r + SIZEOF_BLOCK + SIZEOF_BLOCK + SIZEOF_BLOCK] = block[2];
}

/*
 * xTime
 * @description Finds the product of {02} and the argument to xtime modulo {1b}
 * @param byte x - Byte to multiply
 */
byte xTime(byte x)
{
  bool mod = (x & 0x80) == 0x80;
  x <<= 1;

  if (mod)
    x ^= 0x1b;

  return x;
}

/*
 * Multiply
 * @description Multiplies numbers in the field GF(2^8)
 */
byte gMult(byte a, byte b)
{
  byte lsb;
  byte p = 0x00;

  lsb = b & 0x01; if (lsb == 1) p ^= a; a = xTime(a); b >>= 1;
  lsb = b & 0x01; if (lsb == 1) p ^= a; a = xTime(a); b >>= 1;
  lsb = b & 0x01; if (lsb == 1) p ^= a; a = xTime(a); b >>= 1;
  lsb = b & 0x01; if (lsb == 1) p ^= a; a = xTime(a); b >>= 1;
  lsb = b & 0x01; if (lsb == 1) p ^= a; a = xTime(a); b >>= 1;
  lsb = b & 0x01; if (lsb == 1) p ^= a; a = xTime(a); b >>= 1;
  lsb = b & 0x01; if (lsb == 1) p ^= a; a = xTime(a); b >>= 1;
  lsb = b & 0x01; if (lsb == 1) p ^= a; a = xTime(a); b >>= 1;
  
  return p;
}

/*
 * mixColumns
 * @param *state - Pointer to byte[SIZEOF_STATE] state array to be mixed
 */
void AES::mixColumns(byte *state)
{
  uint8_t c = 0;
  byte col[SIZEOF_BLOCK];
  byte mix[SIZEOF_BLOCK];
  
  memcpy(col, &state[4 * c], SIZEOF_BLOCK);
  mix[0] = gMult(0x02, col[0]) ^ gMult(0x03, col[1]) ^ gMult(0x01, col[2]) ^ gMult(0x01, col[3]);
  mix[1] = gMult(0x01, col[0]) ^ gMult(0x02, col[1]) ^ gMult(0x03, col[2]) ^ gMult(0x01, col[3]);
  mix[2] = gMult(0x01, col[0]) ^ gMult(0x01, col[1]) ^ gMult(0x02, col[2]) ^ gMult(0x03, col[3]);
  mix[3] = gMult(0x03, col[0]) ^ gMult(0x01, col[1]) ^ gMult(0x01, col[2]) ^ gMult(0x02, col[3]);
  memcpy(&state[4 * c], mix, SIZEOF_BLOCK);
  c++;
  
  memcpy(col, &state[4 * c], SIZEOF_BLOCK);
  mix[0] = gMult(0x02, col[0]) ^ gMult(0x03, col[1]) ^ gMult(0x01, col[2]) ^ gMult(0x01, col[3]);
  mix[1] = gMult(0x01, col[0]) ^ gMult(0x02, col[1]) ^ gMult(0x03, col[2]) ^ gMult(0x01, col[3]);
  mix[2] = gMult(0x01, col[0]) ^ gMult(0x01, col[1]) ^ gMult(0x02, col[2]) ^ gMult(0x03, col[3]);
  mix[3] = gMult(0x03, col[0]) ^ gMult(0x01, col[1]) ^ gMult(0x01, col[2]) ^ gMult(0x02, col[3]);
  memcpy(&state[4 * c], mix, SIZEOF_BLOCK);
  c++;
  
  memcpy(col, &state[4 * c], SIZEOF_BLOCK);
  mix[0] = gMult(0x02, col[0]) ^ gMult(0x03, col[1]) ^ gMult(0x01, col[2]) ^ gMult(0x01, col[3]);
  mix[1] = gMult(0x01, col[0]) ^ gMult(0x02, col[1]) ^ gMult(0x03, col[2]) ^ gMult(0x01, col[3]);
  mix[2] = gMult(0x01, col[0]) ^ gMult(0x01, col[1]) ^ gMult(0x02, col[2]) ^ gMult(0x03, col[3]);
  mix[3] = gMult(0x03, col[0]) ^ gMult(0x01, col[1]) ^ gMult(0x01, col[2]) ^ gMult(0x02, col[3]);
  memcpy(&state[4 * c], mix, SIZEOF_BLOCK);
  c++;
  
  memcpy(col, &state[4 * c], SIZEOF_BLOCK);
  mix[0] = gMult(0x02, col[0]) ^ gMult(0x03, col[1]) ^ gMult(0x01, col[2]) ^ gMult(0x01, col[3]);
  mix[1] = gMult(0x01, col[0]) ^ gMult(0x02, col[1]) ^ gMult(0x03, col[2]) ^ gMult(0x01, col[3]);
  mix[2] = gMult(0x01, col[0]) ^ gMult(0x01, col[1]) ^ gMult(0x02, col[2]) ^ gMult(0x03, col[3]);
  mix[3] = gMult(0x03, col[0]) ^ gMult(0x01, col[1]) ^ gMult(0x01, col[2]) ^ gMult(0x02, col[3]);
  memcpy(&state[4 * c], mix, SIZEOF_BLOCK);
}

/*
 * addRoundKey
 * @param byte *state - Pointer to byte[SIZEOF_STATE] state array to add round key (r) to
 * @param uint8_t r - Index of round key to apply
 */
void AES::addRoundKey(byte *state, uint8_t r) {
  uint8_t offset = r * SIZEOF_STATE;

  state[0] ^= _expandedKey[offset];
  state[1] ^= _expandedKey[offset+1];
  state[2] ^= _expandedKey[offset+2];
  state[3] ^= _expandedKey[offset+3];
  state[4] ^= _expandedKey[offset+4];
  state[5] ^= _expandedKey[offset+5];
  state[6] ^= _expandedKey[offset+6];
  state[7] ^= _expandedKey[offset+7];
  state[8] ^= _expandedKey[offset+8];
  state[9] ^= _expandedKey[offset+9];
  state[10] ^= _expandedKey[offset+10];
  state[11] ^= _expandedKey[offset+11];
  state[12] ^= _expandedKey[offset+12];
  state[13] ^= _expandedKey[offset+13];
  state[14] ^= _expandedKey[offset+14];
  state[15] ^= _expandedKey[offset+15];
}

/*
 * cipher
 * @param byte *state - Pointer to byte[SIZEOF_STATE] state array to be ciphered
 * @param byte *key - Pointer to byte[SIZEOF_KEY] key array to expand and cipher with
 */
void AES::cipher(byte *state) {
  // Round 0
  addRoundKey(state, 0);
  
  // Rounds 0-9
  subBytes(state, SIZEOF_STATE); shiftRows(state); mixColumns(state); addRoundKey(state, 1);
  subBytes(state, SIZEOF_STATE); shiftRows(state); mixColumns(state); addRoundKey(state, 2);
  subBytes(state, SIZEOF_STATE); shiftRows(state); mixColumns(state); addRoundKey(state, 3);
  subBytes(state, SIZEOF_STATE); shiftRows(state); mixColumns(state); addRoundKey(state, 4);
  subBytes(state, SIZEOF_STATE); shiftRows(state); mixColumns(state); addRoundKey(state, 5);
  subBytes(state, SIZEOF_STATE); shiftRows(state); mixColumns(state); addRoundKey(state, 6);
  subBytes(state, SIZEOF_STATE); shiftRows(state); mixColumns(state); addRoundKey(state, 7);
  subBytes(state, SIZEOF_STATE); shiftRows(state); mixColumns(state); addRoundKey(state, 8);
  subBytes(state, SIZEOF_STATE); shiftRows(state); mixColumns(state); addRoundKey(state, 9);
  
  // Round 10
  subBytes(state, SIZEOF_STATE);
  shiftRows(state);
  addRoundKey(state, 10);
}

AES.h

/*!
 * The MIT License (MIT)
 * Copyright © 2022 Matt Borja (https://mattborja.dev/)
 * 
 * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
 * 
 * The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
 * 
 * THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 
 */
#ifndef AES_h
#define AES_h

#include "Arduino.h"

#define SIZEOF_BLOCK 4
#define SIZEOF_STATE 16
#define SIZEOF_KEY 16
#define SIZEOF_EXPANDED_KEY 176

const byte SBox[] = {
  0x63, 0x7c, 0x77, 0x7b, 0xf2, 0x6b, 0x6f, 0xc5, 0x30, 0x01, 0x67, 0x2b, 0xfe, 0xd7, 0xab, 0x76,
  0xca, 0x82, 0xc9, 0x7d, 0xfa, 0x59, 0x47, 0xf0, 0xad, 0xd4, 0xa2, 0xaf, 0x9c, 0xa4, 0x72, 0xc0,
  0xb7, 0xfd, 0x93, 0x26, 0x36, 0x3f, 0xf7, 0xcc, 0x34, 0xa5, 0xe5, 0xf1, 0x71, 0xd8, 0x31, 0x15,
  0x04, 0xc7, 0x23, 0xc3, 0x18, 0x96, 0x05, 0x9a, 0x07, 0x12, 0x80, 0xe2, 0xeb, 0x27, 0xb2, 0x75,
  0x09, 0x83, 0x2c, 0x1a, 0x1b, 0x6e, 0x5a, 0xa0, 0x52, 0x3b, 0xd6, 0xb3, 0x29, 0xe3, 0x2f, 0x84,
  0x53, 0xd1, 0x00, 0xed, 0x20, 0xfc, 0xb1, 0x5b, 0x6a, 0xcb, 0xbe, 0x39, 0x4a, 0x4c, 0x58, 0xcf,
  0xd0, 0xef, 0xaa, 0xfb, 0x43, 0x4d, 0x33, 0x85, 0x45, 0xf9, 0x02, 0x7f, 0x50, 0x3c, 0x9f, 0xa8,
  0x51, 0xa3, 0x40, 0x8f, 0x92, 0x9d, 0x38, 0xf5, 0xbc, 0xb6, 0xda, 0x21, 0x10, 0xff, 0xf3, 0xd2,
  0xcd, 0x0c, 0x13, 0xec, 0x5f, 0x97, 0x44, 0x17, 0xc4, 0xa7, 0x7e, 0x3d, 0x64, 0x5d, 0x19, 0x73,
  0x60, 0x81, 0x4f, 0xdc, 0x22, 0x2a, 0x90, 0x88, 0x46, 0xee, 0xb8, 0x14, 0xde, 0x5e, 0x0b, 0xdb,
  0xe0, 0x32, 0x3a, 0x0a, 0x49, 0x06, 0x24, 0x5c, 0xc2, 0xd3, 0xac, 0x62, 0x91, 0x95, 0xe4, 0x79,
  0xe7, 0xc8, 0x37, 0x6d, 0x8d, 0xd5, 0x4e, 0xa9, 0x6c, 0x56, 0xf4, 0xea, 0x65, 0x7a, 0xae, 0x08,
  0xba, 0x78, 0x25, 0x2e, 0x1c, 0xa6, 0xb4, 0xc6, 0xe8, 0xdd, 0x74, 0x1f, 0x4b, 0xbd, 0x8b, 0x8a,
  0x70, 0x3e, 0xb5, 0x66, 0x48, 0x03, 0xf6, 0x0e, 0x61, 0x35, 0x57, 0xb9, 0x86, 0xc1, 0x1d, 0x9e,
  0xe1, 0xf8, 0x98, 0x11, 0x69, 0xd9, 0x8e, 0x94, 0x9b, 0x1e, 0x87, 0xe9, 0xce, 0x55, 0x28, 0xdf,
  0x8c, 0xa1, 0x89, 0x0d, 0xbf, 0xe6, 0x42, 0x68, 0x41, 0x99, 0x2d, 0x0f, 0xb0, 0x54, 0xbb, 0x16
};

const byte RCon[] = { 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36 };
const byte GaloisField[] = { 0x02, 0x01, 0x01, 0x03, 0x03, 0x02, 0x01, 0x01, 0x01, 0x03, 0x02, 0x01, 0x01, 0x01, 0x03, 0x02 };

class AES
{
  public:
    AES(byte *key);
    void cipher(byte *state);

    // TODO
    // void decipher(byte *state);
  
  private:
    byte _expandedKey[SIZEOF_EXPANDED_KEY];
    bool _keyExpanded = false;
    
    void getColumn(byte *result, byte *state, uint8_t column);
    void getRow(byte *result, byte *state, uint8_t row);
    void rotWord(byte *block);
    void subBytes(byte *result, uint8_t len);
    void buildKeyRound(uint8_t w, uint8_t r);
    void keyExpansion(byte *key);
    void shiftRows(byte *state);
    void mixColumns(byte *state);
    void addRoundKey(byte *state, uint8_t r);
};

#endif

TinyAES.ino

/*
 * TinyAES Program
 *
 * The MIT License (MIT)
 * Copyright © 2022 Matt Borja (https://mattborja.dev/)
 * 
 * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
 * 
 * The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
 * 
 * THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 
 */

// Test vector: Rijndael_Animation_v4_eng.swf (OK)
// byte state[] = { 0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d, 0x31, 0x31, 0x98, 0xa2, 0xe0, 0x37, 0x07, 0x34 };
// byte key[] = { 0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c };
// byte expected[] = { 0x39, 0x25, 0x84, 0x1d, 0x02, 0xdc, 0x09, 0xfb, 0xdc, 0x11, 0x85, 0x97, 0x19, 0x6a, 0x0b, 0x32 };

// Test vector: FIPS-197 C.1, AES-128, Nk=4, Nr=10 (OK)
// byte state[] = { 0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff };
// byte key[] = { 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f };
// byte expected[] = { 0x69, 0xc4, 0xe0, 0xd8, 0x6a, 0x7b, 0x04, 0x30, 0xd8, 0xcd, 0xb7, 0x80, 0x70, 0xb4, 0xc5, 0x5a };

#include "Entropy.h"
#include "AES.h"

#define SERIAL_BAUD_RATE 9600
#define DEBUG false

// AES.h
#define SIZEOF_KEY 16
#define SIZEOF_STATE 16

// FIPS-197 example vector (C.1)
byte SYSTEM_KEY[SIZEOF_KEY] = { 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f };
byte EXPECTED[SIZEOF_STATE] = { 0x69, 0xc4, 0xe0, 0xd8, 0x6a, 0x7b, 0x04, 0x30, 0xd8, 0xcd, 0xb7, 0x80, 0x70, 0xb4, 0xc5, 0x5a };

void setup() {
  Serial.begin(SERIAL_BAUD_RATE);
  if (DEBUG) Serial.println("Initializing system...");

  Entropy.initialize();
  // TODO:
  // KeyManager.InitializeSystemKey()
  
  if (DEBUG) Serial.println("System initialized.");
}

void Serial_printhex(byte *state, uint8_t len) {
  for (uint8_t i = 0; i < len; i++) {
    if (state[i] < 0x10)
      Serial.print(0);
      
    Serial.print(state[i], HEX);
  }
}

void RNG (byte *buf, uint8_t len) {
  for (uint8_t i = 0; i < len; i++)
    buf[i] = Entropy.random(WDT_RETURN_BYTE);
}

void loop() {
  // RNG
  byte ENTROPY[SIZEOF_KEY];
  RNG(ENTROPY, SIZEOF_KEY);
  
  // Reset key and state before each cipher
  AES aes(SYSTEM_KEY);
  byte STATE[SIZEOF_STATE] = { 0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff };
  
  aes.cipher(STATE);

  Serial.println();
  Serial.print("\nExpected:     ");
  Serial_printhex(EXPECTED, SIZEOF_STATE);
  
  Serial.print("\nActual:       ");
  Serial_printhex(STATE, SIZEOF_STATE);
  
  Serial.print("\nRandom bytes: ");
  Serial_printhex(ENTROPY, SIZEOF_STATE);

  if (Serial.available() > 0) {
    Serial.print("\nUser input:   ");
      
    // Read up to 255 bytes
    for (uint8_t i = 0; i < 255; i++) {
      byte c[1] = { (byte)Serial.read() };
      if (c[0] == '\n')
        break;

      Serial.print((char)c[0]);
      Serial.print("(");
      Serial_printhex(c, 1);
      Serial.print(") ");
    }
  }

  delay(500);
}