#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <stdint.h>
#include <sys/stat.h>
#include <arpa/inet.h>
#include "../ram/inc/sha256.h"
#include "iot_img_hdr.h"

char default_hash_iv[]={0x12,0x34,0x56,0x78,0x87,0x65,0x43,0x21,0x5A,0x5A,0x5A,0x5A,0x00,0x7E,0xC3,0x00};

#define NEW_LINE "\r\n"

#define OK        0
#define ERROR    (-1)

#define FUNC_CONSTRUCT      "-c"
#define FUNC_ADDHEADER      "-a"
#define FUNC_SHOWHEADER     "-s"

#define KW_SETMANUFACTURE   "-m"

#define PK_NAME_KEY_WORD    "iot_flash"
#define PK_SINGLE_FILE_SIZE 0x1000000 /* sigle file size keeped under 16MB */
#define MAX_IMAGE_NUM       20
#define imgHeaderSize       HEADER_TOLTAL_SIZE

char * global_buffer = NULL;

unsigned int CRC32_Update(unsigned int crc, char *buff, int len);

typedef unsigned short UINT16;
typedef unsigned int UINT32;

typedef struct
{
    char *pName;
    int ver;
    int devType;
    int imgType;
}ftype;

enum{
    RET_OK,
    RET_ENV_FAILED,
    RET_BAD_PARA,
    RET_IMAGE_WRONG,
};

#define GET_BYTE1(i) ((i)&0xFF)
#define GET_BYTE2(i) (((i)>>8)&0xFF)
#define GET_BYTE3(i) (((i)>>16)&0xFF)
#define GET_BYTE4(i) (((i)>>24)&0xFF)

#define PATH_LENGTH 128
#define KEYWRD_LENGTH (PATH_LENGTH - 12) /* "20170901.bin" -> 12 bytes */
#define MANU_NAME_LENGTH 32
typedef struct ah_image_config_
{
    int image_type;
    int version;
    int offset;
    int run_addr;
    int length;
    char path[PATH_LENGTH];
    char fwpath[PATH_LENGTH];
}ah_img_cfg;

typedef struct ah_device_config_
{
    int device_type;
    int encrypt_mode;
    int encrypt_pattern;
    int image_num;
    int version;
    int top_header_enable;
    int package_type;
    int flash_size;
    int psram_size;
    int layout_index;
    int lzma_type;
    int fastboot;
    int fw2_load;
    int img_type;
    int vendor_id;
    int magic_no;
    int crc;
    char key_word[PATH_LENGTH];
    ah_img_cfg image[MAX_IMAGE_NUM];
}ah_cfg;

typedef struct ah_encrypt_t_
{
    uint8_t encrypt_mode;
    uint8_t encrypt_pattern;
    uint32_t vender_id;
    char manu[MANU_NAME_LENGTH];
} ah_encrypt_t;

#define CH_SPC    0x20
#define CH_TAB    0x09

/* encrypt info array. */
ah_encrypt_t encrypt_array[] = {
    {0, 0x55, 18516, "htzd"},
    {3, 0x66, 20810, "qj"},
    {3, 0x99, 20552, "flx"},
    {3, 0xcc, 20549, "spe"},
    {3, 0x33, 18264, "gx"},
    {3, 0xdd, 17492, "dt"},
    {3, 0xa9, 22864, "yp"},
};

int ah_funShowHeader(char argc, char *argv[]);

int ah_atoi(char * str)
{
    int value = 0;
    int sign = 1;
    int radix;


    if(*str == '-')
    {
        sign = -1;
        str++;
    }
    if(*str == '0' && (*(str+1) == 'x' || *(str+1) == 'X'))
    {
        radix = 16;
        str += 2;
    }
    else if(*str == '0')
    {
        radix = 8;
        str++;
    }
    else
    {
        radix = 10;
    }
    while(*str)
    {
        if(radix == 16)
        {
            if(*str >= '0' && *str <= '9')
            {
                value = value * radix + *str - '0';
            }
            else if(*str >= 'A' && *str <= 'F')
            {
                value = value * radix + *str - 'A' + 10;
            }
            else if(*str >= 'a' && *str <= 'f')
            {
                value = value * radix + *str - 'a' + 10;
            }
            else
            {
                break;
            }
        }
        else if(radix == 8)
        {
            if(*str >= '0' && *str <= '7')
            {
                value = value * radix + *str - '0';
            }
            else
            {
                break;
            }
        }
        else
        {
            if(*str >= '0' && *str <= '9')
            {
                value = value * radix + *str - '0';
            }
            else
            {
                break;
            }
        }
        str ++;
    }
    return sign*value;
}

static int ah_isSpace(int c)
{
  switch (c) {
  case 0x20:
  case 0xD:
  case 0xA:
  case 0x9:
    return 1;
  }
  return 0;
}
#define ISSPACE(c) ah_isSpace(c)

/* skip white space */
char *ah_GetValidString(char *str)
{
    char *start, *end;
    start = str;
    end = str+strlen(str)-1;

    while((start<=end) && ISSPACE(*start)) start++;

    while((start<=end) && ISSPACE(*end)) end--;

    if(start > end)
        return NULL;

    *(end+1) = 0x0;

    return start;
}

