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

#include "fsl_device_registers.h"
#include "fsl_debug_console.h"
#include "fsl_shell.h"
#include "pin_mux.h"
#include "clock_config.h"
#include "board.h"

#include "mcuboot_app_support.h"
#include "mflash_drv.h"
#include "xmodem.h"
#include "platform_bindings.h"

#include <ctype.h>

#include "fsl_clock.h"
#include "fsl_reset.h"
#include <stdbool.h>

#include "fsl_romapi.h"
#include "fsl_glikey.h"
#include "fsl_cdog.h"

#include "stdio.h"
/*******************************************************************************
 * Definitions
 ******************************************************************************/

/*******************************************************************************
 * Prototypes
 ******************************************************************************/

static shell_status_t shellCmd_image(shell_handle_t shellHandle, int32_t argc, char **argv);
static shell_status_t shellCmd_xmodem(shell_handle_t shellHandle, int32_t argc, char **argv);
static shell_status_t shellCmd_mem(shell_handle_t shellHandle, int32_t argc, char **argv);
static shell_status_t shellCmd_reboot(shell_handle_t shellHandle, int32_t argc, char **argv);
static shell_status_t shellCmd_swd(shell_handle_t shellHandle, int32_t argc, char **argv);

static int flash_sha256(uint32_t offset, size_t size, uint8_t sha256[32]);


/*******************************************************************************
 * Variables
 ******************************************************************************/
volatile uint32_t s_cdog_start = 0x00U;

static SHELL_COMMAND_DEFINE(image,
                            "\n\"image [info]\"          : Print image information"
                            "\n\"image test [imgNum]\"   : Mark candidate slot of given image number as ready for test"
                            "\n\"image accept [imgNum]\" : Mark active slot of given image number as accepted"
                            "\n\"image erase [imgNum]\"  : Erase candidate slot of given image number"
                            "\n",
                            shellCmd_image,
                            SHELL_IGNORE_PARAMETER_COUNT);

static SHELL_COMMAND_DEFINE(mem,
                            "\n\"mem read addr [size]\" : Read memory at given address"
                            "\n\"mem erase addr \"      : Erase sector containing given address"
                            "\n",
                            shellCmd_mem,
                            SHELL_IGNORE_PARAMETER_COUNT);

static SHELL_COMMAND_DEFINE(xmodem, "\n\"xmodem [imgNum]\": Start receiving with XMODEM-CRC\n", shellCmd_xmodem, SHELL_IGNORE_PARAMETER_COUNT);

static SHELL_COMMAND_DEFINE(reboot, "\n\"reboot\": Triggers software reset\n", shellCmd_reboot, 0);

static SHELL_COMMAND_DEFINE(swd,
                            "\n\"swd unlock [key/password]\" : Unlock the SWD debug port"
                            "\n\"swd lock \"                 : Lock the SWD debug port"
                            "\n",
                            shellCmd_swd,
                            SHELL_IGNORE_PARAMETER_COUNT);

SDK_ALIGN(static uint8_t s_shellHandleBuffer[SHELL_HANDLE_SIZE], 4);
static shell_handle_t s_shellHandle;

/*
 * Buffer used to handover data from XMODEM to flash programming routine.
 * Uses 4B alignment to be compatible with mflash.
 **/
static uint32_t progbuf[1024/sizeof(uint32_t)];

static hashctx_t sha256_xmodem_ctx;

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

static void hexdump(const void *src, size_t size)
{
    const unsigned char *src8 = src;
    const int CNT             = 16;

    for (size_t i = 0; i < size; i++)
    {
        int n = i % CNT;
        if (n == 0)
            PRINTF("%08x  ", (uint32_t)src+i);
        PRINTF("%02X ", src8[i]);
        if ((i && n == CNT - 1) || (i + 1 == size))
        {
            int rem = CNT - 1 - n;
            for (int j = 0; j < rem; j++)
                PRINTF("   ");
            PRINTF("|");
            for (int j = n; j >= 0; j--)
                PUTCHAR(isprint(src8[i - j]) ? src8[i - j] : '.');
            PRINTF("|\n");
        }
    }
    PUTCHAR('\n');
}

static void print_hash(const void *src, size_t size)
{
    const unsigned char *src8 = src;
    for (size_t i = 0; i < size; i++)
    {
        PRINTF("%02X", src8[i]);
    }
}


