/**
 * AES abstraction implementation using Windows CNG (Cryptography API: Next Generation).
 *
 * Requires: Windows Vista / Windows Server 2008 or later
 * Links: bcrypt.lib
 */

#include "brp_lib/aes.h"
#include <windows.h>
#include <bcrypt.h>
#include <stdlib.h>
#include <string.h>


/**
 * Actual key context struct (backend-specific).
 * Stores BCrypt algorithm provider handles and key object.
 * We need separate algorithm handles for ECB and CBC modes.
 */
struct aes_ctx_t {
	BCRYPT_ALG_HANDLE hAlgorithmECB;
	BCRYPT_ALG_HANDLE hAlgorithmCBC;
	BCRYPT_KEY_HANDLE hKeyECB;
	BCRYPT_KEY_HANDLE hKeyCBC;
};


aes_ctx aes_create_ctx(const unsigned char key[AES_KEY_SIZE])
{
	NTSTATUS status;
	aes_ctx ctx = NULL;

	// Validate input
	if (key == NULL)
		return NULL;

	// Allocate context
	ctx = (aes_ctx)malloc(sizeof(struct aes_ctx_t));
	if (ctx == NULL)
		return NULL;

	memset(ctx, 0, sizeof(struct aes_ctx_t));

	// Open AES algorithm provider for ECB mode
	status = BCryptOpenAlgorithmProvider(
		&ctx->hAlgorithmECB,
		BCRYPT_AES_ALGORITHM,
		NULL,  // Default provider
		0      // No flags
	);
	if (!BCRYPT_SUCCESS(status))
	{
		aes_destroy_ctx(ctx);
		return NULL;
	}

	// Set ECB mode
	status = BCryptSetProperty(
		ctx->hAlgorithmECB,
		BCRYPT_CHAINING_MODE,
		(PUCHAR)BCRYPT_CHAIN_MODE_ECB,
		(ULONG)(wcslen(BCRYPT_CHAIN_MODE_ECB) + 1) * sizeof(WCHAR),
		0
	);
	if (!BCRYPT_SUCCESS(status))
	{
		aes_destroy_ctx(ctx);
		return NULL;
	}

	// Generate ECB key object
	status = BCryptGenerateSymmetricKey(
		ctx->hAlgorithmECB,
		&ctx->hKeyECB,
		NULL,  // Let BCrypt allocate key object
		0,
		(PUCHAR)key,
		AES_KEY_SIZE,
		0      // No flags
	);
	if (!BCRYPT_SUCCESS(status))
	{
		aes_destroy_ctx(ctx);
		return NULL;
	}

	// Open AES algorithm provider for CBC mode
	status = BCryptOpenAlgorithmProvider(
		&ctx->hAlgorithmCBC,
		BCRYPT_AES_ALGORITHM,
		NULL,  // Default provider
		0      // No flags
	);
	if (!BCRYPT_SUCCESS(status))
	{
		aes_destroy_ctx(ctx);
		return NULL;
	}

	// Set CBC mode
	status = BCryptSetProperty(
		ctx->hAlgorithmCBC,
		BCRYPT_CHAINING_MODE,
		(PUCHAR)BCRYPT_CHAIN_MODE_CBC,
		(ULONG)(wcslen(BCRYPT_CHAIN_MODE_CBC) + 1) * sizeof(WCHAR),
		0
	);
	if (!BCRYPT_SUCCESS(status))
	{
		aes_destroy_ctx(ctx);
		return NULL;
	}

	// Generate CBC key object
	status = BCryptGenerateSymmetricKey(
		ctx->hAlgorithmCBC,
		&ctx->hKeyCBC,
		NULL,  // Let BCrypt allocate key object
		0,
		(PUCHAR)key,
		AES_KEY_SIZE,
		0      // No flags
	);
	if (!BCRYPT_SUCCESS(status))
	{
		aes_destroy_ctx(ctx);
		return NULL;
	}

	return ctx;
}


