/*
 * Copyright (c) 2016, Freescale Semiconductor, Inc.
 * Copyright 2016-2018 NXP
 * All rights reserved.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 */

/*******************************************************************************
 * Includes
 ******************************************************************************/

#include "fsl_device_registers.h"
#include "fsl_debug_console.h"
#include "board.h"

#include "fsl_hashcrypt.h"
#include "fsl_puf.h"

#include <string.h>

#include "pin_mux.h"
#include "clock_config.h"
/*******************************************************************************
 * Definitions
 ******************************************************************************/

#define TEST_ASSERT(a, str)  \
    if (!(a))                \
    {                        \
        PRINTF("%s error\r\n", str); \
        do                   \
        {                    \
        } while (1);         \
    }

#define PUF_DISCHARGE_TIME 400
#define PUF_INTRINSIC_KEY_SIZE 16
		
/*******************************************************************************
 * Prototypes
 ******************************************************************************/

/*******************************************************************************
 * Variables
 ******************************************************************************/

/*******************************************************************************
 * Code
 ******************************************************************************/

/*!
 * @brief initialize the PUF 
 */
status_t puf_init(uint8_t * pKey, uint32_t keySize)
{
    status_t result;
    uint8_t activationCode[PUF_ACTIVATION_CODE_SIZE];
    uint8_t keyCode0[PUF_GET_KEY_CODE_SIZE_FOR_KEY_SIZE(32)];
    uint8_t tempKey[16];
	
    /* Initialize random number generator used to generate key mask for HW key */
    /* In real application, the seed shall be obtained from a hardware random number generator. */
    srand(0xbabadeda);

    /* PUF SRAM Configuration*/
    puf_config_t conf;
    PUF_GetDefaultConfig(&conf);

		/* deinit PUF for SRAM discharge */
    PUF_Deinit(PUF, &conf);
		
    /* Initialize PUF peripheral */
    result = PUF_Init(PUF, &conf);
    if (result != kStatus_Success)
    {
        PRINTF("Error Initializing PUF!\r\n");
				return result;
    }

    /* Perform enroll to get device specific PUF activation code */
    /* Note: Enroll operation is usually performed only once for each device. */
    /* Activation code is stored and used in Start operation */
    result = PUF_Enroll(PUF, activationCode, sizeof(activationCode));
    if (result != kStatus_Success)
    {
        PRINTF("Error during Enroll!\r\n");
				return result;
    }
    PUF_Deinit(PUF, &conf);

    /* Reinitialize PUF after enroll */
    result = PUF_Init(PUF, &conf);
    if (result != kStatus_Success)
    {
        PRINTF("Error Initializing PUF!\r\n");
				return result;
    }

    /* Start PUF by loading generated activation code */
    result = PUF_Start(PUF, activationCode, sizeof(activationCode));
    if (result != kStatus_Success)
    {
        PRINTF("Error during Start !\r\n");
				return result;
    }

		/* byte reverse key */
		for (int i = 0; i < keySize; i++)	{
			tempKey[i] = pKey[keySize-1-i];
		}
		
    /* Create keycode for user key with index 0 */
    /* Index 0 selects that the key shall be ouptut (by PUF_GetHwKey()) to a SoC specific private hardware bus. */
    result = PUF_SetUserKey(PUF, kPUF_KeyIndex_00, tempKey, 32, keyCode0, sizeof(keyCode0));
    if (result != kStatus_Success)
    {
        PRINTF("Error setting user key!\r\n");
				return result;
    }

    /* Reconstruct key from keyCode0 to HW bus for crypto module */
    result = PUF_GetHwKey(PUF, keyCode0, sizeof(keyCode0), kPUF_KeySlot0, rand());
    if (result != kStatus_Success)
    {
        PRINTF("Error reconstructing key to HW bus!\r\n");
				return result;
    }
		return 0;
}

