#include <stdexcept>
#include <cstring>
#include <cassert>
#include "aes.h"
#include "gf28.h"
#include "rng.h"

using namespace std;

namespace AES{

	//http://csrc.nist.gov/publications/fips/fips197/fips-197.pdf (AES)
	//https://tools.ietf.org/html/rfc3602 (CBC)

	//State is represented in bytes, as is the key schedule (which is represented in words in the AES spec).
	//State is in column-major order.

	//Most functions (and the three tables below) refer to similarly named functions and tables in the AES
	//spec. The core AES algorithm is not documented here.

	uint32_t roundconstant[10]={};
	uint8_t sbox[256]={};
	uint8_t invsbox[256]={}; //look-up table of sbox

	void initTables(){
		//generated tables have been checked with AES spec
		int term=0x01;
		for(int i=0;i<10;i++){
			roundconstant[i]=term<<24;
			term=GF28::multiply(term,0x02);
		}

		for(int i=0;i<256;i++){
			uint8_t inv=GF28::inverse(i);
			uint8_t res=0;
			for(int j=0;j<8;j++){
				uint8_t bit=((inv>>j)&1)^
				            ((inv>>((j+4)%8))&1)^
				            ((inv>>((j+5)%8))&1)^
				            ((inv>>((j+6)%8))&1)^
				            ((inv>>((j+7)%8))&1);
				res|=bit<<j;
			}
			res^=0x63;
			sbox[i]=res;
			invsbox[res]=i;
		}
	}

	uint32_t subWord(uint32_t word){
		return (sbox[word>>24]<<24)|(sbox[(word>>16)&0xff]<<16)|(sbox[(word>>8)&0xff]<<8)|sbox[word&0xff];
	}

	uint32_t rotWord(uint32_t word){
		return (word<<8)|(word>>24);
	}

	void keyExpand(uint8_t *keysched,const uint8_t *key,int keylen,int numrounds){
		memcpy(keysched,key,4*keylen);
		for(int i=keylen;i<4*(numrounds+1);i++){
			uint32_t temp=(keysched[4*i-4]<<24)|(keysched[4*i-3]<<16)|(keysched[4*i-2]<<8)|keysched[4*i-1];
			if(i%keylen==0){
				temp=subWord(rotWord(temp))^roundconstant[i/keylen-1];
			} else if(keylen>6&&i%keylen==4){
				temp=subWord(temp);
			}
			keysched[4*i+0]=keysched[4*(i-keylen)+0]^(temp>>24);
			keysched[4*i+1]=keysched[4*(i-keylen)+1]^((temp>>16)&0xff);
			keysched[4*i+2]=keysched[4*(i-keylen)+2]^((temp>>8)&0xff);
			keysched[4*i+3]=keysched[4*(i-keylen)+3]^(temp&0xff);
		}
	}

	void addRoundKey(uint8_t *state,const uint8_t *roundkey){
		for(int i=0;i<16;i++)state[i]^=roundkey[i];
	}

	void subBytes(uint8_t *state){
		for(int i=0;i<16;i++)state[i]=sbox[state[i]];
	}

	void shiftRows(uint8_t *state){
		uint8_t t=state[1]; state[1]=state[5]; state[5]=state[9]; state[9]=state[13]; state[13]=t;
		swap(state[2],state[10]); swap(state[6],state[14]);
		t=state[3]; state[3]=state[15]; state[15]=state[11]; state[11]=state[7]; state[7]=t;
	}

	void mixColumns(uint8_t *state){
		for(int i=0;i<4;i++){
			uint8_t a=GF28::multiply(0x02,state[4*i+0]) ^ GF28::multiply(0x03,state[4*i+1]) ^ state[4*i+2] ^ state[4*i+3];
			uint8_t b=state[4*i+0] ^ GF28::multiply(0x02,state[4*i+1]) ^ GF28::multiply(0x03,state[4*i+2]) ^ state[4*i+3];
			uint8_t c=state[4*i+0] ^ state[4*i+1] ^ GF28::multiply(0x02,state[4*i+2]) ^ GF28::multiply(0x03,state[4*i+3]);
			uint8_t d=GF28::multiply(0x03,state[4*i+0]) ^ state[4*i+1] ^ state[4*i+2] ^ GF28::multiply(0x02,state[4*i+3]);
			state[4*i+0]=a;
			state[4*i+1]=b;
			state[4*i+2]=c;
			state[4*i+3]=d;
		}
	}

	void invShiftRows(uint8_t *state){
		uint8_t t=state[1]; state[1]=state[13]; state[13]=state[9]; state[9]=state[5]; state[5]=t;
		swap(state[2],state[10]); swap(state[6],state[14]);
		t=state[3]; state[3]=state[7]; state[7]=state[11]; state[11]=state[15]; state[15]=t;
	}

	void invSubBytes(uint8_t *state){
		for(int i=0;i<16;i++)state[i]=invsbox[state[i]];
	}

	void invMixColumns(uint8_t *state){
		for(int i=0;i<4;i++){
			uint8_t a=GF28::multiply(0x0e,state[4*i+0])^
			            GF28::multiply(0x0b,state[4*i+1])^
			            GF28::multiply(0x0d,state[4*i+2])^
			            GF28::multiply(0x09,state[4*i+3]);
			uint8_t b=GF28::multiply(0x09,state[4*i+0])^
			            GF28::multiply(0x0e,state[4*i+1])^
			            GF28::multiply(0x0b,state[4*i+2])^
			            GF28::multiply(0x0d,state[4*i+3]);
			uint8_t c=GF28::multiply(0x0d,state[4*i+0])^
			            GF28::multiply(0x09,state[4*i+1])^
			            GF28::multiply(0x0e,state[4*i+2])^
			            GF28::multiply(0x0b,state[4*i+3]);
			uint8_t d=GF28::multiply(0x0b,state[4*i+0])^
			            GF28::multiply(0x0d,state[4*i+1])^
			            GF28::multiply(0x09,state[4*i+2])^
			            GF28::multiply(0x0e,state[4*i+3]);
			state[4*i+0]=a;
			state[4*i+1]=b;
			state[4*i+2]=c;
			state[4*i+3]=d;
		}
	}