static shell_status_t shellCmd_image(shell_handle_t shellHandle, int32_t argc, char **argv)
{
    int image = 0;
    int ret;
    status_t status;
    uint32_t imgstate;

    if (argc > 3)
    {
        PRINTF("Too many arguments.\n");
        return kStatus_SHELL_Error;
    }

    /* image [info] */

    if (argc == 1 || (argc == 2 && !strcmp(argv[1], "info")))
    {
        bl_print_image_info(flash_sha256);
        return kStatus_SHELL_Success;
    }

    if (argc < 2)
    {
        PRINTF("Wrong arguments. See 'help'\n");
        return kStatus_SHELL_Error;
    }

    if (argc == 3)
    {
        char *parse_end;
        image = strtol(argv[2], &parse_end, 10);
        
        if (image < 0 || image >= MCUBOOT_IMAGE_NUMBER || *parse_end != '\0')
        {
            PRINTF("Wrong image number.\n");
            return kStatus_SHELL_Error;
        }
    }

    status = bl_get_image_state(image, &imgstate);
    if (status != kStatus_Success)
    {
        PRINTF("Failed to get state of image %u (status %d)", image, status);
        return kStatus_SHELL_Error;
    }

    /* image test [imgNum] */

    if (!strcmp(argv[1], "test"))
    {
        status = bl_update_image_state(image, kSwapType_ReadyForTest);
        if (status != kStatus_Success)
        {
            PRINTF("FAILED to mark image state as ReadyForTest (status=%d)\n", status);
            return kStatus_SHELL_Error;
        }
    }

    /* image accept [imgNum] */

    else if (!strcmp(argv[1], "accept"))
    {
        if (imgstate != kSwapType_Testing)
        {
            PRINTF("Image state is not set as Testing. Nothing to accept.\n", status);
            return kStatus_SHELL_Error;
        }

        status = bl_update_image_state(image, kSwapType_Permanent);
        if (status != kStatus_Success)
        {
            PRINTF("FAILED to accept image (status=%d)\n", status);
            return kStatus_SHELL_Error;
        }
    }
    
    /* image erase [imgNum] */
    
    else if (!strcmp(argv[1], "erase"))
    {
        partition_t ptn;

        ret = bl_get_update_partition_info(image, &ptn);
        if (ret != kStatus_Success)
        {
            PRINTF("Failed to determine update partition\n");
            return kStatus_SHELL_Error;
        }

        uint32_t slotaddr     = ptn.start;
        uint32_t slotsize     = ptn.size;
        uint32_t slotcnt      = (slotsize-1 + MFLASH_SECTOR_SIZE) / MFLASH_SECTOR_SIZE;

        PRINTF("Erasing inactive slot...");
        for (int i=0; i < slotcnt; i++)
        {
            ret = mflash_drv_sector_erase(slotaddr);
            if (ret)
            {
                PRINTF("\nFailed to erase sector at 0x%x (ret=%d)\n", slotaddr, ret);
                return kStatus_SHELL_Error;
            }
            slotaddr += MFLASH_SECTOR_SIZE;
        }
        PRINTF("done\n");
    }
    
    else
    {
        PRINTF("Wrong arguments. See 'help'\n");
        return kStatus_SHELL_Error;
    }

    return kStatus_SHELL_Success;
}

static shell_status_t shellCmd_mem(shell_handle_t shellHandle, int32_t argc, char **argv)
{
    int ret;
    uint32_t addr;
    uint32_t size = 128;
    char *parse_end;

    if (argc < 3 || argc > 4)
    {
        PRINTF("Wrong argument count\n");
        return kStatus_SHELL_Error;
    }
 
    addr = strtol(argv[2], &parse_end, 0);
    if (*parse_end != '\0')
    {
        PRINTF("Bad address\n");
        return kStatus_SHELL_Error;
    }
    
    if (argc == 4)
    {
        size = strtol(argv[3], &parse_end, 0);
        if (*parse_end != '\0')
        {
            PRINTF("Bad size\n");
            return kStatus_SHELL_Error;
        }
    }

    /* mem read addr [size] */

    if (!strcmp(argv[1], "read"))
    {
#ifdef MFLASH_PAGE_INTEGRITY_CHECKS
        if (mflash_drv_is_readable(addr) != kStatus_Success)
        {
            PRINTF("Page not readable\n");
            return kStatus_SHELL_Error;
        }
#endif
        hexdump((void *)addr, size);
    }

    /* mem erase addr */

    else if (!strcmp(argv[1], "erase"))
    {
        ret = mflash_drv_sector_erase(addr & ~(MFLASH_SECTOR_SIZE-1));
        if (ret)
        {
            PRINTF("Failed to erase sector (ret=%d)\n", ret);
            return kStatus_SHELL_Error;
        }
    }

    else
    {
        PRINTF("Wrong arguments. See 'help'\n");
        return kStatus_SHELL_Error;
    }

    return kStatus_SHELL_Success;
}