void TestAesPufKeyEcb(void)
{
    static const uint8_t keyAes128[] __attribute__((aligned)) = {0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6,
                                                                 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c};
    static const uint8_t plainAes128[]                        = {0x6b, 0xc1, 0xbe, 0xe2, 0x2e, 0x40, 0x9f, 0x96,
                                          0xe9, 0x3d, 0x7e, 0x11, 0x73, 0x93, 0x17, 0x2a};
    static const uint8_t cipherAes128[]                       = {0x3a, 0xd7, 0x7b, 0xb4, 0x0d, 0x7a, 0x36, 0x60,
                                           0xa8, 0x9e, 0xca, 0xf3, 0x24, 0x66, 0xef, 0x97};
    uint8_t cipher[16];
    uint8_t output[16];
    status_t status;
		hashcrypt_handle_t m_handle;

		/* secret key PUF, 128-bit key */
		SYSCTL0->AESKEY_SRCSEL = 0x0;
    m_handle.keyType = kHASHCRYPT_SecretKey;
    m_handle.keySize = kHASHCRYPT_Aes128;
																					 
		/* PUF init */
		status = puf_init((uint8_t *) keyAes128, 16);
    TEST_ASSERT(0 == status, "PUF");
	
    status = HASHCRYPT_AES_EncryptEcb(HASHCRYPT, &m_handle, plainAes128, cipher, 16);
    TEST_ASSERT(kStatus_Success == status, "ECB");
    TEST_ASSERT(memcmp(cipher, cipherAes128, 16) == 0, "ECB");

    status = HASHCRYPT_AES_DecryptEcb(HASHCRYPT, &m_handle, cipher, output, 16);
    TEST_ASSERT(kStatus_Success == status, "ECB");
    TEST_ASSERT(memcmp(output, plainAes128, 16) == 0, "ECB");

    PRINTF("AES ECB Test - 128-bit PUF key - pass\r\n");
}

void TestAesPufKeyCbc(void)
{
    static const uint8_t keyAes192[] __attribute__((aligned)) = {
		 0x17, 0x16, 0x15, 0x14, 0x13, 0x12, 0x11, 0x10,
		 0x0F, 0x0E, 0x0D, 0x0C, 0x0B, 0x0A, 0x09, 0x08, 0x07, 0x06, 0x05, 0x04, 0x03, 0x02, 0x01, 0x00,
		};

    static const uint8_t plainAes128[]                        = {0x6b, 0xc1, 0xbe, 0xe2, 0x2e, 0x40, 0x9f, 0x96,
                                          0xe9, 0x3d, 0x7e, 0x11, 0x73, 0x93, 0x17, 0x2a};
    static const uint8_t ive[]                                = {0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
                                  0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f};
		/* cipher text based on key = all zeros to avoid programming OTP bits */
    static const uint8_t cipherAes128[] = {0xFC, 0x96, 0x2C, 0xED, 0x40, 0xC6, 0x84, 0x08, 
		                                       0xA2, 0xF5, 0x06, 0xBB, 0x32, 0x10, 0x38, 0x5F };

    uint8_t cipher[16];
    uint8_t output[16];
    status_t status;
		hashcrypt_handle_t m_handle;

		/* secret key PUF, 128-bit key */
		SYSCTL0->AESKEY_SRCSEL = 0x0;
    m_handle.keyType = kHASHCRYPT_SecretKey;
    m_handle.keySize = kHASHCRYPT_Aes192;
																					 
		/* PUF init */
		status = puf_init((uint8_t *) keyAes192, 24);
    TEST_ASSERT(0 == status, "PUF");
	
    status = HASHCRYPT_AES_EncryptCbc(HASHCRYPT, &m_handle, plainAes128, cipher, 16, ive);
    TEST_ASSERT(kStatus_Success == status, "CBC");
    TEST_ASSERT(memcmp(cipher, cipherAes128, 16) == 0, "CBC");

    status = HASHCRYPT_AES_DecryptCbc(HASHCRYPT, &m_handle, cipher, output, 16, ive);
    TEST_ASSERT(kStatus_Success == status, "CBC");
    TEST_ASSERT(memcmp(output, plainAes128, 16) == 0, "CBC");

    PRINTF("AES CBC Test - 192-bit PUF key - pass\r\n");
}