char *ah_GetKeyString(FILE *fp, char *title, char *key)
{
#define LINE_LENGTH 1024
    int  start=0;
    char *pt;
    static char sLine[LINE_LENGTH];

    fseek(fp, 0, SEEK_SET);

    while (NULL != fgets(sLine, LINE_LENGTH, fp)) {

        if (('#' ==sLine[0])
            ||(';' ==sLine[0]))
            continue;

        if('[' ==sLine[0])
        {
            start =0;
        }

        if(start)
        {
            if((NULL != (pt = strchr(sLine, '=')))
                &&(0 == strncmp(key, sLine, pt-sLine)))
            {
                sLine[LINE_LENGTH-1] = 0x0;
                pt = ah_GetValidString(pt + 1);;
                return pt;
            }

        }
        else
        {
            if(0 == strncmp(title, sLine, strlen(title)-1))
            {
                start =1;
            }
        }

    }
    return NULL;
}

void ah_help(void)
{
    printf(NEW_LINE"--Usage of ah:"NEW_LINE);
    printf(NEW_LINE"ah -c config-file");
    printf(NEW_LINE"    e.g. ah -c ./ah_config.cfg");
    printf(NEW_LINE"ah -s file");
    printf(NEW_LINE"    Show the header(s) of a package."NEW_LINE);
    printf(NEW_LINE"ah -a file");
    printf(NEW_LINE"    Add header to fw img."NEW_LINE);
    printf(NEW_LINE);

    return ;
}

static void ah_Sha256(unsigned char* img, int len, unsigned char* sha256)
{
    sha256_context ctx;

    sha256_init( &ctx );
    sha256_starts( &ctx );
    sha256_update( &ctx, img, len );
    sha256_update( &ctx, (unsigned char*)default_hash_iv, 16);
    sha256_finish( &ctx, sha256 );
    sha256_free( &ctx );
}

#define DEVICE_STR              "[DEVICE]"
#define DEVICE_STR_TYPE         "device_type"
#define DEVICE_STR_IMGNUM       "image_num"
#define DEVICE_STR_KEYWRD       "key_word"
#define DEVICE_STR_TOPHDRENA    "top_header_enable"
#define DEVICE_STR_TOPHDRENC    "encrypt_mode"
#define DEVICE_STR_TOPHDRPTN    "encrypt_pattern"
#define DEVICE_STR_TOPHDRVER    "version"
#define DEVICE_STR_PACKAGE_TYPE "package_type"
#define DEVICE_STR_MAGIC_NO     "magic_no"
#define DEVICE_STR_FLASH_SIZE   "flash_size"
#define DEVICE_STR_PSRAM_SIZE   "psram_size"
#define DEVICE_STR_LAYOUT_INDEX "layout_index"
#define DEVICE_STR_LZMA_TYPE    "lzma_type"
#define DEVICE_STR_FASTBOOT     "fastboot"
#define DEVICE_STR_FW2_LOAD     "fw2_load"
#define DEVICE_STR_IMAGE_TYPE   "img_type"
#define DEVICE_STR_VENDOR_ID    "vendor_id"

#define IMAGE_STR               "[IMAGE_%02d]"
//#define IMAGE_STR_PADENA       "pad_enable"
//#define IMAGE_STR_PAD          "pad"
#define IMAGE_STR_TYPE          "image_type"
#define IMAGE_STR_VERSION       "version"
#define IMAGE_STR_PATH          "path"
#define IMAGE_STR_OFFSET        "offset"
#define IMAGE_STR_RUN_ADDR      "run_addr"
#define IMAGE_STR_LENGTH        "length"

