/******************************************************************************************
* Copyright 2015, 2016 Ideetron B.V.
*
* This program is free software: you can redistribute it and/or modify
* it under the terms of the GNU Lesser General Public License as published by
* the Free Software Foundation, either version 3 of the License, or
* (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
* GNU Lesser General Public License for more details.
*
* You should have received a copy of the GNU Lesser General Public License
* along with this program.  If not, see <http://www.gnu.org/licenses/>.
******************************************************************************************/
/******************************************************************************************
*
* File:        AES-128_V10.cpp
* Author:      Gerben den Hartog
* Compagny:    Ideetron B.V.
* Website:     http://www.ideetron.nl/LoRa
* E-mail:      info@ideetron.nl
******************************************************************************************/
/****************************************************************************************
*
* Created on: 			20-10-2015
* Supported Hardware: ID150119-02 Nexus board with RFM95
*
* Firmware Version 1.0
* First version
****************************************************************************************/

#include "AES-128_V10.h"

/*
********************************************************************************************
* Global Variables
********************************************************************************************
*/

unsigned char State[4][4];

unsigned char S_Table[16][16] = {
		{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}
};

/*
*****************************************************************************************
* Description : Function for encrypting data using AES-128
*
* Arguments   : *Data   Data to encrypt is a 16 byte long arry
*               *Key    Key to encrypt data with is a 16 byte long arry
*****************************************************************************************
*/
void AES_Encrypt(unsigned char *Data, unsigned char *Key)
{
	unsigned char i;
	unsigned char Row,Collum;
	unsigned char Round = 0x00;
	unsigned char Round_Key[16];

	//Copy input to State arry
	for(Collum = 0; Collum < 4; Collum++)
	{
		for(Row = 0; Row < 4; Row++)
		{
			State[Row][Collum] = Data[Row + (4*Collum)];
		}
	}

	//Copy key to round key
	for(i = 0; i < 16; i++)
	{
		Round_Key[i] = Key[i];
	}

	//Add round key
	AES_Add_Round_Key(Round_Key);

	//Preform 9 full rounds
	for(Round = 1; Round < 10; Round++)
	{
		//Preform Byte substitution with S table
		for(Collum = 0; Collum < 4; Collum++)
		{
			for(Row = 0; Row < 4; Row++)
			{
				State[Row][Collum] = AES_Sub_Byte(State[Row][Collum]);
			}
		}

		//Preform Row Shift
		AES_Shift_Rows();

		//Mix Collums
		AES_Mix_Collums();

		//Calculate new round key
		AES_Calculate_Round_Key(Round,Round_Key);

		//Add round key
		AES_Add_Round_Key(Round_Key);
	}

	//Last round whitout mix collums
	//Preform Byte substitution with S table
	for(Collum = 0; Collum < 4; Collum++)
	{
		for(Row = 0; Row < 4; Row++)
		{
			State[Row][Collum] = AES_Sub_Byte(State[Row][Collum]);
		}
	}

	//Shift rows
	AES_Shift_Rows();

	//Calculate new round key
	AES_Calculate_Round_Key(Round,Round_Key);

	//Add round Key
	AES_Add_Round_Key(Round_Key);

	//Copy the State into the data array
	for(Collum = 0; Collum < 4; Collum++)
	{
		for(Row = 0; Row < 4; Row++)
		{
			Data[Row + (4*Collum)] = State[Row][Collum];
		}
	}

}

/*
*****************************************************************************************
* Description : Function that add's the round key for the current round
*
* Arguments   : *Round_Key    16 byte long array holding the Round Key
*****************************************************************************************
*/
void AES_Add_Round_Key(unsigned char *Round_Key)
{
	unsigned char Row,Collum;

	for(Collum = 0; Collum < 4; Collum++)
	{
		for(Row = 0; Row < 4; Row++)
		{
			State[Row][Collum] = State[Row][Collum] ^ Round_Key[Row + (4*Collum)];
		}
	}
}

