#include "aes_core.h"
#include <assert.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

static void rk_init_ctr(struct ctr_state *state, const unsigned char *iv) {
  state->num = 0;
  memset(state->ecount, 0, 16);
  memcpy(state->ivec, iv, 16);
}

/* increment counter (128-bit int) by 1 */
static void rk_ctr128_inc(unsigned char *counter) {
  unsigned int n = 16;
  unsigned char c;

  do {
    --n;
    c = counter[n];
    ++c;
    counter[n] = c;
    if (c)
      return;
  } while (n);
}

static void rk_crypto_ctr128_encrypt(const unsigned char *in,
                                     unsigned char *out, int len,
                                     const void *key, unsigned char *ivec,
                                     unsigned char *ecount_buf,
                                     unsigned int *num) {
  int n, l = 0;

  n = *num;
  while (l < len) {
    if (n == 0) {
      rk_aes_encrypt(ivec, ecount_buf, key);
      rk_ctr128_inc(ivec);
    }
    out[l] = in[l] ^ ecount_buf[n];
    ++l;
    n = (n + 1) % 16;
  }

  *num = n;
}

/* XTS makes use of two different keys, usually generated by splitting
 * the supplied block cipher's key in half.
 * Because of the splitting, users wanting AES 256 and AES 128 encryption
 * will need to choose key sizes of 512 bits and 256 bits respectively.
 */
int rk_aes_ctr_encrypt(const unsigned char *in, unsigned char *out,
                       unsigned long length, const unsigned char *key,
                       const int key_len, const unsigned char *ivec,
                       const int enc) {
  RK_AES_KEY ks;
  struct ctr_state state;

  if (key_len != 16 && key_len != 24 && key_len != 32)
    return -1;

  rk_init_ctr(&state, ivec);
  rk_aes_set_encrypt_key(key, key_len * 8, &ks);
  rk_crypto_ctr128_encrypt(in, out, length, &ks, state.ivec, state.ecount,
                           &state.num);

  return 0;
}