#include <assert.h>
#include <inttypes.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>

#include "aes.h"
#include "block.h"
#include "cmac.h"

static const uint8_t zeros[16] = {0};

static uint8_t g_k1[16],
  g_k2[16];

#ifdef HOST_BUILD
void cmac_get_subkeys(uint8_t *dest) {
  /* Testing stub to get subkeys for algo check */
  memcpy(dest, g_k1, 16);
  memcpy(dest+16, g_k2, 16);
  return;
}
#endif /* HOST_BUILD */

void cmac_aes128_init(uint8_t *key) {
  /* Initialize AES engine and cache subkeys */
  aes128_init(key);
  cmac_aes128_expand_key(key, g_k1, g_k2);
}   

void cmac_aes128_expand_key(uint8_t const * const key, uint8_t *k1, uint8_t *k2) {
  /* Generate two required subkeys according to NIST 800-38B */
  uint8_t l[16] = {0},
    Rb[16] = {0, 0, 0, 0, 0, 0, 0, 0,
	      0, 0, 0, 0, 0, 0, 0, 0x87};
  
  aes128_ecb(l, zeros);
  if ((l[0] >> 7) == 0) {
    block_shiftl(k1, l, 1);
  } else {
    uint8_t tmp[16];
    block_shiftl(tmp, l, 1);
    block_xor(k1, tmp, Rb);
  }
  if (!(k1[0] >> 7)) {
    block_shiftl(k2, k1, 1);
  } else {
    uint8_t tmp[16];
    block_shiftl(tmp, k1, 1);
    block_xor(k2, tmp, Rb);
  }
}

static void cmac_truncate(uint8_t *dest, uint8_t *tag, uint_fast8_t tag_len) {
  memmove(dest, tag, tag_len);
}

void cmac_aes128(uint8_t *dest, uint8_t *msg, size_t msg_len, uint_fast8_t tag_len) {
  /* Simulate ceiling integer division by adding a block if remainder */
  size_t num_blocks = msg_len / 16 + (msg_len % 16 ? 1 : 0);
  bool last_block_complete = !(msg_len % 16 ? 1 : 0);
  if (msg_len == 0) {
    num_blocks = 1;
    last_block_complete = false;
  }
  
  uint8_t alt_msg[num_blocks*16],
    *last_block = &alt_msg[(num_blocks-1)*16];
  memset(alt_msg, 0, num_blocks*16);
  memmove(alt_msg, msg, msg_len);

  if (!last_block_complete) {
    /* Padding is single 1 bit, run out on 0s.. find the next byte,
       set it to 0b1000000 */
    alt_msg[msg_len] = 0x80;
    block_xor(last_block, last_block, g_k2);
  } else {
    block_xor(last_block, last_block, g_k1);
  }
  
  uint8_t x[16] = {0},
    y[16] = {0};

  for (uint32_t i = 0; i < num_blocks; i++) {
    uint8_t *block = &alt_msg[i*16];
    block_xor(y, x, block);
    aes128_ecb(x, y);
  }
  cmac_truncate(dest, x, tag_len);
  return;
}