/*
 *--------------------------------------
 * Program Name:
 * Author:
 * License:
 * Description:
 *--------------------------------------
*/

#include <stdint.h>
#include <stdlib.h>
#include <stdbool.h>
#include <string.h>
#include <stdio.h>
#include <encrypt.h>
#include <alloca.h>

#define CEMU_CONSOLE ((char*)0xFB0000)
char *msg = "The lazy fox jumped over the dog!";
char *secondmsg = "The dog got angry and barked";
#define KEYSIZE (256>>3)    // 256 bits converted to bytes

// we will reuse the key and IV for the sake of demo
// in an actual use-case, both key and IV should be generated by a secure RNG
// and also probably shouldn't be globals
uint8_t key[KEYSIZE] = {
	0xEE,0x89,0x19,0xC3,0x8D,0x53,0x7A,0xD6,0x04,0x19,0x9E,0x77,0x0B,0xE0,0xE0,0x4C,0x4C,0x70,0xDB,0xE1,0x22,0x79,0xE1,0x90,0x06,0x1B,0xAF,0x99,0x49,0x8E,0x66,0x73
};
uint8_t iv[CRYPTX_AES_IV_SIZE] = {
	0x79,0xA6,0xDE,0xDF,0xF0,0xA2,0x7C,0x7F,0xEE,0x0B,0x8E,0xF5,0x12,0x63,0xA4,0x8A
};
aes_error_t error;
struct cryptx_aes_ctx ctx;

void hexdump(uint8_t *addr, size_t len, char *label){
    if(label) sprintf(CEMU_CONSOLE, "\n%s\n", label);
    else sprintf(CEMU_CONSOLE, "\n");
    for(size_t rem_len = len, ct=1; rem_len>0; rem_len--, addr++, ct++){
        sprintf(CEMU_CONSOLE, "\\x%02X", *addr);
        if(!(ct%CRYPTX_AES_BLOCK_SIZE)) sprintf(CEMU_CONSOLE, "\n");
    }
    sprintf(CEMU_CONSOLE, "\n");
}


void demo_aes_cbc(void){

	// declare buffers
	size_t ctlen = strlen(msg)/CRYPTX_AES_BLOCK_SIZE*CRYPTX_AES_BLOCK_SIZE+1;
	uint8_t *ebuf = alloca(ctlen);
	uint8_t *dbuf = alloca(ctlen);

	sprintf(CEMU_CONSOLE, "\n-----------------------------------\nCyclic Block Chain (CBC) mode\n\n");

	// init AES context and echo status
	error = cryptx_aes_init(&ctx, key, KEYSIZE, iv, CRYPTX_AES_IV_SIZE, CRYPTX_AES_CBC_FLAGS(0));
	sprintf(CEMU_CONSOLE, "cbc init complete, exit code %u\n", error);

	// encrypt string and echo status
	memset(ebuf, 0, ctlen);
	error = cryptx_aes_encrypt(&ctx, msg, strlen(msg), ebuf);
	sprintf(CEMU_CONSOLE, "cbc encryption done, exit code %u\n", error);
	hexdump(ebuf, ctlen, "-- encrypted msg --");

	// reset context for decryption and echo status
	error = cryptx_aes_init(&ctx, key, KEYSIZE, iv, CRYPTX_AES_IV_SIZE, CRYPTX_AES_CBC_FLAGS(0));
	sprintf(CEMU_CONSOLE, "cbc init complete, exit code %u\n", error);

	// encrypt string and echo status
	memset(dbuf, 0, ctlen);
	error = cryptx_aes_decrypt(&ctx, ebuf, ctlen, dbuf);
	sprintf(CEMU_CONSOLE, "cbc decryption done, exit code %u\n", error);
	sprintf(CEMU_CONSOLE, "%s\n", dbuf);
}

void demo_aes_ctr(void){

	// declare buffers
	size_t ctlen = strlen(msg);
	uint8_t *ebuf = alloca(ctlen);
	uint8_t *dbuf = alloca(ctlen);

	sprintf(CEMU_CONSOLE, "\n-----------------------------------\nCounter (CTR) mode\n\n");

	// init AES context and echo status
	error = cryptx_aes_init(&ctx, key, KEYSIZE, iv, CRYPTX_AES_IV_SIZE, CRYPTX_AES_CTR_FLAGS(8,8));
	sprintf(CEMU_CONSOLE, "ctr init complete, exit code %u\n", error);

	// encrypt string and echo status
	memset(ebuf, 0, ctlen);
	error = cryptx_aes_encrypt(&ctx, msg, strlen(msg), ebuf);
	sprintf(CEMU_CONSOLE, "ctr encryption done, exit code %u\n", error);
	hexdump(ebuf, ctlen, "-- encrypted msg --");

	// reset context for decryption and echo status
	error = cryptx_aes_init(&ctx, key, KEYSIZE, iv, CRYPTX_AES_IV_SIZE, CRYPTX_AES_CTR_FLAGS(8,8));
	sprintf(CEMU_CONSOLE, "ctr init complete, exit code %u\n", error);

	// encrypt string and echo status
	memset(dbuf, 0, ctlen);
	error = cryptx_aes_decrypt(&ctx, ebuf, ctlen, dbuf);
	sprintf(CEMU_CONSOLE, "ctr decryption done, exit code %u\n", error);
	sprintf(CEMU_CONSOLE, "%s\n", dbuf);
}