void TestAesPufKeyCtr(void)
{
    static uint8_t aes_ctr_test01_plaintext[]        = {0x6b, 0xc1, 0xbe, 0xe2, 0x2e, 0x40, 0x9f, 0x96,
                                                 0xe9, 0x3d, 0x7e, 0x11, 0x73, 0x93, 0x17, 0x2a};
    static uint8_t aes_ctr_test01_counter_1[16]      = {0xf0, 0xf1, 0xf2, 0xf3, 0xf4, 0xf5, 0xf6, 0xf7,
                                                   0xf8, 0xf9, 0xfa, 0xfb, 0xfc, 0xfd, 0xfe, 0xff};
    static uint8_t aes_ctr_test01_counter_2[16]      = {0xf0, 0xf1, 0xf2, 0xf3, 0xf4, 0xf5, 0xf6, 0xf7,
                                                   0xf8, 0xf9, 0xfa, 0xfb, 0xfc, 0xfd, 0xfe, 0xff};
    static const uint8_t aes_ctr_test01_key[] __attribute__((aligned)) = {
		 0x1F, 0x1E, 0x1D, 0x1C, 0x1B, 0x1A, 0x19, 0x18, 0x17, 0x16, 0x15, 0x14, 0x13, 0x12, 0x11, 0x10,
		 0x0F, 0x0E, 0x0D, 0x0C, 0x0B, 0x0A, 0x09, 0x08, 0x07, 0x06, 0x05, 0x04, 0x03, 0x02, 0x01, 0x00,
		};
	
		/* cipher text based on key = all zeros to avoid programming OTP bits */
    static const uint8_t aes_ctr_test01_ciphertext[] = {0xaf, 0xa8, 0x99, 0x3e, 0x55, 0x66, 0x14, 0xc2, 
		                                                    0x28, 0x3c, 0xb2, 0x54, 0xbf, 0xdd, 0xf7, 0xf5};

    uint8_t cipher[16]                               = {0};
    uint8_t output[16]                               = {0};
    status_t status;
		hashcrypt_handle_t m_handle;

		/* secret key PUF, 256-bit key */
		SYSCTL0->AESKEY_SRCSEL = 0x0;
    m_handle.keyType = kHASHCRYPT_SecretKey;
    m_handle.keySize = kHASHCRYPT_Aes256;
																					 
		/* PUF init */
		status = puf_init((uint8_t *) aes_ctr_test01_key, 32);
    TEST_ASSERT(0 == status, "PUF");
	
    /* Encrypt */
    status = HASHCRYPT_AES_CryptCtr(HASHCRYPT, &m_handle, aes_ctr_test01_plaintext, cipher, 16,
                                    aes_ctr_test01_counter_1, NULL, NULL);
    TEST_ASSERT(kStatus_Success == status, "CTR");
    TEST_ASSERT(memcmp(cipher, aes_ctr_test01_ciphertext, 16) == 0, "CTR");

    /* Decrypt */
    status = HASHCRYPT_AES_CryptCtr(HASHCRYPT, &m_handle, cipher, output, 16, aes_ctr_test01_counter_2, NULL, NULL);
    TEST_ASSERT(kStatus_Success == status, "CTR");
    TEST_ASSERT(memcmp(output, aes_ctr_test01_plaintext, sizeof(aes_ctr_test01_plaintext)) == 0, "CTR");

    PRINTF("AES CTR Test - 256-bit PUF key - pass\r\n");
}

/*!
 * @brief Main function.
 */
int main(void)
{
    /* Init hardware */
    BOARD_InitPins();
    BOARD_BootClockRUN();
    BOARD_InitDebugConsole();

    /* test description */
    PRINTF("\r\nAES ECB, CBC, CTR testing using PUF key\r\n\r\n");

    /* Initialize Hashcrypt */
    HASHCRYPT_Init(HASHCRYPT);

    /* Call HASH APIs */
    TestAesPufKeyEcb();
    TestAesPufKeyCbc();
    TestAesPufKeyCtr();

    HASHCRYPT_Deinit(HASHCRYPT);

    while (1)
    {
    }
}
