///////////////////////////////////////////////////////////////////////////////
/// @file dePluginRijndael.cpp
///
/// @brief rijndael (aes) encryption plugin
///
/// @author Lightning
///
/// This file is the intellectual property of Novus Delta, LLC.. Usage of the
/// contents of this file is subject to the Destiny3D Member License which
/// can be found at http://www.destiny3d.com.  Any other usage is prohibited.
///
/// This file is distributed "AS IS" without warranty of any kind.  Novus
/// Delta, LLC. does not guarantee the fitness of the contents of this file
/// for any particular purpose.
///
/// Copyright (C) 2001-2003 Novus Delta, LLC. All Rights Reserved.
///
/// <hr>
///                                 Change History
/// <hr>
///
/// @date Feb 2002
/// @author Lightning
/// @remarks Creation
///
/// @date Jun 2003
/// @author Lightning
/// @remarks Renamed from deEncryptionRijndael.cpp to dePluginRijndael.cpp
///
///////////////////////////////////////////////////////////////////////////////

#include <windows.h>
#include "dePlugins_priv.hpp"
#include "Rijndael/Aes.h"

//the rijndael encryption plugin

//ID of rijndael encryption
#ifdef DE3D_CPU_LITTLE_ENDIAN
#define RijndaelID	0x00004952		//RI\0\0, Rijndael, version 0.0
#else
#define RijndaelID	0x52490000		//RI\0\0, Rijndael, version 0.0
#endif

dePluginRijndael::dePluginRijndael()
{
	//setup the defaults
	memset(priv_Key, 0, sizeof(priv_Key));
	memset(priv_IV, 0, sizeof(priv_IV));
	priv_KeySize = 0;
}

dePluginRijndael::~dePluginRijndael()
{
	//clear out all the settings so no one knows what they are
	memset(priv_Key, 0, sizeof(priv_Key));
	memset(priv_IV, 0, sizeof(priv_IV));
	priv_KeySize = 0;
}

void* dePluginRijndael::GetInterface(IdePlugin::interface_t i)
{
	if (i == IdePlugin::iface_IdePlugin)
		return (IdePlugin*)(this);
	else if (i == IdePlugin::iface_IdePluginRijndael)
		return (dePluginRijndael*)(this);
	return NULL;
}

int dePluginRijndael::Release()
{
	delete this;
	return 0;
}

deBoolean dePluginRijndael::EncodeData(void *Buffer, DWORD Length, void *NewBuffer, DWORD *NewLength)
{
	aes_encrypt_ctx		ctx;
	long				XorCounter;
	BYTE *				EncBufferPtr;
	BYTE *				PrevEncBufferPtr;

	//if no valid pointers, return deFALSE
	if(!Buffer || !Length || !NewBuffer || !NewLength)
		return deFALSE;

	//if no data to use, fail
	if(!priv_KeySize)
		return deFALSE;

	//verify that *NewLength is at least the size of the needed length
	if(*NewLength < MaxEncodeLength(Length))
		return deFALSE;

	//reset the encryption struct
	memset(&ctx, 0, sizeof(ctx));

	//setup the encryption key
	if(aes_encrypt_key((BYTE *)priv_Key, priv_KeySize, ctx) == aes_error)
		return deFALSE;

	//the rounds are determined by Rijndael according to the docs

	//setup the beginning part of the buffer for the header
	((RijndaelHeader *)NewBuffer)->ID = RijndaelID;			//id that it is rijndael
	((RijndaelHeader *)NewBuffer)->OriginalSize =  deByteOrderBE32(Length);	//original length

	//set where the encrypted data goes
	EncBufferPtr = &((BYTE *)NewBuffer)[sizeof(RijndaelHeader)];

	//copy the data to encrypt to the new buffer
	memcpy(EncBufferPtr, Buffer, Length);

	//make sure the new length is set correctly
	*NewLength = MaxEncodeLength(Length);

	//set the length to not include the header
	Length = *NewLength - sizeof(RijndaelHeader);

	//loop thru the buffer and encrypt it
	PrevEncBufferPtr = (BYTE *)priv_IV;
	while(Length)
	{
		//xor the data together
		for(XorCounter = 0; XorCounter < (AES_BLOCK_SIZE / 4); XorCounter++)
			((DWORD *)EncBufferPtr)[XorCounter] ^= ((DWORD *)PrevEncBufferPtr)[XorCounter];

		//now, encrypt the data
		if(aes_encrypt(EncBufferPtr, EncBufferPtr, ctx) != aes_good)
		{
			memset(&ctx, 0, sizeof(ctx));
			return deFALSE;
		}
		
		//go to the next block
		PrevEncBufferPtr = EncBufferPtr;
		EncBufferPtr = EncBufferPtr + AES_BLOCK_SIZE;
		Length = Length - AES_BLOCK_SIZE;
	}

	//wipe out the encrypt struct
	memset(&ctx, 0, sizeof(ctx));

	return deTRUE;
}