int aes_encrypt_block(
	aes_ctx ctx,
	const unsigned char plaintext[AES_BLOCK_SIZE],
	unsigned char ciphertext[AES_BLOCK_SIZE])
{
	NTSTATUS status;
	ULONG cbResult;

	// Validate inputs
	if (ctx == NULL || plaintext == NULL || ciphertext == NULL)
		return AES_ERR_INVALID_PARAM;

	// Encrypt single block using ECB mode (no IV needed)
	status = BCryptEncrypt(
		ctx->hKeyECB,
		(PUCHAR)plaintext,
		AES_BLOCK_SIZE,
		NULL,  // No padding info
		NULL,  // No IV for ECB
		0,
		ciphertext,
		AES_BLOCK_SIZE,
		&cbResult,
		0      // No padding
	);

	if (!BCRYPT_SUCCESS(status))
	{
		// Clear output buffer on error
		SecureZeroMemory(ciphertext, AES_BLOCK_SIZE);
		return AES_ERR_CRYPTO_FAILURE;
	}

	// Verify output length
	if (cbResult != AES_BLOCK_SIZE)
	{
		SecureZeroMemory(ciphertext, AES_BLOCK_SIZE);
		return AES_ERR_CRYPTO_FAILURE;
	}

	return AES_OK;
}


int aes_cbc_encrypt(
	aes_ctx ctx,
	const unsigned char * plaintext,
	size_t plaintext_len,
	unsigned char * ciphertext,
	const unsigned char iv[AES_BLOCK_SIZE])
{
	NTSTATUS status;
	ULONG cbResult;
	unsigned char iv_copy[AES_BLOCK_SIZE];

	// Validate inputs
	if (ctx == NULL || plaintext == NULL || ciphertext == NULL || iv == NULL)
		return AES_ERR_INVALID_PARAM;

	// Validate length is non-zero and block-aligned
	if (plaintext_len == 0 || plaintext_len % AES_BLOCK_SIZE != 0)
		return AES_ERR_INVALID_LENGTH;

	// Check for size_t to ULONG overflow (ULONG is 32-bit on Windows)
	if (plaintext_len > ULONG_MAX)
		return AES_ERR_INVALID_LENGTH;

	// BCrypt modifies the IV, so we need to make a copy
	memcpy(iv_copy, iv, AES_BLOCK_SIZE);

	// Encrypt using pre-configured CBC mode
	status = BCryptEncrypt(
		ctx->hKeyCBC,
		(PUCHAR)plaintext,
		(ULONG)plaintext_len,
		NULL,      // No padding info
		iv_copy,   // IV (will be modified)
		AES_BLOCK_SIZE,
		ciphertext,
		(ULONG)plaintext_len,
		&cbResult,
		0          // No padding (data must be block-aligned)
	);

	// Clear IV copy from stack
	SecureZeroMemory(iv_copy, AES_BLOCK_SIZE);

	if (!BCRYPT_SUCCESS(status))
	{
		// Clear output buffer on error
		SecureZeroMemory(ciphertext, plaintext_len);
		return AES_ERR_CRYPTO_FAILURE;
	}

	// Verify output length matches input
	if (cbResult != (ULONG)plaintext_len)
	{
		SecureZeroMemory(ciphertext, plaintext_len);
		return AES_ERR_CRYPTO_FAILURE;
	}

	return AES_OK;
}


int aes_cbc_decrypt(
	aes_ctx ctx,
	const unsigned char * ciphertext,
	size_t ciphertext_len,
	unsigned char * plaintext,
	const unsigned char iv[AES_BLOCK_SIZE])
{
	NTSTATUS status;
	ULONG cbResult;
	unsigned char iv_copy[AES_BLOCK_SIZE];

	// Validate inputs
	if (ctx == NULL || ciphertext == NULL || plaintext == NULL || iv == NULL)
		return AES_ERR_INVALID_PARAM;

	// Validate length is non-zero and block-aligned
	if (ciphertext_len == 0 || ciphertext_len % AES_BLOCK_SIZE != 0)
		return AES_ERR_INVALID_LENGTH;

	// Check for size_t to ULONG overflow (ULONG is 32-bit on Windows)
	if (ciphertext_len > ULONG_MAX)
		return AES_ERR_INVALID_LENGTH;

	// BCrypt modifies the IV, so we need to make a copy
	memcpy(iv_copy, iv, AES_BLOCK_SIZE);

	// Decrypt using pre-configured CBC mode
	status = BCryptDecrypt(
		ctx->hKeyCBC,
		(PUCHAR)ciphertext,
		(ULONG)ciphertext_len,
		NULL,      // No padding info
		iv_copy,   // IV (will be modified)
		AES_BLOCK_SIZE,
		plaintext,
		(ULONG)ciphertext_len,
		&cbResult,
		0          // No padding (data must be block-aligned)
	);

	// Clear IV copy from stack
	SecureZeroMemory(iv_copy, AES_BLOCK_SIZE);

	if (!BCRYPT_SUCCESS(status))
	{
		// Clear output buffer on error
		SecureZeroMemory(plaintext, ciphertext_len);
		return AES_ERR_CRYPTO_FAILURE;
	}

	// Verify output length matches input
	if (cbResult != (ULONG)ciphertext_len)
	{
		SecureZeroMemory(plaintext, ciphertext_len);
		return AES_ERR_CRYPTO_FAILURE;
	}

	return AES_OK;
}


