#include "crypto.h"
#include "rng.h"
#include <stdint.h>
#include <stdbool.h>

#include "inc/hw_types.h"
#include "inc/hw_aes.h"
#include "inc/hw_memmap.h"

#include "driverlib/aes.h"
#include "driverlib/rom_map.h"
#include "driverlib/prcm.h"


static volatile bool g_bContextInIntFlag;
static volatile bool g_bDataInIntFlag;
static volatile bool g_bContextOutIntFlag;
static volatile bool g_bDataOutIntFlag;

static uint32_t aes_mode=AES_CFG_MODE_CBC, aes_key_size, aes_blocksize;
static uint8_t *aes_key;

void AESCrypt(uint32_t dir, uint8_t *iv,
              uint8_t *data, uint8_t *result,
              uint32_t dataLength) {
    // Step1:  Enable Interrupts
    // Step2:  Wait for Context Ready Interrupt
    // Step3:  Set the Configuration Parameters (Direction,AES Mode and Key Size)
    // Step4:  Set the Initialisation Vector
    // Step5:  Write Key
    // Step6:  Start the crypt process

    // Clear the flags.
    g_bContextInIntFlag = false;
    g_bDataInIntFlag = false;
    g_bContextOutIntFlag = false;
    g_bDataOutIntFlag = false;

    // Enable all interrupts.
    MAP_AESIntEnable(AES_BASE, AES_INT_CONTEXT_IN | AES_INT_CONTEXT_OUT | AES_INT_DATA_IN | AES_INT_DATA_OUT);
    // Wait for the context in flag, the flag will be set in the Interrupt handler.
    while(!g_bContextInIntFlag) {}
    MAP_AESConfigSet(AES_BASE, dir | aes_mode | aes_key_size);
    MAP_AESIVSet(AES_BASE, iv);
    MAP_AESKey1Set(AES_BASE, aes_key, aes_key_size);
    MAP_AESDataProcess(AES_BASE, (unsigned char *) data, (unsigned char *) result, dataLength);
}

uint8_t * AESEncrypt(uint8_t *message, uint32_t messageLength) {
    uint32_t msgLen = messageLength + aes_blocksize;
    uint32_t ivblocks = aes_blocksize/4;
    uint8_t *result = (uint8_t *) malloc(msgLen);
    uint32_t *IVrand = (uint32_t *) result;
    uint8_t *data = &result[aes_blocksize];
    int i;
    for(i=0;i<ivblocks;i++) {
        IVrand[i] = random();
    }
    AESCrypt(AES_DIR_ENCRYPT, result, message, data, messageLength);
    return result;
}

// Message must be larger than 2 * aes_key_length (in bytes)
uint8_t * AESDecrypt(uint8_t *message, uint32_t messageLength) {
    uint8_t *data = &message[aes_blocksize];
    uint32_t msgLen = messageLength - aes_blocksize;
    uint8_t *result = (uint8_t *) malloc(msgLen);
    AESCrypt(AES_DIR_DECRYPT, message, data, result, msgLen);
    return result;
}

void AESIntHandler(void){
    uint32_t uiIntStatus;
    // Read the AES masked interrupt status.
    uiIntStatus = MAP_AESIntStatus(AES_BASE, true);

    // Set Different flags depending on the interrupt source.
    if(uiIntStatus & AES_INT_CONTEXT_IN) {
        MAP_AESIntDisable(AES_BASE, AES_INT_CONTEXT_IN);
        g_bContextInIntFlag = true;
    }
    if(uiIntStatus & AES_INT_DATA_IN) {
        MAP_AESIntDisable(AES_BASE, AES_INT_DATA_IN);
        g_bDataInIntFlag = true;
    }
    if(uiIntStatus & AES_INT_CONTEXT_OUT) {
        MAP_AESIntDisable(AES_BASE, AES_INT_CONTEXT_OUT);
        g_bContextOutIntFlag = true;
    }
    if(uiIntStatus & AES_INT_DATA_OUT) {
        MAP_AESIntDisable(AES_BASE, AES_INT_DATA_OUT);
        g_bDataOutIntFlag = true;
    }
}

uint32_t getBlockSize(void) {
    return aes_blocksize;
}

void AESSetup(uint32_t key_size, uint8_t *key){
    aes_key_size = key_size;
    switch(key_size) {
    case AES_KEY_128BIT:
        aes_blocksize = 16;
        break;
    case AES_KEY_192BIT:
        aes_blocksize = 24;
        break;
    case AES_KEY_256BIT:
        aes_blocksize = 32;
        break;
    }
    aes_key = key;
    // Enable AES Module
    MAP_PRCMPeripheralClkEnable(PRCM_DTHE, PRCM_RUN_MODE_CLK);
    // Enable AES interrupts
    MAP_AESIntRegister(AES_BASE, AESIntHandler);
}