void demo_aes_gcm(void){

	// declare buffers
	size_t ctlen = strlen(msg);
	char *associated = "Some random header";
	uint8_t *ebuf = alloca(ctlen);
	uint8_t *dbuf = alloca(ctlen);
	uint8_t odigest[CRYPTX_AES_BLOCK_SIZE], vdigest[CRYPTX_AES_BLOCK_SIZE];

	sprintf(CEMU_CONSOLE, "\n-----------------------------------\nGalois Counter (GCM) mode\n\n");

	// init AES context and echo status
	error = cryptx_aes_init(&ctx, key, KEYSIZE, iv, 13, CRYPTX_AES_GCM_FLAGS);
	sprintf(CEMU_CONSOLE, "gcm init complete, exit code %u\n", error);

	// update context for associated data
	cryptx_aes_update_aad(&ctx, associated, strlen(associated));

	// encrypt string and echo status
	memset(ebuf, 0, ctlen);
	error = cryptx_aes_encrypt(&ctx, msg, strlen(msg), ebuf);
	sprintf(CEMU_CONSOLE, "gcm encryption done, exit code %u\n", error);
	hexdump(ebuf, ctlen, "-- encrypted msg --");

	// now echo digest/auth tag
	error = cryptx_aes_digest(&ctx, odigest);
	sprintf(CEMU_CONSOLE, "gcm digest return done, exit code %u\n", error);
	hexdump(odigest, CRYPTX_AES_BLOCK_SIZE, "-- digest of aad + ciphertext --");

	// reset context for decryption and echo status
	error = cryptx_aes_init(&ctx, key, KEYSIZE, iv, 13, CRYPTX_AES_GCM_FLAGS);
	sprintf(CEMU_CONSOLE, "gcm init complete, exit code %u\n", error);

	// #######################################################
	// to test an invalid ciphertext uncomment the line below
	// ebuf[10] ^= 0xff;

	// this is a macro to do a few things. See the header file
	if(!cryptx_aes_verify(&ctx, associated, strlen(associated), ebuf, ctlen, odigest)){
		sprintf(CEMU_CONSOLE, "auth tag invalid. not decrypting.\n");
		return;
	}
	else
		sprintf(CEMU_CONSOLE, "auth tag valid. proceeding.\n");

	// update context for associated data
	cryptx_aes_update_aad(&ctx, associated, strlen(associated));

	// decrypt string and echo status
	memset(dbuf, 0, ctlen);
	error = cryptx_aes_decrypt(&ctx, ebuf, ctlen, dbuf);
	sprintf(CEMU_CONSOLE, "gcm decryption done, exit code %u\n", error);
	sprintf(CEMU_CONSOLE, "%s\n", dbuf);



	// now echo digest/auth tag
	error = cryptx_aes_digest(&ctx, vdigest);
	sprintf(CEMU_CONSOLE, "gcm digest return done, exit code %u\n", error);
	hexdump(vdigest, CRYPTX_AES_BLOCK_SIZE, "-- digest of aad + ciphertext --");
}



int main(void)
{
    sprintf(CEMU_CONSOLE, "\n------------------------------\nCryptX AES Demo\n------------------------------\n");

   // sprintf(CEMU_CONSOLE, "\n----- CTR Mode -----\n");

    // generate random key and IV
    if(!cryptx_csrand_init(SAMPLING_FAST)) return 1;          // <<<----- DONT FORGET THIS
    // !!!! NEVER PROCEED WITH ANYTHING CRYPTOGRAPHIC !!!!
    // !!!! IF THE CSRNG FAILS TO INIT !!!!


    // show the IV and key for testing purposes
	sprintf(CEMU_CONSOLE, "\n-----------------------------------\nSecrets for Testing\n\n");
    hexdump(key, KEYSIZE, "-- aes key --");
    hexdump(iv, CRYPTX_AES_IV_SIZE, "-- initialization vector --");

	demo_aes_cbc();
	demo_aes_ctr();
	demo_aes_gcm();

}