static int process_received_data(uint32_t dst_addr, uint32_t offset, uint32_t size)
{
    int ret;
    uint32_t *data = progbuf;
    uint32_t addr = dst_addr + offset;
    
    /* 1kB programming buffer should be ok with all page size alignments */
      
    while (size)
    {
        size_t chunk = (size < MFLASH_PAGE_SIZE) ? size : MFLASH_PAGE_SIZE;
               
        /* mlfash takes entire page, in case of last data of smaller size it will
           program more data, which shouln't be a problem as the space allocated
           for the image slot is page aligned */
        
        ret = mflash_drv_page_program(addr, data);
        if (ret)
        {
            PRINTF("Failed to program flash at %x (ret %d)\n", addr, ret);
            return -1;
        }
        
        sha256_update(&sha256_xmodem_ctx, data, chunk);
        addr += chunk;
        data += chunk/sizeof(uint32_t);
        size -= chunk;
    }
    
    return 0;
}

static shell_status_t shellCmd_xmodem(shell_handle_t shellHandle, int32_t argc, char **argv)
{
    int image = 0;
    long recvsize;    
    uint8_t sha256_recv[32], sha256_flash[32];
    partition_t prt_ota;
    
    if (argc > 3)
    {
        PRINTF("Too many arguments.\n");
        return kStatus_SHELL_Error;
    }
    
    if (argc == 3)
    {
        char *parse_end;
        image = strtol(argv[2], &parse_end, 10);
        
        if (image < 0 || image >= MCUBOOT_IMAGE_NUMBER || *parse_end != '\0')
        {
            PRINTF("Wrong image number.\n");
            return kStatus_SHELL_Error;
        }
    }
       
    if (bl_get_update_partition_info(image, &prt_ota) != kStatus_Success)
    {
        PRINTF("FAILED to determine address for download\n");
        return kStatus_SHELL_Error;
    }
    
    PRINTF("Started xmodem download into flash at 0x%X\n", prt_ota.start);
    
    struct xmodem_cfg cfg = {
        .putc = xmodem_putc,
        .getc = xmodem_getc,
        .canread = xmodem_canread,
        .canread_retries = xmodem_canread_retries,
        .dst_addr = prt_ota.start,
        .maxsize = prt_ota.size,
        .buffer = (uint8_t*)progbuf,
        .buffer_size = sizeof(progbuf),
        .buffer_full_callback = process_received_data
    };
    
    sha256_init(&sha256_xmodem_ctx);
    
    PRINTF("Initiated XMODEM-CRC transfer. Receiving... (Press 'x' to cancel)\n");
    
    recvsize = xmodem_receive(&cfg);
    
    /* With some terminals it takes a while before they recover receiving to the console */
    SDK_DelayAtLeastUs(100000, SystemCoreClock);
    
    if (recvsize < 0)
    {
        PRINTF("\nTransfer failed (%d)\n", recvsize);
        return kStatus_SHELL_Error;
    }
       
    PRINTF("\nReceived %u bytes\n", recvsize);    
    
    sha256_finish(&sha256_xmodem_ctx, sha256_recv);
    flash_sha256(prt_ota.start, recvsize, sha256_flash);    
    
    PRINTF("SHA256 of received data: ");
    print_hash(sha256_recv, 10);
    PRINTF("...\n");
    
    PRINTF("SHA256 of flashed data:  ");
    print_hash(sha256_flash, 10);
    PRINTF("...\n");

    return kStatus_SHELL_Success;
}


static shell_status_t shellCmd_reboot(shell_handle_t shellHandle, int32_t argc, char **argv)
{
    PRINTF("System reset!\n");
    NVIC_SystemReset();

    /* return kStatus_SHELL_Success; */
}

