/*
 * Copyright 1997-2005 Markus Hahn 
 * 
 * Licensed under the Apache License, Version 2.0 (the "License"); 
 * you may not use this file except in compliance with the License. 
 * You may obtain a copy of the License at 
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "CipherServer.h"
#include "cpkernel.h"
#include "Yarrow.h"

#include "AES.h"
#include "Blowfish.h"
#include "cast.h"
#include "arcfour.h"
#include "Serpent.h"
#include "TMS.h"
#include "TripleDES.h"
#include "Twofish.h"

#include <stdlib.h>
#include <string.h>
#include <memory.h>

// fixed key length (i.n.) for extended selftest
#define FIXKEYLEN   16

//////////////////////////////////////////////////////////////////////////////

struct CIPHERCTX
{
	// function addresses
	Cipher_GetCipherInfo*         pGetCipherInfo;
	Cipher_SelfTest*              pSelfTest;
	Cipher_CreateWorkContext*     pCreateWorkContext;
	Cipher_ResetWorkContext*      pResetWorkContext;
	Cipher_DestroyWorkContext*    pDestroyWorkContext;
	Cipher_EncryptBuffer*         pEncryptBuffer;
	Cipher_DecryptBuffer*         pDecryptBuffer;

	//  cipher information
	CIPHERINFOBLOCK infoblock;

	// function reference for random number generation
	Cipher_RandomGenerator* pRandomGenerator;

	// data reference for the random generator
	const void* pRndGenData;

	// to check if the internal random generator is used
	BYTEBOOL blUseInternalRndGen;
};

//////////////////////////////////////////////////////////////////////////////

struct CIPHERSESSION
{
	// crypt mode, see CIPHERSERVER_MODE_xxx constants
	WORD32 lCryptMode;

	// pointer to a cipher context
	PCIPHERCTX pCipherContext;

	// pointer to an algorithm's runtime memory area
	void* pWorkContext;
};

//////////////////////////////////////////////////////////////////////////////

// lookup table to fetch the function addresses

typedef struct
{
	char* szName;
	Cipher_GetCipherInfo*       pGetCipherInfo;
	Cipher_SelfTest*            pSelfTest;
	Cipher_CreateWorkContext*   pCreateWorkContext;
	Cipher_ResetWorkContext*    pResetWorkContext;
	Cipher_DestroyWorkContext*  pDestroyWorkContext;
	Cipher_EncryptBuffer*       pEncryptBuffer;
	Cipher_DecryptBuffer*       pDecryptBuffer;
}
CIPHERLOOKUP;

//////////////////////////////////////////////////////////////////////////////

#define NUM_OF_CIPHERS  7

CIPHERLOOKUP cipherFaces[NUM_OF_CIPHERS] = {

	{       TMS_CIPHERNAME,       TMS_GetCipherInfo,       TMS_SelfTest,       TMS_CreateWorkContext,       TMS_ResetWorkContext,       TMS_DestroyWorkContext,       TMS_EncryptBuffer,       TMS_DecryptBuffer },
	{  BLOWFISH_CIPHERNAME,  Blowfish_GetCipherInfo,  Blowfish_SelfTest,  Blowfish_CreateWorkContext,  Blowfish_ResetWorkContext,  Blowfish_DestroyWorkContext,  Blowfish_EncryptBuffer,  Blowfish_DecryptBuffer },
	{      CAST_CIPHERNAME,      CAST_GetCipherInfo,      CAST_SelfTest,      CAST_CreateWorkContext,      CAST_ResetWorkContext,      CAST_DestroyWorkContext,      CAST_EncryptBuffer,      CAST_DecryptBuffer },
	{   ARCFOUR_CIPHERNAME,   ARCFOUR_GetCipherInfo,   ARCFOUR_SelfTest,   ARCFOUR_CreateWorkContext,   ARCFOUR_ResetWorkContext,   ARCFOUR_DestroyWorkContext,   ARCFOUR_EncryptBuffer,   ARCFOUR_DecryptBuffer },
	{   SERPENT_CIPHERNAME,   Serpent_GetCipherInfo,   Serpent_SelfTest,   Serpent_CreateWorkContext,   Serpent_ResetWorkContext,   Serpent_DestroyWorkContext,   Serpent_EncryptBuffer,   Serpent_DecryptBuffer },
	{       AES_CIPHERNAME,       AES_GetCipherInfo,       AES_SelfTest,       AES_CreateWorkContext,       AES_ResetWorkContext,       AES_DestroyWorkContext,       AES_EncryptBuffer,       AES_DecryptBuffer },
	{ TRIPLEDES_CIPHERNAME, TripleDES_GetCipherInfo, TripleDES_SelfTest, TripleDES_CreateWorkContext, TripleDES_ResetWorkContext, TripleDES_DestroyWorkContext, TripleDES_EncryptBuffer, TripleDES_DecryptBuffer },
	//{   TWOFISH_CIPHERNAME,   Twofish_GetCipherInfo,   Twofish_SelfTest,   Twofish_CreateWorkContext,   Twofish_ResetWorkContext,   Twofish_DestroyWorkContext,   Twofish_EncryptBuffer,   Twofish_DecryptBuffer }
};
//////////////////////////////////////////////////////////////////////////////

// simple table storing the cipher names
char* cipherNames[NUM_OF_CIPHERS] =
{
	TMS_CIPHERNAME,
	BLOWFISH_CIPHERNAME,
	CAST_CIPHERNAME,
	ARCFOUR_CIPHERNAME,
	SERPENT_CIPHERNAME,
	AES_CIPHERNAME,
	TRIPLEDES_CIPHERNAME,
	//TWOFISH_CIPHERNAME
};
//
//#define NUM_OF_CIPHERS  2
//
//CIPHERLOOKUP cipherFaces[NUM_OF_CIPHERS] = {
//
//   // {       TMS_CIPHERNAME,       TMS_GetCipherInfo,       TMS_SelfTest,       TMS_CreateWorkContext,       TMS_ResetWorkContext,       TMS_DestroyWorkContext,       TMS_EncryptBuffer,       TMS_DecryptBuffer },
//    {  BLOWFISH_CIPHERNAME,  Blowfish_GetCipherInfo,  Blowfish_SelfTest,  Blowfish_CreateWorkContext,  Blowfish_ResetWorkContext,  Blowfish_DestroyWorkContext,  Blowfish_EncryptBuffer,  Blowfish_DecryptBuffer },
//    //{      CAST_CIPHERNAME,      CAST_GetCipherInfo,      CAST_SelfTest,      CAST_CreateWorkContext,      CAST_ResetWorkContext,      CAST_DestroyWorkContext,      CAST_EncryptBuffer,      CAST_DecryptBuffer },
//    //{   ARCFOUR_CIPHERNAME,   ARCFOUR_GetCipherInfo,   ARCFOUR_SelfTest,   ARCFOUR_CreateWorkContext,   ARCFOUR_ResetWorkContext,   ARCFOUR_DestroyWorkContext,   ARCFOUR_EncryptBuffer,   ARCFOUR_DecryptBuffer },
//    //{   SERPENT_CIPHERNAME,   Serpent_GetCipherInfo,   Serpent_SelfTest,   Serpent_CreateWorkContext,   Serpent_ResetWorkContext,   Serpent_DestroyWorkContext,   Serpent_EncryptBuffer,   Serpent_DecryptBuffer },
//    {       AES_CIPHERNAME,       AES_GetCipherInfo,       AES_SelfTest,       AES_CreateWorkContext,       AES_ResetWorkContext,       AES_DestroyWorkContext,       AES_EncryptBuffer,       AES_DecryptBuffer },
//   // { TRIPLEDES_CIPHERNAME, TripleDES_GetCipherInfo, TripleDES_SelfTest, TripleDES_CreateWorkContext, TripleDES_ResetWorkContext, TripleDES_DestroyWorkContext, TripleDES_EncryptBuffer, TripleDES_DecryptBuffer },
//   // {   TWOFISH_CIPHERNAME,   Twofish_GetCipherInfo,   Twofish_SelfTest,   Twofish_CreateWorkContext,   Twofish_ResetWorkContext,   Twofish_DestroyWorkContext,   Twofish_EncryptBuffer,   Twofish_DecryptBuffer }
//};
////////////////////////////////////////////////////////////////////////////////
//
//// simple table storing the cipher names
//char* cipherNames[NUM_OF_CIPHERS] =
//{
//   // TMS_CIPHERNAME,
//    BLOWFISH_CIPHERNAME,
//   // CAST_CIPHERNAME,
//  //  ARCFOUR_CIPHERNAME,
//  //  SERPENT_CIPHERNAME,
//    AES_CIPHERNAME,
//  //  TRIPLEDES_CIPHERNAME,
//  //  TWOFISH_CIPHERNAME
//};


//////////////////////////////////////////////////////////////////////////////

WORD32 CRYPTPAK_API CipherServer_GetCipherNames
(char*** pszList)
{
	*pszList = cipherNames;
	return NUM_OF_CIPHERS;
}

//////////////////////////////////////////////////////////////////////////////

WORD32 CRYPTPAK_API CipherServer_GetCipherInfo
(const char* pCipherName,
 CIPHERINFOBLOCK* pInfoBlock)
{
	Cipher_GetCipherInfo* pGetCipherInfo = NULL;
	int nI;

	// get the selftest function address
	for (nI = 0; nI < NUM_OF_CIPHERS; nI++)
	{
		if (strcmp(cipherFaces[nI].szName, pCipherName) == 0)
		{
			pGetCipherInfo = cipherFaces[nI].pGetCipherInfo;
			break;
		}
	}

	if (pGetCipherInfo == NULL) {
		return CIPHERSERVER_ERROR_CIPHERNOTFOUND;
	}

	// prepare the information block
	pInfoBlock->lSizeOf = sizeof(CIPHERINFOBLOCK);

	// call the retrieved function to get the information block
	switch ((*pGetCipherInfo)(pInfoBlock))
	{
	case CIPHER_ERROR_NOERROR :
		return CIPHERSERVER_ERROR_NOERROR;
	case CIPHER_ERROR_INVALID :
		return CIPHERSERVER_ERROR_INVALIDCIPHER;
	default:
		// unknown return value
		return CIPHERSERVER_ERROR_ERROR;
	}
}

//////////////////////////////////////////////////////////////////////////////

// internal random generator
void CRYPTPAK_CALLCONV _internalRandomGenerator
(WORD8* pTargetBuffer,
 WORD32 lNumOfRandomBytes,
 const void* pData)
{
	// just map the call
	Yarrow_GetData((PYARROWCTX) pData, pTargetBuffer, lNumOfRandomBytes);
}

//////////////////////////////////////////////////////////////////////////////

WORD32 CRYPTPAK_API CipherServer_Create
(const char* pCipherName,
 PCIPHERCTX* pCtxPtr,
 Cipher_RandomGenerator* pRandGenFunc,
 const void* pRandGenData,
 const void* pRandSeed,
 WORD32 lRandSeedLen)
{
	PCIPHERCTX pNewCtx;
	WORD32 lRetCode;
	PYARROWCTX pYrwCtx;
	int nI;

	// create a new cipher context
#ifdef KERNEL_COMPILE
	pNewCtx = ( PCIPHERCTX )ExAllocatePool( NonPagedPool, sizeof(CIPHERCTX)  );
#else
	pNewCtx = (PCIPHERCTX) malloc(sizeof(CIPHERCTX));
#endif
	if (pNewCtx == NULL)
		return CIPHERSERVER_ERROR_OUTOFMEMORY;

	// get the function addresses
	nI = 0;
	while (nI < NUM_OF_CIPHERS)
	{
		if (strcmp(cipherFaces[nI].szName, pCipherName) == 0)
			break;
		nI++;
	}
	if (nI == NUM_OF_CIPHERS)
	{
		return CIPHERSERVER_ERROR_CIPHERNOTFOUND;
	}

	pNewCtx->pGetCipherInfo      = cipherFaces[nI].pGetCipherInfo;
	pNewCtx->pSelfTest           = cipherFaces[nI].pSelfTest;
	pNewCtx->pCreateWorkContext  = cipherFaces[nI].pCreateWorkContext;
	pNewCtx->pResetWorkContext   = cipherFaces[nI].pResetWorkContext;
	pNewCtx->pDestroyWorkContext = cipherFaces[nI].pDestroyWorkContext;
	pNewCtx->pEncryptBuffer      = cipherFaces[nI].pEncryptBuffer;
	pNewCtx->pDecryptBuffer      = cipherFaces[nI].pDecryptBuffer;

	// call the info function to get the information block
	pNewCtx->infoblock.lSizeOf = sizeof(CIPHERINFOBLOCK);
	switch ((*pNewCtx->pGetCipherInfo)(&pNewCtx->infoblock))
	{
	case CIPHER_ERROR_NOERROR :
		lRetCode = CIPHERSERVER_ERROR_NOERROR;
		break;
	case CIPHER_ERROR_INVALID :
		lRetCode = CIPHERSERVER_ERROR_INVALIDCIPHER;
		break;
	default:
		// unknown return value
		lRetCode = CIPHERSERVER_ERROR_ERROR;
	}
	if (lRetCode != CIPHERSERVER_ERROR_NOERROR)
	{
#ifdef KERNEL_COMPILE
		ExFreePool( pNewCtx );
#else
		free(pNewCtx);
#endif
		return lRetCode;
	}

	// must we create our own random generator?
	if (pRandGenFunc == CIPHER_NULL)
	{
		// create a random context
		pYrwCtx = Yarrow_Create(pRandSeed, lRandSeedLen);
		if (pYrwCtx == NULL)
		{
#ifdef KERNEL_COMPILE
			ExFreePool( pNewCtx );
#else
			free(pNewCtx);
#endif
			return CIPHERSERVER_ERROR_OUTOFMEMORY;
		}

		// store the context pointer
		pNewCtx->pRndGenData = pYrwCtx;

		// store the address of the internal random generator
		pNewCtx->pRandomGenerator = _internalRandomGenerator;

		// set the internal usage flag
		pNewCtx->blUseInternalRndGen = BOOL_TRUE;
	}
	else
	{
		// just link to the external random generator
		pNewCtx->pRandomGenerator = pRandGenFunc;
		pNewCtx->pRndGenData = pRandGenData;
		pNewCtx->blUseInternalRndGen = BOOL_FALSE;
	}

	// return the cipher handle
	*pCtxPtr = pNewCtx;

	// cipher successfully loaded
	return CIPHERSERVER_ERROR_NOERROR;
}

//////////////////////////////////////////////////////////////////////////////

WORD32 CRYPTPAK_API CipherServer_Destroy
(PCIPHERCTX pCtx)
{
	// release the internal rng, if necessary
	PYARROWCTX yctx;

	if (pCtx->blUseInternalRndGen)
	{
		yctx = (PYARROWCTX) pCtx->pRndGenData;

		if (yctx)
		{
			Yarrow_Destroy(yctx);
		}
	}

	// clear the context
	memset(pCtx, 0, sizeof(CIPHERCTX));

	// free the context and quit
#ifdef KERNEL_COMPILE
	ExFreePool( pCtx );
#else
	free(pCtx);
#endif
	return CIPHERSERVER_ERROR_NOERROR;
}

//////////////////////////////////////////////////////////////////////////////

WORD32 CRYPTPAK_API CipherServer_ExecuteSelfTest
(PCIPHERCTX pCtx,
 BYTEBOOL blExtendedTest)
{
	void* pTempWorkCtx;
	void* pInitData;
	WORD32 lResult;
	WORD32 lKeySize;
	WORD32 lI, lJ, lModeFlags;
	WORD8* keybuf;
	WORD8* twoblocks;
	int nTestBlockSize;
	int nI;
	PCIPHERSESSION chandle;

	// create a temporary work context
	if (pCtx->infoblock.lContextSize)
	{
#ifdef KERNEL_COMPILE
		pTempWorkCtx = ExAllocatePool( NonPagedPool, pCtx->infoblock.lContextSize  );
#else
		pTempWorkCtx = malloc(pCtx->infoblock.lContextSize);
#endif
		
		if (pTempWorkCtx == NULL)
			return CIPHERSERVER_ERROR_OUTOFMEMORY;
	}
	else
	{
		pTempWorkCtx = NULL;
	}

	// execute the cipher selftest
	if ((*pCtx->pSelfTest)(pTempWorkCtx) == CIPHER_ERROR_INVALID)
	{
		// selftest failed
#ifdef KERNEL_COMPILE
		ExFreePool( pTempWorkCtx );
#else
		free(pTempWorkCtx);
#endif
		return CIPHERSERVER_ERROR_INVALIDCIPHER;
	}

	// free the temporary work context
#ifdef KERNEL_COMPILE
	ExFreePool( pTempWorkCtx );
#else
	free(pTempWorkCtx);
#endif

	// execute the extended test, i.n.
	if (blExtendedTest)
	{
		// test both legacy and standard implementation
		for (lJ = 0, lModeFlags = 0; lJ < 2; lJ++)
		{
			if (lJ)
			{
				lModeFlags |= CIPHER_MODE_FLAG_LEGACY;
			}

			// encrypt two blocks using a simple key
			if (pCtx->infoblock.blOwnHasher)
				lKeySize = FIXKEYLEN;
			else
				lKeySize = pCtx->infoblock.lKeySize;

			if (lKeySize)
			{
#ifdef KERNEL_COMPILE
				keybuf = ( WORD8* )ExAllocatePool( NonPagedPool, lKeySize  );
#else
				keybuf = (WORD8*) malloc(lKeySize);