int ah_GetConfig(char *pcfgName, ah_cfg *pcfg)
{
    FILE *fconfig;
    char *pkey;
    char imageName[64];
    ah_cfg *cfg=pcfg;
    int tmp0,tmp1,tmp2,tmp3,i;

    if(NULL == (fconfig=fopen(pcfgName, "r")))
    {
        printf("\r\n Can not open config: %s!\r\n", pcfgName);
        return -1;
    }

    if(NULL != (pkey = ah_GetKeyString(fconfig, DEVICE_STR, DEVICE_STR_TYPE)))
    {
        cfg->device_type = ah_atoi(pkey);
    }
    else
    {
        printf("\r\n Can not find %s in %s !\r\n", DEVICE_STR_TYPE, DEVICE_STR);
        return -1;
    }

    if(NULL != (pkey = ah_GetKeyString(fconfig, DEVICE_STR, DEVICE_STR_IMGNUM)))
    {
        cfg->image_num = ah_atoi(pkey);
    }
    else
    {
        printf("\r\n Can not find %s in %s !\r\n", DEVICE_STR_IMGNUM, DEVICE_STR);
        return -1;
    }

    if(NULL != (pkey = ah_GetKeyString(fconfig, DEVICE_STR, DEVICE_STR_KEYWRD)))
    {
        if(strlen(pkey) <= KEYWRD_LENGTH)
        {
            strcpy(cfg->key_word, pkey);
        }
        else
        {
            printf("\r\n Keywords over length: %s !\r\n", pkey);
            return -1;
        }
    }
    else
    {
        printf("\r\n Can not find %s in %s !\r\n", DEVICE_STR_KEYWRD, DEVICE_STR);
        return -1;
    }


    if(NULL != (pkey = ah_GetKeyString(fconfig, DEVICE_STR, DEVICE_STR_TOPHDRENA)))
    {
        cfg->top_header_enable = ah_atoi(pkey);
    }
    else
    {
        printf("\r\n Can not find %s in %s !\r\n", DEVICE_STR_TOPHDRENA, DEVICE_STR);
        return -1;
    }

    if(cfg->top_header_enable)
    {
        if(NULL != (pkey = ah_GetKeyString(fconfig, DEVICE_STR, DEVICE_STR_TOPHDRENC)))
        {
            cfg->encrypt_mode = ah_atoi(pkey);
        }
        else
        {
            printf("\r\n Can not find %s in %s !\r\n", DEVICE_STR_TOPHDRENC, DEVICE_STR);
            return -1;
        }

        if(NULL != (pkey = ah_GetKeyString(fconfig, DEVICE_STR, DEVICE_STR_TOPHDRPTN)))
        {
            cfg->encrypt_pattern = ah_atoi(pkey);
        }
        else
        {
            printf("\r\n Can not find %s in %s !\r\n", DEVICE_STR_TOPHDRPTN, DEVICE_STR);
            return -1;
        }

        if(NULL != (pkey = ah_GetKeyString(fconfig, DEVICE_STR, DEVICE_STR_TOPHDRVER)))
        {
            sscanf(pkey, "%d.%d.%d.%d", &tmp0, &tmp1, &tmp2, &tmp3);
            cfg->version = ((tmp0&0x1F)<<27)|((tmp1&0x7F)<<20)|((tmp2&0xF)<<16)|(tmp3&0xFFFF);
        }
        else
        {
            printf("\r\n Can not find %s in %s !\r\n", DEVICE_STR_TOPHDRVER, DEVICE_STR);
            return -1;
        }

        if(NULL != (pkey = ah_GetKeyString(fconfig, DEVICE_STR, DEVICE_STR_FLASH_SIZE)))
        {
            cfg->flash_size = ah_atoi(pkey);
        }
        if(NULL != (pkey = ah_GetKeyString(fconfig, DEVICE_STR, DEVICE_STR_LAYOUT_INDEX)))
        {
            cfg->layout_index = ah_atoi(pkey);
        }
        if(NULL != (pkey = ah_GetKeyString(fconfig, DEVICE_STR, DEVICE_STR_PSRAM_SIZE)))
        {
            cfg->psram_size = ah_atoi(pkey);
        }
        if(NULL != (pkey = ah_GetKeyString(fconfig, DEVICE_STR, DEVICE_STR_PACKAGE_TYPE)))
        {
            cfg->package_type = ah_atoi(pkey);
        }
        else
        {
            printf("\r\n Can not find %s in %s !\r\n", DEVICE_STR_PACKAGE_TYPE, DEVICE_STR);
            return -1;
        }

        if(NULL != (pkey = ah_GetKeyString(fconfig, DEVICE_STR, DEVICE_STR_MAGIC_NO)))
        {
            cfg->magic_no = ah_atoi(pkey);
        }
        else
        {
            printf("\r\n Can not find %s in %s !\r\n", DEVICE_STR_MAGIC_NO, DEVICE_STR);
            return -1;
        }
        if(NULL != (pkey = ah_GetKeyString(fconfig, DEVICE_STR, DEVICE_STR_LZMA_TYPE)))
        {
            cfg->lzma_type = ah_atoi(pkey);
        }
        if(NULL != (pkey = ah_GetKeyString(fconfig, DEVICE_STR, DEVICE_STR_FASTBOOT)))
        {
            cfg->fastboot = ah_atoi(pkey);
        }
        if (NULL != (pkey = ah_GetKeyString(fconfig, DEVICE_STR, DEVICE_STR_FW2_LOAD)))
        {
            cfg->fw2_load = ah_atoi(pkey);
        }
        if (NULL != (pkey = ah_GetKeyString(fconfig, DEVICE_STR, DEVICE_STR_IMAGE_TYPE)))
        {
            cfg->img_type = ah_atoi(pkey);
        }

        if(NULL != (pkey = ah_GetKeyString(fconfig, DEVICE_STR, DEVICE_STR_VENDOR_ID)))
        {
            cfg->vendor_id = ah_atoi(pkey);
        }
    }

    for(i = 0; i < cfg->image_num; i++)
    {
        sprintf(imageName, IMAGE_STR, i+1);
        if(NULL != (pkey = ah_GetKeyString(fconfig, imageName, IMAGE_STR_TYPE)))
        {
            cfg->image[i].image_type = ah_atoi(pkey);
        }
        else
        {
            printf("\r\n Can not find %s in %s !\r\n", IMAGE_STR_TYPE, imageName);
            return -1;
        }

        if(NULL != (pkey = ah_GetKeyString(fconfig, imageName, IMAGE_STR_VERSION)))
        {
            sscanf(pkey, "%d.%d.%d.%d", &tmp0, &tmp1, &tmp2, &tmp3);
            cfg->image[i].version = ((tmp0&0xF)<<27)|((tmp1&0xF)<<20)|
                ((tmp2&0xF)<<16)|(tmp3&0xFF);
        }
        else
        {
            printf("\r\n Can not find %s in %s !\r\n", IMAGE_STR_VERSION, imageName);
            return -1;
        }

        if(NULL != (pkey = ah_GetKeyString(fconfig, imageName, IMAGE_STR_PATH)))
        {
            if(strlen(pkey) <= PATH_LENGTH)
            {
                strcpy(cfg->image[i].path, pkey);
                strcpy(cfg->image[i].fwpath, pkey);
                strcat(cfg->image[i].fwpath, ".fw");
            }
            else
            {
                printf("\r\n Path over length: %s !\r\n", pkey);
                return -1;
            }
        }
        else
        {
            printf("\r\n Can not find %s in %s !\r\n", IMAGE_STR_PATH, imageName);
            return -1;
        }

        if(NULL != (pkey = ah_GetKeyString(fconfig, imageName, IMAGE_STR_OFFSET)))
        {
            cfg->image[i].offset = ah_atoi(pkey);
        }
        else
        {
            printf("\r\n Can not find %s in %s !\r\n", IMAGE_STR_OFFSET, imageName);
            return -1;
        }

        if ((cfg->image[i].image_type == imgV1SBL)
            || (cfg->image[i].image_type == imgV1FRWCUS)
            || (cfg->image[i].image_type == imgV1FRWPLC)) {
            if(NULL != (pkey = ah_GetKeyString(fconfig, imageName, IMAGE_STR_RUN_ADDR)))
            {
                cfg->image[i].run_addr = ah_atoi(pkey);
            }
        } else {
            cfg->image[i].run_addr = 0;
        }

        if(NULL != (pkey = ah_GetKeyString(fconfig, imageName, IMAGE_STR_LENGTH)))
        {
            cfg->image[i].length = ah_atoi(pkey);
        }
        else
        {
            printf("\r\n Can not find %s in %s !\r\n", IMAGE_STR_LENGTH, imageName);
            return -1;
        }

#if 0
        if(NULL != (pkey = ah_GetKeyString(fconfig, imageName, IMAGE_STR_PADENA)))
        {
            cfg->image[i].pad_enable = ah_atoi(pkey);
        }
        else
        {
            printf("\r\n Can not find %s in %s !\r\n", IMAGE_STR_PADENA, DEVICE_STR);
            return -1;
        }

        if(cfg->image[i].pad_enable)
        {
            if(NULL != (pkey = ah_GetKeyString(fconfig, imageName, IMAGE_STR_PAD)))
            {
                cfg->image[i].pad = ah_atoi(pkey);
            }
            else
            {
                printf("\r\n Can not find %s in %s !\r\n", IMAGE_STR_PAD, DEVICE_STR);
                return -1;
            }
        }
#endif
    }

    return 0;
}