static int flash_sha256(uint32_t offset, size_t size, uint8_t sha256[32])
{
    uint32_t buf[128 / sizeof(uint32_t)];
    status_t status;
    hashctx_t sha256ctx;

    sha256_init(&sha256ctx);

    while (size > 0)
    {
        size_t chunk = (size > sizeof(buf)) ? sizeof(buf) : size;
        /* mflash demands size to be in multiples of 4 */
        size_t chunkAlign4 = (chunk + 3) & (~3);

        status = mflash_drv_read(offset, buf, chunkAlign4);
        if (status != kStatus_Success)
        {
            return status;
        }

        sha256_update(&sha256ctx, (unsigned char *)buf, chunk);

        size -= chunk;
        offset += chunk;
    }

    sha256_finish(&sha256ctx, sha256);

    return kStatus_Success;
}

void glikey_write_enable(uint32_t index)
{
    status_t status;

    status = GLIKEY_IsLocked(GLIKEY0);
    if (status != kStatus_GLIKEY_NotLocked)
    {
        PRINTF("\n! GLIKEY locked !\n");
        return;
    }

    GLIKEY_SyncReset(GLIKEY0);
    GLIKEY_StartEnable(GLIKEY0, index);
    GLIKEY_ContinueEnable(GLIKEY0, GLIKEY_CODEWORD_STEP1);
    GLIKEY_ContinueEnable(GLIKEY0, GLIKEY_CODEWORD_STEP2);
    GLIKEY_ContinueEnable(GLIKEY0, GLIKEY_CODEWORD_STEP3);
    GLIKEY_ContinueEnable(GLIKEY0, GLIKEY_CODEWORD_STEP4);
    GLIKEY_ContinueEnable(GLIKEY0, GLIKEY_CODEWORD_STEP5);
    GLIKEY_ContinueEnable(GLIKEY0, GLIKEY_CODEWORD_STEP6);
    GLIKEY_ContinueEnable(GLIKEY0, GLIKEY_CODEWORD_STEP7);
    GLIKEY_ContinueEnable(GLIKEY0, GLIKEY_CODEWORD_STEP_EN);
}

static shell_status_t shellCmd_swd(shell_handle_t shellHandle, int32_t argc, char **argv)
{
    int ret;
    uint32_t addr;
    uint32_t size = 128;
    char *parse_end;
    uint8_t uuid[16];
    uint8_t uuid_from_shell[16];
    uint32_t temp = 0;

    if (argc != 2 && argc != 3)
    {
        PRINTF("Wrong argument count\n");
        PRINTF("%d\r\n", argc);
        return kStatus_SHELL_Error;
    }

    memset(uuid, 0x55, sizeof(uuid));
    memset(uuid_from_shell, 0xAA, sizeof(uuid_from_shell));

    if (argc == 3)
    {
        if (strlen(argv[2]) != 32)
        {
            goto unlock_failed;
        }

        for (uint8_t i = 0; i < 32; i += 2)
        {
            sscanf(&argv[2][i], "%2x", &temp);
            uuid_from_shell[i / 2] = temp;
        }

        ROMAPI_GetUUID(uuid);
    }

    /* swd unlock [key/password] */

    if (!strcmp(argv[1], "unlock"))
    {
    	CDOG_Start(CDOG, 0xFFFFFFU, s_cdog_start);

    	for(uint8_t check_pos = 0; check_pos < 16; check_pos++)
    	{
            s_cdog_start += uuid_from_shell[check_pos];
            CDOG_Add(CDOG, uuid[check_pos]);
    	}

        /* Check the password*/
        if (memcmp(uuid_from_shell, uuid, 16) == 0)
        {
        	CDOG_Stop(CDOG, s_cdog_start);

            PRINTF("DEBUG_LOCK_EN = 0x%X\r\n", SYSCON->DEBUG_LOCK_EN);
            /* Enable debug*/
            glikey_write_enable(1);
            SYSCON->DEBUG_FEATURES = SYSCON_DEBUG_FEATURES_CPU0_DBGEN(2) |
                                     SYSCON_DEBUG_FEATURES_CPU0_NIDEN(2);
            SYSCON->DEBUG_FEATURES_DP = SYSCON_DEBUG_FEATURES_DP_CPU0_DBGEN(2) |
                                        SYSCON_DEBUG_FEATURES_DP_CPU0_NIDEN(2);

            SYSCON->SWD_ACCESS_CPU0 = 0x12345678;

            GLIKEY_EndOperation(GLIKEY0);

            PRINTF("swd unlock done!\n");
        }
        /* key/password is wrong */
        else
        {
        	goto unlock_failed;
        }
    }

    /* swd lock */

    else if (!strcmp(argv[1], "lock"))
    {
        /* lock the swd */
        glikey_write_enable(1);
        SYSCON->DEBUG_FEATURES = SYSCON_DEBUG_FEATURES_CPU0_DBGEN(1) |
                                 SYSCON_DEBUG_FEATURES_CPU0_NIDEN(1);
        SYSCON->DEBUG_FEATURES_DP = SYSCON_DEBUG_FEATURES_DP_CPU0_DBGEN(1) |
                                    SYSCON_DEBUG_FEATURES_DP_CPU0_NIDEN(1);
        SYSCON->SWD_ACCESS_CPU0 = 0x0;

        GLIKEY_EndOperation(GLIKEY0);

        PRINTF("swd lock done!\n");
    }

    else
    {
        PRINTF("Wrong arguments. See 'help'\n");
        return kStatus_SHELL_Error;
    }

    return kStatus_SHELL_Success;

unlock_failed:
	PRINTF("key/password is wrong!\n");
	/* Mitigation of violent attacks */
	SDK_DelayAtLeastUs(10 * 1000 * 1000, SystemCoreClock);
	return kStatus_SHELL_Success;
}

