// aes.c

#include "pch.h"

#ifdef SCPR

BOOLEAN CipherInitComplete = FALSE;

void
CprAESInit()
{
	g_Data.MaxAesSessions = 256;
	g_Data.CurAesSessions = 0;
	aes_cipher_init();
}

void
CprPortAESInit(
	PCPR_DEVICE_EXTENSION pDevExt
   )
{
	if ((pDevExt->AESKeyLength > 256) || (pDevExt->AESKeyLength <= 0))
	{
		pDevExt->AESKeyLength = 256;
	}

	aes_key_init(
		&pDevExt->AES_Encryption_Key,
		pDevExt->AESKey,
		pDevExt->AESKeyLength
		);

	pDevExt->EncryptLeft = 0;
	pDevExt->DecryptLeft = 0;
	pDevExt->InitRcvLen = 0;
	pDevExt->AESInitVectorFlag = FALSE;
}

void
CprSendInitVector(
	PCPR_DEVICE_EXTENSION pDevExt
	)
{
	CprCreateAESInitVector(pDevExt);
	CprUartWriteDirect(&pDevExt->Uart, pDevExt->EncryptVector, AES_INIT_VECTOR_SIZE, TRUE, FALSE);
	pDevExt->AESInitVectorFlag = TRUE;
}

NTSTATUS
CprCreateAESInitVector (
	PCPR_DEVICE_EXTENSION pDevExt
	)
{
    int i;

	for (i = 0; i < AES_INIT_VECTOR_SIZE; i++)
	{
		pDevExt->EncryptVector[i] = (BYTE)genrand_int32();
		pDevExt->DecryptVector[i] = (BYTE)pDevExt->EncryptVector[i];		
	}

    return STATUS_SUCCESS;
}

NTSTATUS
CprUpdateAESInitVector (
	PCPR_DEVICE_EXTENSION pDevExt,
	BYTE newVector[],
	int nAesBytes
	)
{
    NTSTATUS Status = STATUS_UNSUCCESSFUL;
    int i;

	//DbgPrint("KeyLen %d MaxAesSessions %d CurAesSessions %d\n", pDevExt->AESKeyLength, g_Data.MaxAesSessions, g_Data.CurAesSessions);
	//for (i = 0; i < 32; i++)
	//{
 //       DbgPrint(" %02X", pDevExt->KeyBytes[i]);
	//}
	//DbgPrint("\n");

	for (i = 0; (i < nAesBytes) && (pDevExt->InitRcvLen < AES_INIT_VECTOR_SIZE); i++)
	{
		pDevExt->EncryptVector[pDevExt->InitRcvLen] = newVector[i];
		pDevExt->DecryptVector[pDevExt->InitRcvLen] = newVector[i];
        //DbgPrint(" %02X", newVector[i]);
		pDevExt->InitRcvLen++;
	}
	//DbgPrint("\n");

	if (pDevExt->InitRcvLen == AES_INIT_VECTOR_SIZE)
	{
		pDevExt->AESInitVectorFlag = TRUE;
	}

	Status = STATUS_SUCCESS;
    return Status;
}

void CprEncrypt(
	PCPR_DEVICE_EXTENSION	pDevExt,
    UCHAR					*input,
    UCHAR					*output,
	int						length
	)
{
	aes_byte_encrypt(&pDevExt->AES_Encryption_Key,
		&pDevExt->EncryptVector[0],
		input,
		length,
		output,
		(PBYTE)&pDevExt->EncryptLeft,
		1);
}

void CprDecrypt(
	PCPR_DEVICE_EXTENSION	pDevExt,
    UCHAR					*input,
    UCHAR					*output,
	int						length
	)
{
	aes_byte_decrypt(&pDevExt->AES_Encryption_Key,
		&pDevExt->DecryptVector[0],
		input,
		length,
		output,
		(PBYTE)&pDevExt->DecryptLeft,
		2);
}

#endif // SCPR