#define PACKAGE_HEADER_SIZE         96



// 填充额外的oem数据
int ah_add_oem_ext(char *buf)
{
    int file_size;
    char *buf_start=buf;
    FILE *fp = NULL;
    struct stat st;
    const char *file_name="oem_ext.bin";
    const char *oem_magic_str="_oem_ext";
    if((0 != stat(file_name, &st))
        ||(NULL == (fp=fopen(file_name, "r"))))
    {
        printf(NEW_LINE"Cannot open image file(%s)."NEW_LINE, file_name);
        return 0;
    }
    memcpy(buf,oem_magic_str,strlen(oem_magic_str));
    buf += strlen(oem_magic_str);
    fseek(fp, 0, SEEK_END);
    file_size = ftell(fp);
    fseek(fp, 0, SEEK_SET);
    fread(buf, 1, file_size, fp);
    buf += file_size;
    fclose(fp);
    fp = NULL;
    return buf-buf_start;
}




// 这里添加的是img的头
void ah_fillHeader(char* p_buf, ah_cfg *cfg, int cnt, int fileSize, int _CRC,
    unsigned char *sha)
{
    int i;

    /* Fill header structure start */

    /* DEVICT TYPE */
    p_buf[0] = (char)GET_BYTE1(cfg->device_type);

    /* IMG TYPE */
    p_buf[1] = (char)GET_BYTE1(cfg->image[cnt].image_type);

    /* ENCRYPT VALUE */
    if (cfg->encrypt_mode) {
        p_buf[2] = (char)GET_BYTE1(cfg->encrypt_pattern);
    } else {
        p_buf[2] = 0;
    }

    /* CFG */
    // 快速启动也是在这里配置的
    if (cfg->fastboot == 1) {
        p_buf[3] = 1;
    } else {
        p_buf[3] = 0;
    }

    /* IMG SIZE */
    p_buf[4] = (char)GET_BYTE4(fileSize);
    p_buf[5] = (char)GET_BYTE3(fileSize);
    p_buf[6] = (char)GET_BYTE2(fileSize);
    p_buf[7] = (char)GET_BYTE1(fileSize);

    /* ZIP TYPE */
    if ((cfg->image[cnt].image_type == imgV1FRWCUS) &&
        (cfg->fastboot == 1)) {
        p_buf[8] = 0;
    } else {
        p_buf[8] = (char)GET_BYTE1(cfg->lzma_type);
    }

    /* FLASH SIZE */
    p_buf[9] = (char)GET_BYTE1(cfg->flash_size);

    /* PSRAM SIZE */
    p_buf[10] = (char)GET_BYTE1(cfg->psram_size);

    /* HDR VER */
    p_buf[11] = hdrVer_10;

    /* IMG VER */
    // 用大端方式存储多字节
    p_buf[12] = (char)GET_BYTE4(cfg->image[cnt].version);
    p_buf[13] = (char)GET_BYTE3(cfg->image[cnt].version);
    p_buf[14] = (char)GET_BYTE2(cfg->image[cnt].version);
    p_buf[15] = (char)GET_BYTE1(cfg->image[cnt].version);

    /* IMG CRC */
    p_buf[16] = (char)GET_BYTE4(_CRC);
    p_buf[17] = (char)GET_BYTE3(_CRC);
    p_buf[18] = (char)GET_BYTE2(_CRC);
    p_buf[19] = (char)GET_BYTE1(_CRC);

    /* Guard */
    p_buf[20] = (char)GET_BYTE4(IMAGE_GUARD);
    p_buf[21] = (char)GET_BYTE3(IMAGE_GUARD);
    p_buf[22] = (char)GET_BYTE2(IMAGE_GUARD);
    p_buf[23] = (char)GET_BYTE1(IMAGE_GUARD);

    /* flash layout index */
    p_buf[24] = cfg->layout_index;

    /* reserved, default 0 */
    p_buf[25] = 0;
    p_buf[26] = 0;
    p_buf[27] = 0;

    if ((cfg->image[cnt].image_type == imgV1SBL)
        ||(cfg->image[cnt].image_type == imgV1FRWPLC)
        ||(cfg->image[cnt].image_type == imgV1FRWCUS)) {
        p_buf[28] = (char)GET_BYTE4(cfg->image[cnt].run_addr);
        p_buf[29] = (char)GET_BYTE3(cfg->image[cnt].run_addr);
        p_buf[30] = (char)GET_BYTE2(cfg->image[cnt].run_addr);
        p_buf[31] = (char)GET_BYTE1(cfg->image[cnt].run_addr);
    }

    /* SHA256 */
    for(i=32; i<64; i++)
    {
        p_buf[i] = sha[i-32];
    }

    return;
}