/*
*****************************************************************************************
* Description : Function that substitutes a byte with a byte from the S_Table
*
* Arguments   : Byte    The byte that will be substituted
*
* Return      : The return is the found byte in the S_Table
*****************************************************************************************
*/
unsigned char AES_Sub_Byte(unsigned char Byte)
{
	unsigned char S_Row,S_Collum;
	unsigned char S_Byte;

	//Split byte up in Row and Collum
	S_Row = ((Byte >> 4) & 0x0F);
	S_Collum = (Byte & 0x0F);

	//Find the correct byte in the S_Table
	S_Byte = S_Table[S_Row][S_Collum];

	return S_Byte;
}

/*
*****************************************************************************************
* Description : Function that preforms the shift row operation described in the AES standard
*****************************************************************************************
*/
void AES_Shift_Rows()
{
	unsigned char Buffer;

	//Row 0 doesn't change

	//Shift Row 1 one left
	//Store firt byte in buffer
	Buffer = State[1][0];
	//Shift all bytes
	State[1][0] = State[1][1];
	State[1][1] = State[1][2];
	State[1][2] = State[1][3];
	State[1][3] = Buffer;

	//Shift row 2 two left
	Buffer = State[2][0];
	State[2][0] = State[2][2];
	State[2][2] = Buffer;
	Buffer = State[2][1];
	State[2][1] = State[2][3];
	State[2][3] = Buffer;

	//Shift row 3 three left
	Buffer = State[3][3];
	State[3][3] = State[3][2];
	State[3][2] = State[3][1];
	State[3][1] = State[3][0];
	State[3][0] = Buffer;
}

/*
*****************************************************************************************
* Description : Function that preforms the Mix Collums operation described in the AES standard
*****************************************************************************************
*/
void AES_Mix_Collums()
{
	unsigned char Row,Collum;
	unsigned char a[4], b[4];
	for(Collum = 0; Collum < 4; Collum++)
	{
		for(Row = 0; Row < 4; Row++)
		{
			a[Row] = State[Row][Collum];
			b[Row] = (State[Row][Collum] << 1);

			if((State[Row][Collum] & 0x80) == 0x80)
			{
				b[Row] = b[Row] ^ 0x1B;
			}
		}
		State[0][Collum] = b[0] ^ a[1] ^ b[1] ^ a[2] ^ a[3];
		State[1][Collum] = a[0] ^ b[1] ^ a[2] ^ b[2] ^ a[3];
		State[2][Collum] = a[0] ^ a[1] ^ b[2] ^ a[3] ^ b[3];
		State[3][Collum] = a[0] ^ b[0] ^ a[1] ^ a[2] ^ b[3];
	}
}

/*
*****************************************************************************************
* Description : Function that calculaties the round key for the current round
*
* Arguments   :   Round         Number of current Round
*                *Round_Key     16 byte long array holding the Round Key
*****************************************************************************************
*/
void AES_Calculate_Round_Key(unsigned char Round, unsigned char *Round_Key)
{
	unsigned char i,j;
	unsigned char b;
	unsigned char Temp[4];
	unsigned char Buffer;
	unsigned char Rcon;

	//Calculate first Temp
	//Copy laste byte from previous key
	for(i = 0; i < 4; i++)
	{
		Temp[i] = Round_Key[i+12];
	}

	//Rotate Temp
	Buffer = Temp[0];
	Temp[0] = Temp[1];
	Temp[1] = Temp[2];
	Temp[2] = Temp[3];
	Temp[3] = Buffer;

	//Substitute Temp
	for(i = 0; i < 4; i++)
	{
		Temp[i] = AES_Sub_Byte(Temp[i]);
	}

	//Calculate Rcon
	Rcon = 0x01;
	while(Round != 1)
	{
		b = Rcon & 0x80;
		Rcon = Rcon << 1;
		if(b == 0x80)
		{
			Rcon = Rcon ^ 0x1b;
		}
		Round--;
	}

	//XOR Rcon
	Temp[0] = Temp[0] ^ Rcon;

	//Calculate new key
	for(i = 0; i < 4; i++)
	{
		for(j = 0; j < 4; j++)
		{
			Round_Key[j + (4*i)] = Round_Key[j + (4*i)] ^ Temp[j];
			Temp[j] = Round_Key[j + (4*i)];
		}
	}
}