deBoolean dePluginRijndael::DecodeData(void *Buffer, DWORD Length, void *NewBuffer, DWORD *NewLength)
{
	aes_decrypt_ctx		ctx;
	long				XorCounter;
	BYTE *				DecBufferPtr;
	BYTE *				PrevDecBufferPtr;

	//if no valid pointers, return deFALSE
	if(!Buffer || !Length || !NewBuffer || !NewLength)
		return deFALSE;

	//if no data to use, fail
	if(!priv_KeySize)
		return deFALSE;

	//verify that the buffer is ours
	if(((RijndaelHeader *)Buffer)->ID != RijndaelID)
		return deFALSE;

	//verify that *NewLength is at least the size of the needed length
	if(*NewLength < deByteOrderBE32(((RijndaelHeader *)Buffer)->OriginalSize))
		return deFALSE;

	//reset the decryption struct
	memset(&ctx, 0, sizeof(ctx));

	//setup the decryption key
	if(aes_decrypt_key((BYTE *)priv_Key, priv_KeySize, ctx) == aes_error)
		return deFALSE;

	//the rounds are determined by Rijndael according to the docs

	//set the length to not include the header
	Length = Length - sizeof(RijndaelHeader);

	//verify the length is valid for the block size (should be a multiple)
	if(Length % AES_BLOCK_SIZE)
		return deFALSE;

	//loop thru the buffer and decrypt it
	DecBufferPtr = &((BYTE *)Buffer)[sizeof(RijndaelHeader)];
	PrevDecBufferPtr = (BYTE *)priv_IV;
	while(Length)
	{
		//decrypt the data
		if(aes_decrypt(DecBufferPtr, (BYTE *)NewBuffer, ctx) != aes_good)
		{
			memset(&ctx, 0, sizeof(ctx));
			return deFALSE;
		}

		//remove the xor
		for(XorCounter = 0; XorCounter < (AES_BLOCK_SIZE / 4); XorCounter++)
			((DWORD *)NewBuffer)[XorCounter] ^= ((DWORD *)PrevDecBufferPtr)[XorCounter];

		//go to the next block
		PrevDecBufferPtr = DecBufferPtr;
		DecBufferPtr = DecBufferPtr + AES_BLOCK_SIZE;
		NewBuffer = (BYTE *)NewBuffer + AES_BLOCK_SIZE;
		Length = Length - AES_BLOCK_SIZE;
	}

	//wipe out the decrypt struct
	memset(&ctx, 0, sizeof(ctx));

	//set the new length to the original length
	*NewLength = deByteOrderBE32(((RijndaelHeader *)Buffer)->OriginalSize);

	return deTRUE;
}

DWORD dePluginRijndael::MaxEncodeLength(DWORD Length)
{
	//return the maximum length for encrypted data
	DWORD	NeededLength;

	//figure out how many blocks is needed
	NeededLength = Length / AES_BLOCK_SIZE;
	if(Length % AES_BLOCK_SIZE)
		NeededLength++;

	//add in the length of the header
	NeededLength = (NeededLength * AES_BLOCK_SIZE) + sizeof(RijndaelHeader);

	return NeededLength;
}

DWORD dePluginRijndael::MaxDecodeLength(void *Buffer)
{
	DWORD	NeededLength;
	DWORD	OriginalSize;

	//return the maximum length for decryption

	//if no buffer, return 0
	if(!Buffer)
		return 0;

	//verify the header
	if(((RijndaelHeader *)Buffer)->ID != RijndaelID)
		return 0;

	//now, figure out just the needed size to bring the original length up to a block
	//multiple

	OriginalSize = deByteOrderBE32(((RijndaelHeader *)Buffer)->OriginalSize);

	//figure out the next multiple of a block with length
	NeededLength = OriginalSize / AES_BLOCK_SIZE;
	if(OriginalSize % AES_BLOCK_SIZE)
		NeededLength++;

	//neededlength = number of characters to round to a block. add the original
	//length for total needed length
	NeededLength = NeededLength * AES_BLOCK_SIZE;

	return NeededLength;
}

deBoolean dePluginRijndael::SetKey(void *KeyBuffer, long Bitsize)
{
	//what key to use

	//verify the bitsize
	if(Bitsize != 128 && Bitsize != 192 && Bitsize != 256)
		return deFALSE;

	//verify the buffer
	if(!KeyBuffer)
		return deFALSE;

	//set the bitsize
	priv_KeySize = Bitsize / 8;

	//copy the key to the key buffer
	memcpy(priv_Key, KeyBuffer, Bitsize / 8);
	return deTRUE;
}

deBoolean dePluginRijndael::SetIV(void *IVBuffer, long Size)
{
	//what IV to use

	//see if they passed in a buffer
	if(!IVBuffer)
	{
		//no buffer, blank out the IV data
		memset(priv_IV, 0, AES_BLOCK_SIZE);
	}
	else
	{
		//if the length isnt valid then fail
		if(Size != AES_BLOCK_SIZE)
			return deFALSE;

		//copy it
		memcpy(priv_IV, IVBuffer, AES_BLOCK_SIZE);
	}

	//return true
	return deTRUE;
}
