#include "aes.h"

// should probably just pre-generate them, will do it later.

static u8	pow_tab[256];
static u8	log_tab[256];
static u8	sbx_tab[256];
static u8	isb_tab[256];
static u32	rco_tab[7];
static u32	ft_tab[4][256];
static u32	it_tab[4][256];
static u32	fl_tab[4][256];
static u32	il_tab[4][256];

#define ff_mult(a,b)	(((a) && (b)) ? pow_tab[(log_tab[a] + log_tab[b]) % 255] : 0)
#define byte(x, n)		(((x) >> (n << 3)) & 0xFF)
#define f_rn(y,x,n,k)	(y[n] = ft_tab[0][byte(x[n],0)] ^ ft_tab[1][byte(x[(n + 1) & 3],1)] ^ ft_tab[2][byte(x[(n + 2) & 3],2)] ^ ft_tab[3][byte(x[(n + 3) & 3],3)] ^ k[n])
#define f_rl(y,x,n,k)	(y[n] = fl_tab[0][byte(x[n],0)] ^ fl_tab[1][byte(x[(n + 1) & 3],1)] ^ fl_tab[2][byte(x[(n + 2) & 3],2)] ^ fl_tab[3][byte(x[(n + 3) & 3],3)] ^ k[n])
#define i_rn(y,x,n,k)	(y[n] = it_tab[0][byte(x[n],0)] ^ it_tab[1][byte(x[(n + 3) & 3],1)] ^ it_tab[2][byte(x[(n + 2) & 3],2)] ^ it_tab[3][byte(x[(n + 1) & 3],3)] ^ k[n])
#define i_rl(y,x,n,k)	(y[n] = il_tab[0][byte(x[n],0)] ^ il_tab[1][byte(x[(n + 3) & 3],1)] ^ il_tab[2][byte(x[(n + 2) & 3],2)] ^ il_tab[3][byte(x[(n + 1) & 3],3)] ^ k[n])
#define ls_box(x)		(fl_tab[0][byte(x,0)] ^ fl_tab[1][byte(x,1)] ^ fl_tab[2][byte(x,2)] ^ fl_tab[3][byte(x,3)])
#define star_x(x)		(((x) & 0x7F7F7F7F) << 1) ^ ((((x) & 0x80808080) >> 7) * 0x1B)
#define imix_col(y,x)	(u = star_x(x), v = star_x(u), w = star_x(v), t = w ^ (x), (y) = u ^ v ^ w, (y) ^= rotr32 (u ^ t, 8) ^ rotr32 (v ^ t, 16) ^ rotr32 (t, 24))
#define f_nround(y,x,k)	(f_rn (y, x, 0, k), f_rn (y, x, 1, k), f_rn (y, x, 2, k), f_rn (y, x, 3, k), k += 4)
#define f_lround(y,x,k)	(f_rl (y, x, 0, k), f_rl (y, x, 1, k), f_rl (y, x, 2, k), f_rl (y, x, 3, k))
#define i_nround(y,x,k)	(i_rn (y, x, 0, k), i_rn (y, x, 1, k), i_rn (y, x, 2, k), i_rn (y, x, 3, k), k -= 4)
#define i_lround(y,x,k)	(i_rl (y, x, 0, k), i_rl (y, x, 1, k), i_rl (y, x, 2, k), i_rl (y, x, 3, k))

static void __cdecl generate_rijndael_tables (void)
{
	u32			i, t;
	u8			p, q;
	
	/* log and power tables for GF(2**8) finite field with */
	/* 0x11B as modular polynomial - the simplest prmitive */
	/* root is 0x11, used here to generate the tables	   */
	
	for (i = 0, p = 1; i < 256; i++)
	{
		pow_tab[i] = (u8) p; log_tab[p] = (u8) i;
		p = p ^ (p << 1) ^ (p & 0x80 ? 0x01B : 0);
	}
	
	log_tab[1] = 0; p = 1;
	
	for (i = 0; i < 7; i++)
	{
		rco_tab[i] = p;
		p = (p << 1) ^ (p & 0x80 ? 0x1b : 0);
	}
	
	/* note that the affine byte transformation matrix in   */
	/* rijndael specification is in big endian format with  */
	/* bit 0 as the most significant bit. In the remainder  */
	/* of the specification the bits are numbered from the  */
	/* least significant end of a byte.					 */
	
	for (i = 0; i < 256; ++i)
	{   
		p = (i ? pow_tab[255 - log_tab[i]] : 0); q = p;
		q = (q >> 7) | (q << 1); p ^= q;
		q = (q >> 7) | (q << 1); p ^= q;
		q = (q >> 7) | (q << 1); p ^= q;
		q = (q >> 7) | (q << 1); p ^= q ^ 0x63;
		sbx_tab[i] = (u8) p; isb_tab[p] = (u8) i;
	}
	
	for (i = 0; i < 256; ++i)
	{
		fl_tab[0][i] = t = p = sbx_tab[i];
		fl_tab[1][i] = rotl32 (t,  8);
		fl_tab[2][i] = rotl32 (t, 16);
		fl_tab[3][i] = rotl32 (t, 24);
		
		ft_tab[0][i] = t = ((u32) ff_mult(2, p)) | ((u32) p <<  8) | ((u32) p << 16) | ((u32) ff_mult(3, p) << 24);
		ft_tab[1][i] = rotl32 (t,  8);
		ft_tab[2][i] = rotl32 (t, 16);
		ft_tab[3][i] = rotl32 (t, 24);
		
		il_tab[0][i] = t = p = isb_tab[i];
		il_tab[1][i] = rotl32(t,  8);
		il_tab[2][i] = rotl32(t, 16);
		il_tab[3][i] = rotl32(t, 24);
		
		it_tab[0][i] = t = ((u32) ff_mult (14, p)) | ((u32) ff_mult (9, p) << 8) | ((u32) ff_mult (13, p) << 16) | ((u32) ff_mult (11, p) << 24);
		it_tab[1][i] = rotl32(t,  8);
		it_tab[2][i] = rotl32(t, 16);
		it_tab[3][i] = rotl32(t, 24);
	}
}

