#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include "aes_core.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;
}