int ah_addHeader(char argc, char *argv[])
{
    FILE *fps = NULL, *fpd = NULL;
    int cnt, _CRC, fileSize, crcSize, ret =RET_OK;
    ah_cfg cfg;
    unsigned char sha[32];
    struct stat st;
    char *sh_argv[3];
    char *p_buf = NULL;
    int len = 0;

    if(argc < 2)
    {
        return RET_BAD_PARA;
    }

    memset(&cfg, 0x0, sizeof(ah_cfg));

    /* Get information from input */
    if(NULL == (global_buffer = (char*)malloc(PK_SINGLE_FILE_SIZE)))
    {
        ret = RET_ENV_FAILED;
        goto ERROR_HANDLE;
    }

    if(0 != ah_GetConfig(argv[2], &cfg))
    {
        ret = RET_ENV_FAILED;
        goto ERROR_HANDLE;
    }

    if(cfg.image_num>MAX_IMAGE_NUM || cfg.image_num<0)
    {
        printf(NEW_LINE"Invalid config file(%s)."NEW_LINE, argv[2]);
        ret = RET_ENV_FAILED;
        goto ERROR_HANDLE;
    }

    for(cnt = 0; cnt < cfg.image_num; cnt++)
    {
        // 不为fw 或 是fw且为快速启动时 不添加img_header
        if(((cfg.image[cnt].image_type != imgV1FRWPLC) &&
            (cfg.image[cnt].image_type != imgV1FRWCUS)) ||
            ((cfg.image[cnt].image_type == imgV1FRWCUS) &&
            (cfg.fastboot == 1))) {
            continue;
        }

        memset(global_buffer, 0x0, PK_SINGLE_FILE_SIZE);

        p_buf = global_buffer;

        if((0 != stat(cfg.image[cnt].path, &st))
            ||(NULL == (fps=fopen(cfg.image[cnt].path, "r"))))
        {
            printf(NEW_LINE"Cannot open image file(%s)."NEW_LINE, cfg.image[cnt].path);
            ret = RET_ENV_FAILED;
            goto ERROR_HANDLE;
        }

        /* Get file size */
        fseek(fps, 0, SEEK_END);
        fileSize = ftell(fps);
        fseek(fps, 0, SEEK_SET);

        /* Read file into buffer */
        fread(p_buf + imgHeaderSize, 1, fileSize, fps);
        fclose(fps);
        fps = NULL;

        crcSize = IMAGE_LEN_FOR_CRC(fileSize);

        /* Clear pad bytes to '0' */
        memset(p_buf + imgHeaderSize + fileSize, 0x0, crcSize - fileSize);

        if((crcSize+imgHeaderSize) > PK_SINGLE_FILE_SIZE)
        {
            printf(NEW_LINE"Image oversize(%s, size:%d, limited:%d)."NEW_LINE,
                cfg.image[cnt].path, crcSize, cfg.image[cnt].length - imgHeaderSize);
            ret = RET_IMAGE_WRONG;
            goto ERROR_HANDLE;
        }

        /* for CRC */
        _CRC= 0xFFFFFFFF^CRC32_Update(0xFFFFFFFF, p_buf + imgHeaderSize, crcSize);

        ah_Sha256((unsigned char*)(p_buf + imgHeaderSize), fileSize, sha);

        /* fill header. */
        ah_fillHeader(p_buf, &cfg, cnt, fileSize, _CRC, sha);

        /* write size */
        len = crcSize + imgHeaderSize;

        /*write buffer into file*/
        if(NULL == (fpd=fopen(cfg.image[cnt].path, "w")))
        {
            printf(NEW_LINE"Cannot open package file(%s)."NEW_LINE,
                cfg.image[cnt].path);
            goto ERROR_HANDLE;
        }

        fseek(fpd, 0, SEEK_SET);
        fwrite(global_buffer, 1, len, fpd);
        fclose(fpd);
        fpd = NULL;

        /* show info */
        sh_argv[2] = cfg.image[cnt].path;
        ah_funShowHeader(3, sh_argv);
        printf(NEW_LINE"Add header to %s success"NEW_LINE, cfg.image[cnt].path);
    }

    ret = RET_OK;

ERROR_HANDLE:

    if(global_buffer)
    {
        free(global_buffer);
        global_buffer = NULL;
    }
    if(fpd)
    {
        fclose(fpd);
        fpd = NULL;
    }
    if(fps)
    {
        fclose(fps);
        fps = NULL;

    }
    return ret;
}