aes_keystream * __cdecl rijndael_setkey (const u32 *key, aes_keystream *ks)
{
	u32			i, t, u, v, w;
	
	if (!pow_tab[0]) generate_rijndael_tables();
	
	for (i = 0; i < 8; i++) ks->rijndael.e_key[i] = t = lsf32 (key[i]);
	
	for (i = 0; i < 7; i++)
	{
		t = ls_box(rotr32(t,8))^rco_tab[i]; t ^= ks->rijndael.e_key[i*8];
		ks->rijndael.e_key[i*8 +  8] = t;	t ^= ks->rijndael.e_key[i*8 + 1];
		ks->rijndael.e_key[i*8 +  9] = t;	t ^= ks->rijndael.e_key[i*8 + 2];
		ks->rijndael.e_key[i*8 + 10] = t;	t ^= ks->rijndael.e_key[i*8 + 3];
		ks->rijndael.e_key[i*8 + 11] = t;	t  = ks->rijndael.e_key[i*8 + 4] ^ ls_box(t);
		ks->rijndael.e_key[i*8 + 12] = t;	t ^= ks->rijndael.e_key[i*8 + 5];
		ks->rijndael.e_key[i*8 + 13] = t;	t ^= ks->rijndael.e_key[i*8 + 6];
		ks->rijndael.e_key[i*8 + 14] = t;	t ^= ks->rijndael.e_key[i*8 + 7];
		ks->rijndael.e_key[i*8 + 15] = t;
	}
	memcpy (ks->rijndael.d_key, ks->rijndael.e_key, 16);
	
	for (i = 4; i < 56; i++) imix_col (ks->rijndael.d_key[i], ks->rijndael.e_key[i]);
	
	return ks;
}

u32 * __cdecl rijndael_encrypt (u32 *one_block, const aes_keystream *ks)
{
	u32				b0[4], b1[4];
	const u32		*kp = ks->rijndael.e_key + 4;
	
	b0[0] = lsf32 (one_block[0]) ^ ks->rijndael.e_key[0];
	b0[1] = lsf32 (one_block[1]) ^ ks->rijndael.e_key[1];
	b0[2] = lsf32 (one_block[2]) ^ ks->rijndael.e_key[2];
	b0[3] = lsf32 (one_block[3]) ^ ks->rijndael.e_key[3];
	
	f_nround (b1, b0, kp); f_nround (b0, b1, kp);
	f_nround (b1, b0, kp); f_nround (b0, b1, kp);
	f_nround (b1, b0, kp); f_nround (b0, b1, kp);
	f_nround (b1, b0, kp); f_nround (b0, b1, kp);
	f_nround (b1, b0, kp); f_nround (b0, b1, kp);
	f_nround (b1, b0, kp); f_nround (b0, b1, kp);
	f_nround (b1, b0, kp); f_lround (b0, b1, kp);
	
	one_block[0] = lsf32 (b0[0]);
	one_block[1] = lsf32 (b0[1]);
	one_block[2] = lsf32 (b0[2]);
	one_block[3] = lsf32 (b0[3]);
	
	return one_block;
}

u32 * __cdecl rijndael_decrypt (u32 *one_block, const aes_keystream *ks)
{
	u32				b0[4], b1[4];
	const u32		*kp = ks->rijndael.d_key + 52;
	
	b0[0] = lsf32 (one_block[0]) ^ ks->rijndael.e_key[56];
	b0[1] = lsf32 (one_block[1]) ^ ks->rijndael.e_key[57];
	b0[2] = lsf32 (one_block[2]) ^ ks->rijndael.e_key[58];
	b0[3] = lsf32 (one_block[3]) ^ ks->rijndael.e_key[59];
	
	i_nround (b1, b0, kp); i_nround (b0, b1, kp);
	i_nround (b1, b0, kp); i_nround (b0, b1, kp);
	i_nround (b1, b0, kp); i_nround (b0, b1, kp);
	i_nround (b1, b0, kp); i_nround (b0, b1, kp);
	i_nround (b1, b0, kp); i_nround (b0, b1, kp);
	i_nround (b1, b0, kp); i_nround (b0, b1, kp);
	i_nround (b1, b0, kp); i_lround (b0, b1, kp);
	
	one_block[0] = lsf32 (b0[0]);
	one_block[1] = lsf32 (b0[1]);
	one_block[2] = lsf32 (b0[2]);
	one_block[3] = lsf32 (b0[3]);
	
	return one_block;
}