	//The AES block encryption algorithm specified in the AES spec
	void encryptBlock(uint8_t *state,const uint8_t *keysched,const uint8_t *data,int numrounds){
		memcpy(state,data,16);

		addRoundKey(state,keysched);
		for(int round=0;round<numrounds-1;round++){
			subBytes(state);
			shiftRows(state);
			mixColumns(state);
			addRoundKey(state,keysched+16*(round+1));
		}
		subBytes(state);
		shiftRows(state);
		addRoundKey(state,keysched+16*numrounds);
	}

	//The AES block decryption algorithm specified in the AES spec
	void decryptBlock(uint8_t *state,const uint8_t *keysched,const uint8_t *data,int numrounds){
		memcpy(state,data,16);

		addRoundKey(state,keysched+16*numrounds);
		for(int round=numrounds-2;round>=0;round--){
			invShiftRows(state);
			invSubBytes(state);
			addRoundKey(state,keysched+16*(round+1));
			invMixColumns(state);
		}
		invShiftRows(state);
		invSubBytes(state);
		addRoundKey(state,keysched);
	}

	//Encrypt variably sized data with a key (assumed to be of valid length).
	//This applies the CBC algorithm: first generate an IV and add it to the
	//ciphertext as the first block. Then when encrypting the next block, first
	//xor it with the previous encrypted block. With this encryption method, all
	//of the encrypted text is influenced by the random IV, and therefore
	//different each time.
	//Padding is standard: if n>0 bytes are needed to reach a multiple of 16 bytes,
	//n bytes of padding are added with the value n. This is easily reversible.
	string encryptCBC(const string &data,const string &key,int numrounds){
		if(roundconstant[0]==0)initTables();

		int sz=data.size();
		if(sz==0)return {}; //if nothing to encrypt, don't even give an IV
		int blocks=sz/16;
		int padding=16-sz%16;
		string res;
		assert((sz+padding)%16==0);
		res.reserve(16+sz+padding);

		res.resize(16);
		CryptoRng crng;
		*(uint32_t*)&res[0]=crng.get(); //IV
		*(uint32_t*)&res[4]=crng.get(); //endianness doesn't matter, since the data is random anyway
		*(uint32_t*)&res[8]=crng.get();
		*(uint32_t*)&res[12]=crng.get();

		uint8_t keysched[16*(numrounds+1)];
		keyExpand(keysched,(const uint8_t*)key.data(),key.size()/4,numrounds);

		uint8_t buf[16],inbuf[16];
		for(int i=0;i<blocks;i++){
			memcpy(inbuf,data.data()+16*i,16);
			for(int j=0;j<16;j++)inbuf[j]^=res[res.size()-16+j]; //the CBC xor step
			encryptBlock(buf,keysched,inbuf,numrounds);
			res.insert(res.size(),(char*)buf,16);
		}

		if(padding<16)memcpy(inbuf,data.data()+16*blocks,16-padding);
		memset(inbuf+16-padding,padding,padding); //apply the padding
		for(int j=0;j<16;j++)inbuf[j]^=res[res.size()-16+j]; //the CBC xor step
		encryptBlock(buf,keysched,inbuf,numrounds); //encrypt the final padded block
		res.insert(res.size(),(char*)buf,16);
		return res;
	}

	//The inverse of encryptCBC
	string decryptCBC(const string &data,const string &key,int numrounds){
		if(roundconstant[0]==0)initTables();

		if(data.size()==0)return {};
		if(data.size()%16!=0)throw invalid_argument("AES encrypted data not multiple of block size");
		int blocks=data.size()/16-1; //the IV is not counted as a block
		string res(16*blocks,'\0');

		uint8_t keysched[16*(numrounds+1)];
		keyExpand(keysched,(const uint8_t*)key.data(),key.size()/4,numrounds);

		for(int i=blocks-1;i>=0;i--){
			decryptBlock((uint8_t*)&res[16*i],keysched,(const uint8_t*)data.data()+(16+16*i),numrounds);
			for(int j=0;j<16;j++)res[16*i+j]^=data[16*i+j]; //CBC: xor with the previous (remember the IV taking up space)
		}
		int padsize=res.back();
		if(padsize>16||padsize<0)throw invalid_argument("Malformed AES padding");
		res.resize(res.size()-padsize);
		return res;
	}

	//Public interface to encryptCBC
	string encrypt(const string &data,const string &key,Algorithm algo){
		int increment;
		switch(algo){
			case AES_128_CBC: increment=0; break;
			case AES_192_CBC: increment=1; break;
			case AES_256_CBC: increment=2; break;
			default: assert(false);
		}
		assert((int)key.size()==4*(4+2*increment));
		return encryptCBC(data,key,10+2*increment);
	}

	//Public interface to decryptCBC
	string decrypt(const string &data,const string &key,Algorithm algo){
		int increment;
		switch(algo){
			case AES_128_CBC: increment=0; break;
			case AES_192_CBC: increment=1; break;
			case AES_256_CBC: increment=2; break;
			default: assert(false);
		}
		if((int)key.size()!=4*(4+2*increment)){
			throw invalid_argument("Invalid AES key length");
		}
		return decryptCBC(data,key,10+2*increment);
	}

}