static void cdog_initial(void)
{
    cdog_config_t conf;
    status_t result = kStatus_Fail;

    CDOG_GetDefaultConfig(&conf);

    conf.lock       = (uint8_t)kCDOG_LockCtrl_Lock;    /* Lock control */
    conf.timeout    = (uint8_t)kCDOG_FaultCtrl_EnableReset; /* Timeout control */
    conf.miscompare = (uint8_t)kCDOG_FaultCtrl_EnableReset; /* Miscompare control */
    conf.sequence   = (uint8_t)kCDOG_FaultCtrl_EnableReset; /* Sequence control */
    conf.state      = (uint8_t)kCDOG_FaultCtrl_EnableReset; /* State control */
    conf.address    = (uint8_t)kCDOG_FaultCtrl_EnableReset; /* Address control */

    conf.irq_pause  = (uint8_t)kCDOG_IrqPauseCtrl_Pause;   /* IRQ pause control */
    conf.debug_halt = (uint8_t)kCDOG_DebugHaltCtrl_Pause;  /* Debug halt control */

    /* Clears pending FLAGS and sets CONTROL register */
    result = CDOG_Init(CDOG, &conf);
    if (result != kStatus_Success)
    {
        PRINTF("Error while CDOG Init. CDOG was probably not in IDLE mode due SW reset.\r\n");
        while(1);
    }
}

/*!
 * @brief Main function
 */
int main(void)
{
    int ret;
    s_shellHandle = &s_shellHandleBuffer[0];

    /* Init board hardware. */
    BOARD_InitPins();
    BOARD_InitBootClocks();
    BOARD_InitDebugConsole();
    
    ret = mflash_drv_init();
    if (ret)
    {
        PRINTF("Failed to init flash driver\n");
    }

    PRINTF("\n"
           "*************************************\n"
           "* Basic MCUBoot application example *\n"
           "*************************************\n\n");
    
    PRINTF("Built " __DATE__ " " __TIME__ "\n");
    
    PRINTF("Please use ISP blhost with 'get-property 18' to get uuid\r\n");
    PRINTF("--- !Please DO NOT use UUID as password in REAL products! ---");

    ret = SHELL_Init(s_shellHandle, g_serialHandle, "$ ");
    if (ret != kStatus_SHELL_Success)
    {
        PRINTF("Failed to init shell\n");
        goto failed_init;
    }

    SHELL_RegisterCommand(s_shellHandle, SHELL_COMMAND(image));
    SHELL_RegisterCommand(s_shellHandle, SHELL_COMMAND(xmodem));
    SHELL_RegisterCommand(s_shellHandle, SHELL_COMMAND(mem));
    SHELL_RegisterCommand(s_shellHandle, SHELL_COMMAND(reboot));
    SHELL_RegisterCommand(s_shellHandle, SHELL_COMMAND(swd));


    while (1)
    {
        SHELL_Task(s_shellHandle);
    }

failed_init:
    while (1)
    {        
    }
}
