/**
 * AES abstraction implementation using macOS CommonCrypto.
 *
 * Requires: macOS 10.7 or later
 * Framework: Security.framework (included by default)
 */

#include "brp_lib/aes.h"
#include <CommonCrypto/CommonCryptor.h>
#include <stdlib.h>
#include <string.h>
#include <limits.h>

/**
 * Secure memory zeroing (cannot be optimized away).
 * macOS doesn't provide secure_zero, so we implement it here.
 */
static void secure_zero(void *ptr, size_t len)
{
	if (ptr == NULL || len == 0)
		return;

	// Use memset_s if available (C11 Annex K)
	#if defined(__STDC_LIB_EXT1__)
		memset_s(ptr, len, 0, len);
	#else
		// Fallback: Use volatile pointer to prevent optimization
		volatile unsigned char *p = (volatile unsigned char *)ptr;
		while (len--) {
			*p++ = 0;
		}
	#endif
}


/**
 * Actual key context struct (backend-specific).
 * CommonCrypto doesn't require persistent key schedule storage,
 * but we store the key for use in operations.
 */
struct aes_ctx_t {
	unsigned char key[AES_KEY_SIZE];
};


aes_ctx aes_create_ctx(const unsigned char key[AES_KEY_SIZE])
{
	aes_ctx ctx;

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

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

	// Zero entire struct to avoid uninitialized padding bytes
	memset(ctx, 0, sizeof(struct aes_ctx_t));

	// Store key for later use
	memcpy(ctx->key, key, AES_KEY_SIZE);

	return ctx;
}


int aes_encrypt_block(
	aes_ctx ctx,
	const unsigned char plaintext[AES_BLOCK_SIZE],
	unsigned char ciphertext[AES_BLOCK_SIZE])
{
	size_t dataOutMoved;

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

	// Use ECB mode for single block encryption
	CCCryptorStatus status = CCCrypt(
		kCCEncrypt,
		kCCAlgorithmAES,
		kCCOptionECBMode,  // ECB mode
		ctx->key,
		kCCKeySizeAES128,
		NULL,              // No IV for ECB
		plaintext,
		AES_BLOCK_SIZE,
		ciphertext,
		AES_BLOCK_SIZE,
		&dataOutMoved
	);

	if (status != kCCSuccess)
	{
		// Clear output buffer on error
		secure_zero(ciphertext, AES_BLOCK_SIZE);
		return AES_ERR_CRYPTO_FAILURE;
	}

	// Verify output length
	if (dataOutMoved != AES_BLOCK_SIZE)
	{
		secure_zero(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])
{
	size_t dataOutMoved;
	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 integer overflow (size_t to int conversion safety)
	if (plaintext_len > INT_MAX)
		return AES_ERR_INVALID_LENGTH;

	// CommonCrypto may modify the IV, so make a copy
	memcpy(iv_copy, iv, AES_BLOCK_SIZE);

	// Use CBC mode (default when kCCOptionECBMode is not set)
	CCCryptorStatus status = CCCrypt(
		kCCEncrypt,
		kCCAlgorithmAES,
		0,                 // No options = CBC mode, no padding
		ctx->key,
		kCCKeySizeAES128,
		iv_copy,
		plaintext,
		plaintext_len,
		ciphertext,
		plaintext_len,
		&dataOutMoved
	);

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

	if (status != kCCSuccess)
	{
		// Clear output buffer on error
		secure_zero(ciphertext, plaintext_len);
		return AES_ERR_CRYPTO_FAILURE;
	}

	// Verify output length matches input
	if (dataOutMoved != plaintext_len)
	{
		secure_zero(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])
{
	size_t dataOutMoved;
	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 integer overflow (size_t to int conversion safety)
	if (ciphertext_len > INT_MAX)
		return AES_ERR_INVALID_LENGTH;

	// CommonCrypto may modify the IV, so make a copy
	memcpy(iv_copy, iv, AES_BLOCK_SIZE);

	// Use CBC mode (default when kCCOptionECBMode is not set)
	CCCryptorStatus status = CCCrypt(
		kCCDecrypt,
		kCCAlgorithmAES,
		0,                 // No options = CBC mode, no padding
		ctx->key,
		kCCKeySizeAES128,
		iv_copy,
		ciphertext,
		ciphertext_len,
		plaintext,
		ciphertext_len,
		&dataOutMoved
	);

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

	if (status != kCCSuccess)
	{
		// Clear output buffer on error
		secure_zero(plaintext, ciphertext_len);
		return AES_ERR_CRYPTO_FAILURE;
	}

	// Verify output length matches input
	if (dataOutMoved != ciphertext_len)
	{
		secure_zero(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])
{
	unsigned char * temp_buffer = NULL;
	size_t dataOutMoved;
	CCCryptorStatus status;
	unsigned char mac_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 integer overflow (size_t to int conversion safety)
	if (plaintext_len > INT_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(mac_copy, mac, AES_BLOCK_SIZE);

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

	// Encrypt in CBC mode using mac as IV
	status = CCCrypt(
		kCCEncrypt,
		kCCAlgorithmAES,
		0,                 // CBC mode, no padding
		ctx->key,
		kCCKeySizeAES128,
		mac_copy,          // Use mac copy as IV input
		plaintext,
		plaintext_len,
		temp_buffer,
		plaintext_len,
		&dataOutMoved
	);

	// Clear mac copy from stack
	secure_zero(mac_copy, AES_BLOCK_SIZE);

	if (status == kCCSuccess && dataOutMoved == 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
		secure_zero(mac, AES_BLOCK_SIZE);
	}

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

	return (status == kCCSuccess && dataOutMoved == plaintext_len) ? AES_OK : AES_ERR_CRYPTO_FAILURE;
}


void aes_destroy_ctx(aes_ctx ctx)
{
	if (ctx != NULL)
	{
		// Securely wipe the key using secure_zero (cannot be optimized away)
		secure_zero(ctx->key, AES_KEY_SIZE);
		free(ctx);
	}
}
