/******************************************************************************
* C language implementations of the AES-128 and AES-256 key schedules to match
* the barrel-shiftrows representation. Note that those implementations rely on
* Look-Up Tables (LUT).
*
* See the paper at https://eprint.iacr.org/2020/1123.pdf for more details.
*
* @author 	Alexandre Adomnicai, Nanyang Technological University, Singapore
*			alexandre.adomnicai@ntu.edu.sg
*
* @date		August 2020
******************************************************************************/
#include "aes.h"
#include "internal-aes.h"

/******************************************************************************
* LUT of the AES Sbox used in the key schedule.
******************************************************************************/
static unsigned char sbox_lut[256] = {
	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
};

/******************************************************************************
* AES round constants used in the key schedule.
******************************************************************************/
static unsigned char rcon[11] = {
	0x00, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36
};

/******************************************************************************
* Packing routine used to rearrange the rkeys to match the barrel-shiftrows.
* It is about twice more efficient than the 'packing' func. This optimization
* is possible because the 8 16-byte blocs to pack are all equal.
* However, if one wants to encrypt 8 blocs in parallel with different keys, the
* 'packing' function is necessary.
******************************************************************************/
static void packing_rkey(uint32_t* rkey_bsr, const unsigned char* rkey) {
	uint32_t tmp, tmp0, tmp1, tmp2, tmp3;
	tmp0 = LE_LOAD_32(rkey);
	tmp1 = LE_LOAD_32(rkey + 4);
	tmp2 = LE_LOAD_32(rkey + 8);
	tmp3 = LE_LOAD_32(rkey + 12);
	SWAPMOVE(tmp0, tmp1, 0x00ff00ff, 8);
	SWAPMOVE(tmp2, tmp3, 0x00ff00ff, 8);
	SWAPMOVE(tmp0, tmp2, 0x0000ffff, 16);
	SWAPMOVE(tmp1, tmp3, 0x0000ffff, 16);
	for(int i = 0; i < 8; i++) {
		rkey_bsr[i] 	= (tmp0 & (0x80808080 >> i)) << i;
		rkey_bsr[i] 	|= (rkey_bsr[i] >> 1);
		rkey_bsr[i] 	|= (rkey_bsr[i] >> 2);
		rkey_bsr[i] 	|= (rkey_bsr[i] >> 4);
		rkey_bsr[i+8] 	= (tmp1 & (0x80808080 >> i)) << i;
		rkey_bsr[i+8] 	|= (rkey_bsr[i+8] >> 1);
		rkey_bsr[i+8] 	|= (rkey_bsr[i+8] >> 2);
		rkey_bsr[i+8] 	|= (rkey_bsr[i+8] >> 4);
		rkey_bsr[i+16] 	= (tmp2 & (0x80808080 >> i)) << i;
		rkey_bsr[i+16] 	|= (rkey_bsr[i+16] >> 1);
		rkey_bsr[i+16] 	|= (rkey_bsr[i+16] >> 2);
		rkey_bsr[i+16] 	|= (rkey_bsr[i+16] >> 4);
		rkey_bsr[i+24] 	= (tmp3 & (0x80808080 >> i)) << i;
		rkey_bsr[i+24] 	|= (rkey_bsr[i+24] >> 1);
		rkey_bsr[i+24] 	|= (rkey_bsr[i+24] >> 2);
		rkey_bsr[i+24] 	|= (rkey_bsr[i+24] >> 4);
	}
}

void aes128_keyschedule_lut(uint32_t* rkeys_bsr, const unsigned char* key) {
	uint32_t rkeys[44];
	// key schedule in the classical representation
	rkeys[0] = LE_LOAD_32(key);
	rkeys[1] = LE_LOAD_32(key + 4);
	rkeys[2] = LE_LOAD_32(key + 8);
	rkeys[3] = LE_LOAD_32(key + 12);
	for(int i = 4; i < 44; i+=4) {
		rkeys[i] = rkeys[i-4] ^ rcon[i/4];
		rkeys[i] ^= (sbox_lut[rkeys[i-1] & 0xff] << 24);
		rkeys[i] ^= sbox_lut[(rkeys[i-1] >> 8) & 0xff];
		rkeys[i] ^= (sbox_lut[rkeys[i-1] >> 24] << 16);
		rkeys[i] ^= (sbox_lut[(rkeys[i-1] >> 16) & 0xff] << 8);
		rkeys[i+1] = rkeys[i] ^ rkeys[i-3];
		rkeys[i+2] = rkeys[i+1] ^ rkeys[i-2];
		rkeys[i+3] = rkeys[i+2] ^ rkeys[i-1];
	}
	// packing all round keys to match the fully-fixsliced representation
	packing_rkey(rkeys_bsr, (unsigned char*)rkeys);
	for(int i = 1; i < 11; i++) {
		packing_rkey(rkeys_bsr + i*32, (unsigned char*)(rkeys + i*4));
		for(int j = 0; j < 32; j+=8) {
			rkeys_bsr[i*32+j+1] ^= 0xffffffff; 	// NOT to speed up SBox calculations
			rkeys_bsr[i*32+j+2] ^= 0xffffffff; 	// NOT to speed up SBox calculations
			rkeys_bsr[i*32+j+6] ^= 0xffffffff; 	// NOT to speed up SBox calculations
			rkeys_bsr[i*32+j+7] ^= 0xffffffff; 	// NOT to speed up SBox calculations
		}
	}
}