int ah_funConstruct(char argc, char *argv[])
{
    FILE *fps = NULL, *fpd = NULL;
    int cnt, _CRC, fileSize, crcSize, ret =RET_OK, i;
    char pkName[256];
    time_t rawtime;
    struct tm *timeinfo;
    ah_cfg cfg;
    unsigned char sha[32];
    struct stat st;
    char *sh_argv[3];
    char *p_buf = NULL;
    char *p_saved_buf = NULL;
    iot_pkg_hdr_t *hdr = NULL;
    char *payload = NULL;
    int pad_len = 0;
    int len = 0;
    int offset = 0;
    int total_size = 0;
    ah_encrypt_t *p_encrypt = NULL;

    if(argc < 2)
    {
        return RET_BAD_PARA;
    }

    if ((argc > 4) && (0 == strcmp(KW_SETMANUFACTURE, argv[3]))) {
        for (int i = 0; i < sizeof(encrypt_array)/ sizeof(encrypt_array[0]); i++) {
            if (0 == strcmp(encrypt_array[i].manu, argv[4])) {
                p_encrypt = &encrypt_array[i];
                break;
            }
        }
    }

    memset(&cfg, 0x0, sizeof(ah_cfg));

    /* Get local time */
    time ( &rawtime );
    timeinfo = localtime ( &rawtime );

    /* Get information from input */
    if(NULL == (global_buffer = (char*)malloc(PK_SINGLE_FILE_SIZE)))
    {
        ret = RET_ENV_FAILED;
        goto ERROR_HANDLE;
    }

    if(0 != ah_GetConfig(argv[2], &cfg))
    {
        ret = RET_ENV_FAILED;
        goto ERROR_HANDLE;
    }

    if (p_encrypt) {
        cfg.encrypt_mode = p_encrypt->encrypt_mode;
        cfg.encrypt_pattern = p_encrypt->encrypt_pattern;
        cfg.vendor_id = p_encrypt->vender_id;
    }

    if(cfg.image_num>MAX_IMAGE_NUM || cfg.image_num<0)
    {
        printf(NEW_LINE"Invalid config file(%s)."NEW_LINE, argv[2]);
        ret = RET_ENV_FAILED;
        goto ERROR_HANDLE;
    }

    /* make package name */
    // key_word 被gen_img.sh写入在tmp.cfg文件中 这里添加生成日期
    sprintf(pkName, "./%s%04d%02d%02d.bin", cfg.key_word,\
            timeinfo->tm_year+1900, timeinfo->tm_mon+1, timeinfo->tm_mday);

    memset(global_buffer, 0x0, PK_SINGLE_FILE_SIZE);

    if(cfg.top_header_enable)
    {
        /*image header*/
        hdr = (iot_pkg_hdr_t*)global_buffer;
        hdr->enctype = cfg.encrypt_mode;
        hdr->pattern = cfg.encrypt_pattern;
        hdr->dev_type = htons(cfg.package_type);
        hdr->magic_no = htonl(cfg.magic_no);
        hdr->version = htonl(cfg.version);
        hdr->vendor_id = htons(cfg.vendor_id);
        hdr->img_flag = (cfg.flash_size) << 4 | (cfg.lzma_type)<< 2 | (cfg.fw2_load)<< 3;
        hdr->img_type = cfg.img_type;

        /*image payload*/
        payload = global_buffer + sizeof(iot_pkg_hdr_t);
    }
    else
    {
        /*image payload*/
        payload = global_buffer;
    }

    p_buf = payload;

    /* 2400 aligned */
    if((cfg.image_num > 0) &&
        ((cfg.image[0].image_type == imgV1CFG)
        ||(cfg.image[0].image_type == imgV1FRWPLC)
        ||(cfg.image[0].image_type == imgV1FRWCUS)))
    {
            len = p_buf - global_buffer;
            if(len%2400)
            {
                len = 2400 - (len%2400);
                for(i=0; i<len; i++)
                {
                    *p_buf = 0;
                    p_buf++;
                }
                total_size += len;
            }
    }
    //printf(NEW_LINE"pkName :%s"NEW_LINE, pkName);
    for(cnt = 0; cnt < cfg.image_num; cnt++)
    {
        if((0 != stat(cfg.image[cnt].path, &st))
            ||(NULL == (fps=fopen(cfg.image[cnt].path, "r"))))
        {
            printf(NEW_LINE"Cannot open image file(%s)."NEW_LINE, cfg.image[cnt].path);
            ret = RET_ENV_FAILED;
            goto ERROR_HANDLE;
        }

        /* Get file size */
        fseek(fps, 0, SEEK_END);
        fileSize = ftell(fps);
        fseek(fps, 0, SEEK_SET);

        /* Read file into buffer */
        fread(p_buf + imgHeaderSize, 1, fileSize, fps);
        fclose(fps);
        fps = NULL;

        memset(p_buf, 0x0, imgHeaderSize);
        p_saved_buf = p_buf;
        if((cfg.image[cnt].image_type == imgV1OEM) || (cfg.image[cnt].image_type == imgV1CFG))
        {
            // OEM 和 CFG 分区强行扩展至 cfg文件指定的大小
            //printf(NEW_LINE"    add size :%d"NEW_LINE, fileSize);
            if(cfg.image[cnt].length%IMAGE_ALIGNED_SIZE)
            {
                printf(NEW_LINE"Image space length invalid(0x%x)."NEW_LINE, cfg.image[cnt].length);
            }
            if(cfg.image[cnt].length - imgHeaderSize < fileSize)
            {
                printf(NEW_LINE"Image oversize(%s, size:%d, limited:%d)."NEW_LINE,
                    cfg.image[cnt].path, fileSize, cfg.image[cnt].length - imgHeaderSize);
                ret = RET_IMAGE_WRONG;
                goto ERROR_HANDLE;
            }
            crcSize = cfg.image[cnt].length - imgHeaderSize;
            pad_len = crcSize - fileSize;
            memset(p_buf + fileSize + imgHeaderSize, 0x0, pad_len);
            fileSize = crcSize;
        }
        else
        {
            pad_len = 0;
            crcSize = IMAGE_LEN_FOR_CRC(fileSize);
            /* Clear pad bytes to '0' */
            memset(p_buf + imgHeaderSize + fileSize, 0x0, crcSize - fileSize);
        }

        if((crcSize+imgHeaderSize)>cfg.image[cnt].length
            || (crcSize+imgHeaderSize) > PK_SINGLE_FILE_SIZE)
        {
            printf(NEW_LINE"Image oversize(%s, size:%d, limited:%d)."NEW_LINE,
                cfg.image[cnt].path, crcSize, cfg.image[cnt].length - imgHeaderSize);
            ret = RET_IMAGE_WRONG;
            goto ERROR_HANDLE;
        }

        /* for CRC */
        _CRC= 0xFFFFFFFF^CRC32_Update(0xFFFFFFFF, p_buf + imgHeaderSize, crcSize);

        ah_Sha256((unsigned char*)(p_buf + imgHeaderSize), fileSize, sha);

        /* fill header. */
        ah_fillHeader(p_buf, &cfg, cnt, fileSize, _CRC, sha);

        /* write image */
        len = crcSize + imgHeaderSize;

        /*update buffer pointer*/
        p_buf += len;

        // 在这里添加 oem_ext.bin
        if(cfg.image[cnt].image_type == imgV1FRWPLC){
            p_buf+=ah_add_oem_ext(p_buf);
        }

        /*new bin file's size*/
        len = p_buf - p_saved_buf;
        total_size += len;

        if(cfg.top_header_enable)
        {
            /*update image header*/
            offset = p_saved_buf - global_buffer;
            if(cfg.image[cnt].image_type == imgV1SBL)
            {
                hdr->sbl_start = htonl(offset);
                hdr->sbl_len = htonl(crcSize+imgHeaderSize);
            }
            if(cfg.image[cnt].image_type == imgV1OEM)
            {
                hdr->oem_start = htonl(offset);
                hdr->oem_len = htonl(crcSize+imgHeaderSize);
            }
            if(cfg.image[cnt].image_type == imgV1CFG)
            {
                hdr->pib_start = htonl(offset);
                hdr->pib_len = htonl(crcSize+imgHeaderSize);
            }

            if(cfg.image[cnt].image_type == imgV1FRWPLC)
            {
                static uint8_t fw_packed = 0;
                if(!fw_packed) {
                    hdr->fw_start = htonl(offset);
                    hdr->fw_len = htonl(crcSize+imgHeaderSize);
                    fw_packed = 1;
                } else {
                    hdr->fw2_start = htonl(offset);
                    hdr->fw2_len = htonl(crcSize+imgHeaderSize);
               }
            }

            if(cfg.image[cnt].image_type ==  imgV1FRWCUS)
            {
                hdr->cus_start = htonl(offset);
                hdr->cus_len = htonl(crcSize+imgHeaderSize);
            }
        }

         /* 2400 aligned */
        if((cnt+1 < cfg.image_num)
            &&((cfg.image[cnt+1].image_type == imgV1CFG)
            ||(cfg.image[cnt+1].image_type == imgV1FRWPLC)
            ||(cfg.image[cnt+1].image_type == imgV1FRWCUS)))
        {
            len = p_buf - global_buffer;
            if(len%2400)
            {
                len = 2400 - (len%2400);
                for(i=0; i<len; i++)
                {
                    *p_buf = 0;
                    p_buf++;
                }
                total_size += len;
            }
        }
    }

    if(cfg.top_header_enable)
    {
        char *st = NULL;
        char *end = NULL;
        uint8_t ptn  = 0;

        hdr->file_len = htonl(total_size);

        /*calculate payload's CRC*/
        cfg.crc = 0xFFFFFFFF^CRC32_Update(0xFFFFFFFF, payload , total_size);
        hdr->file_crc = htonl(cfg.crc);

        // 这里根据加密类型对整个镜像进行加密
        if(1 == cfg.encrypt_mode)
        {
            st = payload;
            end = payload + total_size;
            ptn = (cfg.encrypt_pattern&0xFF);
            while(st <= end) {
                *(st) += ptn;
                st++;
            }
        }
        else if(2 == cfg.encrypt_mode)
        {
            st = payload;
            end = payload + total_size;
            ptn = (cfg.encrypt_pattern&0xFF);
            while(st <= end) {
                *(st) -= ptn;
                st++;
            }
        }
        else if(3 == cfg.encrypt_mode)
        {
            st = global_buffer;
            end = payload + total_size;
            ptn = (cfg.encrypt_pattern&0xFF);
            while(st <= end) {
                *(st) ^= ptn;
                st++;
            }
        }
    }

    /*write buffer into file*/
    if(NULL == (fpd=fopen(pkName, "w+")))
    {
        printf(NEW_LINE"Cannot open package file(%s)."NEW_LINE, pkName);
        goto ERROR_HANDLE;
    }
    if(cfg.top_header_enable)
    {
        fwrite(global_buffer, 1, total_size + PACKAGE_HEADER_SIZE, fpd);
    }
    else
    {
        fwrite(global_buffer, 1, total_size, fpd);
    }
    fclose(fpd);
    fpd = NULL;

    /* show info */
    sh_argv[2] = pkName;
    ah_funShowHeader(3, sh_argv);
    printf(NEW_LINE"Package created - %s"NEW_LINE, pkName);
    ret = RET_OK;

ERROR_HANDLE:

    if(global_buffer)
    {
        free(global_buffer);
        global_buffer = NULL;
    }

    if(fpd)
    {
        fclose(fpd);
        if(RET_OK != ret)
            remove(pkName);
        fpd = NULL;
    }
    if(fps)
    {
        fclose(fps);
        fps = NULL;

    }
    return ret;
}