int aes_cbc_mac(
	aes_ctx ctx,
	const unsigned char * plaintext,
	size_t plaintext_len,
	unsigned char mac[AES_BLOCK_SIZE])
{
	NTSTATUS status;
	ULONG cbResult;
	unsigned char * temp_buffer = NULL;
	unsigned char iv_copy[AES_BLOCK_SIZE];

	// Validate inputs
	if (ctx == NULL || plaintext == NULL || mac == NULL)
		return AES_ERR_INVALID_PARAM;

	// Validate length is at least one block and block-aligned
	if (plaintext_len < AES_BLOCK_SIZE || plaintext_len % AES_BLOCK_SIZE != 0)
		return AES_ERR_INVALID_LENGTH;

	// Check for size_t to ULONG overflow (ULONG is 32-bit on Windows)
	if (plaintext_len > ULONG_MAX)
		return AES_ERR_INVALID_LENGTH;

	// CBC-MAC: Encrypt in CBC mode and keep only the last block

	// Make a copy of the IV (mac serves as both input IV and output MAC)
	memcpy(iv_copy, mac, AES_BLOCK_SIZE);

	// Allocate temporary buffer for ciphertext
	temp_buffer = (unsigned char *)malloc(plaintext_len);
	if (temp_buffer == NULL)
	{
		SecureZeroMemory(iv_copy, AES_BLOCK_SIZE);
		return AES_ERR_CRYPTO_FAILURE;
	}

	// Encrypt in CBC mode using pre-configured CBC handle
	status = BCryptEncrypt(
		ctx->hKeyCBC,
		(PUCHAR)plaintext,
		(ULONG)plaintext_len,
		NULL,        // No padding info
		iv_copy,     // IV (will be modified)
		AES_BLOCK_SIZE,
		temp_buffer,
		(ULONG)plaintext_len,
		&cbResult,
		0            // No padding
	);

	// Clear IV copy from stack
	SecureZeroMemory(iv_copy, AES_BLOCK_SIZE);

	if (BCRYPT_SUCCESS(status) && cbResult == (ULONG)plaintext_len)
	{
		// Copy last block to MAC output
		memcpy(mac, temp_buffer + plaintext_len - AES_BLOCK_SIZE, AES_BLOCK_SIZE);
	}
	else
	{
		// Clear MAC output on error
		SecureZeroMemory(mac, AES_BLOCK_SIZE);
	}

	// Securely wipe temporary buffer before freeing
	SecureZeroMemory(temp_buffer, plaintext_len);
	free(temp_buffer);

	return (BCRYPT_SUCCESS(status) && cbResult == (ULONG)plaintext_len) ? AES_OK : AES_ERR_CRYPTO_FAILURE;
}


void aes_destroy_ctx(aes_ctx ctx)
{
	if (ctx != NULL)
	{
		if (ctx->hKeyECB != NULL)
			BCryptDestroyKey(ctx->hKeyECB);

		if (ctx->hKeyCBC != NULL)
			BCryptDestroyKey(ctx->hKeyCBC);

		if (ctx->hAlgorithmECB != NULL)
			BCryptCloseAlgorithmProvider(ctx->hAlgorithmECB, 0);

		if (ctx->hAlgorithmCBC != NULL)
			BCryptCloseAlgorithmProvider(ctx->hAlgorithmCBC, 0);

		// Securely wipe the context
		SecureZeroMemory(ctx, sizeof(struct aes_ctx_t));
		free(ctx);
	}
}