void aes256_keyschedule_lut(uint32_t* rkeys_bsr, const unsigned char* key) {
	uint32_t rkeys[60];
	// key schedule in the classical representation
	rkeys[0] = LE_LOAD_32(key);
	rkeys[1] = LE_LOAD_32(key + 4);
	rkeys[2] = LE_LOAD_32(key + 8);
	rkeys[3] = LE_LOAD_32(key + 12);
	rkeys[4] = LE_LOAD_32(key + 16);
	rkeys[5] = LE_LOAD_32(key + 20);
	rkeys[6] = LE_LOAD_32(key + 24);
	rkeys[7] = LE_LOAD_32(key + 28);
	// loop over double rounds for the round function
	for(int i = 8; i < 56; i+=8) {
		rkeys[i] = rkeys[i-8] ^ rcon[i/8];
		rkeys[i] ^= (sbox_lut[rkeys[i-1] & 0xff] << 24);
		rkeys[i] ^= sbox_lut[(rkeys[i-1] >> 8) & 0xff];
		rkeys[i] ^= (sbox_lut[rkeys[i-1] >> 24] << 16);
		rkeys[i] ^= (sbox_lut[(rkeys[i-1] >> 16) & 0xff] << 8);
		rkeys[i+1] = rkeys[i] ^ rkeys[i-7];
		rkeys[i+2] = rkeys[i+1] ^ rkeys[i-6];
		rkeys[i+3] = rkeys[i+2] ^ rkeys[i-5];
		rkeys[i+4] = rkeys[i-4];
		rkeys[i+4] ^= sbox_lut[rkeys[i+3] & 0xff];
		rkeys[i+4] ^= sbox_lut[(rkeys[i+3] >> 8) & 0xff] << 8;
		rkeys[i+4] ^= sbox_lut[(rkeys[i+3] >> 16) & 0xff] << 16;
		rkeys[i+4] ^= sbox_lut[rkeys[i+3] >> 24] << 24;
		rkeys[i+5] = rkeys[i+4] ^ rkeys[i-3];
		rkeys[i+6] = rkeys[i+5] ^ rkeys[i-2];
		rkeys[i+7] = rkeys[i+6] ^ rkeys[i-1];
	}
	// last round
	rkeys[56] = rkeys[48] ^ rcon[7];
	rkeys[56] ^= (sbox_lut[rkeys[55] & 0xff] << 24);
	rkeys[56] ^= sbox_lut[(rkeys[55] >> 8) & 0xff];
	rkeys[56] ^= (sbox_lut[rkeys[55] >> 24] << 16);
	rkeys[56] ^= (sbox_lut[(rkeys[55] >> 16) & 0xff] << 8);
	rkeys[57] = rkeys[56] ^ rkeys[49];
	rkeys[58] = rkeys[57] ^ rkeys[50];
	rkeys[59] = rkeys[58] ^ rkeys[51];
	// packing all round keys to match the fully-fixsliced representation
	packing_rkey(rkeys_bsr, (unsigned char*)rkeys);
	for(int i = 1; i < 15; i++) {
		packing_rkey(rkeys_bsr + i*32, (unsigned char*)(rkeys + i*4));
		for(int j = 0; j < 32; j+=8) {
			rkeys_bsr[i*32+j+1] ^= 0xffffffff; 	// NOT to speed up SBox calculations
			rkeys_bsr[i*32+j+2] ^= 0xffffffff; 	// NOT to speed up SBox calculations
			rkeys_bsr[i*32+j+6] ^= 0xffffffff; 	// NOT to speed up SBox calculations
			rkeys_bsr[i*32+j+7] ^= 0xffffffff; 	// NOT to speed up SBox calculations
		}
	}
}