int ah_funShowHeader(char argc, char *argv[])
{
    FILE *fps = NULL;
    char pbuf[imgHeaderSize];
    int imgNo=0, imgType, tmp1, tmp2, tmp3, offset, start, device;

    if(argc < 3)
    {
        return RET_BAD_PARA;
    }

    if(NULL == (fps=fopen(argv[2], "r")))
    {
        printf(NEW_LINE"Cannot open package file(%s)."NEW_LINE, argv[2]);
        return RET_BAD_PARA;
    }

    do
    {
        imgNo++;

        while(32 == (tmp1 = fread(pbuf, 1, 32, fps)))
        {
            if(1==imgNo)
            {
                tmp3 = IMGHDR_V1_GET_GUARD(pbuf);
                if(tmp3 == (((int)IMAGE_GUARD)&0xFFFFFFFF))
                {
                    device = IMGHDR_V1_GET_DEVTYPE(pbuf);
                    if((((int)devV1Kunlun3)&0xFFFF) == device)
                    {
                        break;
                    }
                }
            }
            else
            {
                tmp2 = IMGHDR_V1_GET_DEVTYPE(pbuf);
                tmp3 = IMGHDR_V1_GET_GUARD(pbuf);

                if((tmp2 == device)
                    &&(tmp3 == (((int)IMAGE_GUARD)&0xFFFFFFFF)))
                    break;
            }
        }
        if(32 != tmp1)
        {
            fclose(fps);
            return RET_OK;
        }
        tmp1 = fread(pbuf+32, 1, 32, fps);
        if(32 != tmp1)
        {
            fclose(fps);
            printf(NEW_LINE"Maybe the package is damaged."NEW_LINE NEW_LINE);
            return RET_OK;
        }

        /* Print headers */
        start = ftell(fps) - imgHeaderSize;
        printf(NEW_LINE"Image NO.%02d @ 0x%08x:", imgNo, start);

        tmp1 = IMGHDR_V1_GET_DEVTYPE(pbuf);
        printf(NEW_LINE"    device:0x%02x", tmp1&0xFF);

        tmp1 = IMGHDR_V1_GET_IMGTYPE(pbuf);
        printf(NEW_LINE"    image:0x%02x", tmp1&0xFF);
        imgType = tmp1&0xFF;

        tmp1 = IMGHDR_V1_GET_ENCVAL(pbuf);
        printf(NEW_LINE"    encval:0x%02x", tmp1&0xFF);

        tmp1 =  IMGHDR_V1_GET_IMGSIZE(pbuf);
        offset =  IMAGE_LEN_FOR_CRC(tmp1);
        printf(NEW_LINE"    size:%d Bytes", tmp1);

        tmp1 = IMGHDR_V1_GET_ZIPTYPE(pbuf);
        printf(NEW_LINE"    zip type:%d", tmp1&0xFF);

        tmp1 = IMGHDR_V1_GET_FLASHSIZE(pbuf);
        printf(NEW_LINE"    flash size:%dM", tmp1&0xFF);

        tmp1 = IMGHDR_V1_GET_PSRAMSIZE(pbuf);
        printf(NEW_LINE"    psram size:%dM", tmp1&0xFF);

        tmp1 = IMGHDR_V1_GET_HDRVER(pbuf);
        printf(NEW_LINE"    hdr version:0x%02X", tmp1&0xFF);

        tmp1 = IMGHDR_V1_GET_IMGVER(pbuf);
        printf(NEW_LINE"    version:%d.%d.%d", (tmp1>>12)&0xF,
            (tmp1>>8)&0xF, tmp1&0xFF);

        tmp1 = IMGHDR_V1_GET_IMGCRC(pbuf);
        printf(NEW_LINE"    img crc:0x%08x", tmp1);

        tmp1 = IMGHDR_V1_GET_GUARD(pbuf);
        printf(NEW_LINE"    guard:0x%08x", tmp1);

        tmp1 = IMGHDR_V1_GET_LAYOUTIDX(pbuf);
        printf(NEW_LINE"    flash layout index:0x%02x", tmp1);

        if ((imgType == imgV1SBL) || (imgType == imgV1FRWPLC)
            || (imgType == imgV1FRWCUS)) {
            tmp1 = IMGHDR_V1_GET_RUNADDR(pbuf);
            printf(NEW_LINE"    run address:0x%08x", tmp1);
        }

        printf(NEW_LINE"    SHA256:");
        for(tmp1 = 32; tmp1 < 64; tmp1++)
        {
            tmp2 = ((int)pbuf[tmp1])&0xFF;
            printf("%02x", tmp2);
        }
        printf(NEW_LINE);
        /* skip size */
        if(fseek(fps, offset, SEEK_CUR))
        {
            fclose(fps);
            return RET_OK;
        }

    }while(1);

    return RET_OK;
}

int main(int argc, char *argv[])
{
    int ret = RET_BAD_PARA;

    if(argc <= 1)
    {
        ah_help();
        goto err_out;
    }

    if(0 == strcmp(FUNC_CONSTRUCT, argv[1]))
    {
        // -c 生成最终文件
        ret =ah_funConstruct(argc, argv);
    }
    else if(0 == strcmp(FUNC_SHOWHEADER, argv[1]))
    {
        ret =ah_funShowHeader(argc, argv);
    }
    else if(0 == strcmp(FUNC_ADDHEADER, argv[1]))
    {
        // -a 把bin文件添加img_header 然后用原文件名保存
        ret =ah_addHeader(argc, argv);
    }

    if(RET_OK == ret)
    {
        return 0;
    }

err_out:
    printf(NEW_LINE"Something wrong !"NEW_LINE NEW_LINE);
    return ret;
}