#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define MBEDTLS

#ifdef MBEDTLS

#include <mbedtls/aes.h>    // might need to be in double-quotes
#include <mbedtls/config.h>

#else

#include <tiny-AES-c/aes.h>

#endif

void print_buffer(uint8_t *buffer, size_t buff_len)
{
	size_t i;
	for (i = 0; i < buff_len; ++i) {
		putchar(buffer[i]);
	}
}

// TODO: write generic AES function that uses mbedtls or tiny-AES-c
bool aes_128_enc_ecb(uint8_t key[16], uint8_t *plaintext, size_t plaintext_len,
                                    uint8_t *ciphertext, size_t ciphertext_len)
{
	size_t i;
	mbedtls_aes_context aes_ctx;

	if (plaintext_len > ciphertext_len) {
		fprintf(stderr, "Plaintext longer than ciphertext buffer\n");
		return false;
	}

	if (mbedtls_aes_setkey_enc(&aes_ctx, key, 128) != 0) {
		fprintf(stderr, "Setting AES encryption key\n");
		return false;
	}

	for (i = 0; i < plaintext_len; i = i + 16) {
		// TODO: handle ciphertext not block size-aligned
		if (mbedtls_aes_crypt_ecb(&aes_ctx, MBEDTLS_AES_ENCRYPT,
					  plaintext + i, ciphertext + i) != 0) {
			fprintf(stderr, "Performing AES encryption\n");
			return false;
		}
	}

	return true;
}


int main()
{
	// for testing (TODO: take key, plaintext as args)
	uint8_t key[16];
	uint8_t message[32];
	uint8_t ciph_out[48];
	memset(key, 'B', sizeof(key));
	memset(message, 'A', sizeof(message));
	memset(ciph_out, 0, sizeof(ciph_out));

	if (!aes_128_enc_ecb(key, message, sizeof(message),
	                           ciph_out, sizeof(ciph_out))) {
		return EXIT_FAILURE;
	}

	print_buffer(ciph_out, sizeof(ciph_out));

	return EXIT_SUCCESS;                              
}