/*
 * Copyright (c) Hisilicon Technologies Co., Ltd. 2019-2021. All rights reserved.
 * Description: stbm r2r keyladder sample.
 * Author: Hisilicon
 * Create: 2019-09-19
 */
#include "securec.h"
#include "uapi_system.h"
#include "uapi_memory.h"
#include "uapi_klad.h"
#include "uapi_cipher.h"
#include "parse_config_file.h"

#define klad_err_print_hex(val)   printf("[%-32s][line:%04d]%s = 0x%08x\n", __FUNCTION__, __LINE__, #val, val)
#define klad_err_print_info(val)  printf("[%-32s][line:%04d]%s\n", __FUNCTION__, __LINE__, val)
#define klad_err_print_val(val)   printf("[%-32s][line:%04d]%s = %d\n", __FUNCTION__, __LINE__, #val, val)
#define klad_err_print_point(val) printf("[%-32s][line:%04d]%s = %p\n", __FUNCTION__, __LINE__, #val, val)
#define klad_print_error_code(err_code) \
    printf("[%-32s][line:%04d]return [0x%08x]\n", __FUNCTION__, __LINE__, err_code)
#define klad_print_error_func(func, err_code) \
    printf("[%-32s][line:%04d]call [%s] return [0x%08x]\n", __FUNCTION__, __LINE__, #func, err_code)

#define KEY_LEN 16
#define DATA_LEN 32
#define td_handle_makehandle(mod, privatedata, chnid) \
    (td_handle)((((mod) & 0xff) << 24) | ((((privatedata) & 0xff)<< 16)) | (((chnid) & 0xffff)))

typedef struct {
    td_handle handle_klad;
    td_handle handle_cipher;
    td_handle handle_ks;
} td_handle_klad;

static td_u8 g_data_in[DATA_LEN] = {0};

static td_u8 g_session_key[KEY_LEN] = {0};

static td_u8 g_content_key[KEY_LEN] = {0};

static td_void print_buffer(const td_char *string, const td_u8 *input, td_u32 length)
{
    td_u32 i;

    if (string != NULL) {
        printf("%s\n", string);
    }

    for (i = 0; i < length; i++) {
        if ((i % KEY_LEN == 0) && (i != 0)) {
            printf("\n");
        }
        printf("0x%02x ", input[i]);
    }
    printf("\n");

    return;
}

static td_s32 set_key(td_handle klad)
{
    td_s32 ret;
    errno_t errno;
    uapi_klad_session_key session_key = {0};
    uapi_klad_content_key content_key = {0};

    session_key.level = UAPI_KLAD_LEVEL1;
    session_key.key_size = KEY_LEN;
    session_key.alg = UAPI_KLAD_ALG_TYPE_AES;
    errno = memcpy_s(session_key.key, UAPI_KLAD_MAX_KEY_LEN, g_session_key, KEY_LEN);
    if (errno != EOK) {
        klad_print_error_func(memcpy_s, errno);
        return TD_FAILURE;
    }

    ret = uapi_klad_set_session_key(klad, &session_key);
    if (ret != TD_SUCCESS) {
        klad_print_error_func(uapi_klad_set_session_key, ret);
        return ret;
    }

    content_key.odd = 0;
    content_key.key_size = KEY_LEN;
    content_key.alg = UAPI_KLAD_ALG_TYPE_AES;
    errno = memcpy_s(content_key.key, UAPI_KLAD_MAX_KEY_LEN, g_content_key, KEY_LEN);
    if (errno != EOK) {
        klad_print_error_func(memcpy_s, errno);
        return TD_FAILURE;
    }

    ret = uapi_klad_set_content_key(klad, &content_key);
    if (ret != TD_SUCCESS) {
        klad_print_error_func(uapi_klad_set_content_key, ret);
        return ret;
    }

    return TD_SUCCESS;
}

static td_s32 cfg_cipher(td_handle cipher, uapi_cipher_symc_alg alg,
                         uapi_cipher_symc_work_mode mode, const td_u8 *iv, td_u32 iv_len)
{
    td_s32 ret;

    uapi_cipher_symc_ctrl_t cipher_ctrl = {0};

    cipher_ctrl.symc_alg = alg;
    cipher_ctrl.work_mode = mode;

    if (cipher_ctrl.symc_alg == UAPI_CIPHER_SYMC_ALG_AES) {
        cipher_ctrl.symc_key_length = UAPI_CIPHER_SYMC_KEY_128BIT;
        cipher_ctrl.symc_bit_width = UAPI_CIPHER_SYMC_BIT_WIDTH_128BIT;
        if (iv != TD_NULL) {
            cipher_ctrl.iv_change_flag = UAPI_CIPHER_SYMC_IV_CHANGE_ONE_PKG;
            cipher_ctrl.iv_length = iv_len;
            if (memcpy_s(cipher_ctrl.iv, sizeof(cipher_ctrl.iv), iv, iv_len) != EOK) {
                return TD_FAILURE;
            }
        }
    } else if (cipher_ctrl.symc_alg == UAPI_CIPHER_SYMC_ALG_TDES) {
        cipher_ctrl.symc_key_length = UAPI_CIPHER_SYMC_KEY_128BIT;
        cipher_ctrl.symc_bit_width = UAPI_CIPHER_SYMC_BIT_WIDTH_64BIT;
        if (iv != TD_NULL) {
            cipher_ctrl.iv_change_flag = UAPI_CIPHER_SYMC_IV_CHANGE_ONE_PKG;
            cipher_ctrl.iv_length = iv_len;
            if (memcpy_s(cipher_ctrl.iv, sizeof(cipher_ctrl.iv), iv, iv_len) != EOK) {
                return TD_FAILURE;
            }
        }
    } else if (cipher_ctrl.symc_alg == UAPI_CIPHER_SYMC_ALG_SM4) {
        if (iv != TD_NULL) {
            cipher_ctrl.iv_change_flag = UAPI_CIPHER_SYMC_IV_CHANGE_ONE_PKG;
            cipher_ctrl.iv_length = iv_len;
            if (memcpy_s(cipher_ctrl.iv, sizeof(cipher_ctrl.iv), iv, iv_len) != EOK) {
                return TD_FAILURE;
            }
        }
    } else {
        return TD_FAILURE;
    }

    ret = uapi_cipher_symc_set_config(cipher, &cipher_ctrl);
    if (ret != TD_SUCCESS) {
        klad_print_error_func(uapi_cipher_symc_set_config, ret);
        return ret;
    }

    return TD_SUCCESS;
}

static td_s32 priv_klad_init(td_handle_klad *klad_handle)
{
    td_s32 ret;
    uapi_cipher_symc_attr cipher_attr = {0};

    ret = uapi_sys_init();
    if (ret != TD_SUCCESS) {
        klad_print_error_func(uapi_sys_init, ret);
        return ret;
    }

    ret = uapi_cipher_symc_init();
    if (ret != TD_SUCCESS) {
        klad_print_error_func(uapi_cipher_init, ret);
        goto sys_deinit;
    }

    ret = uapi_klad_init();
    if (ret != TD_SUCCESS) {
        klad_print_error_func(uapi_klad_init, ret);
        goto cipher_deinit;
    }

    ret = uapi_klad_create(&klad_handle->handle_klad);
    if (ret != TD_SUCCESS) {
        klad_print_error_func(uapi_klad_create, ret);
        goto klad_deinit;
    }

    cipher_attr.symc_type = UAPI_CIPHER_SYMC_TYPE_NORMAL;
    cipher_attr.symc_alg = UAPI_CIPHER_SYMC_ALG_AES;
    cipher_attr.work_mode = UAPI_CIPHER_SYMC_WORK_MODE_CBC;
    cipher_attr.is_long_term = TD_TRUE;
    ret = uapi_cipher_symc_create(&klad_handle->handle_cipher, &cipher_attr);
    if (ret != TD_SUCCESS) {
        goto klad_destroy;
    }

    return TD_SUCCESS;

klad_destroy:
    uapi_klad_destroy(klad_handle->handle_klad);
klad_deinit:
    uapi_klad_deinit();
cipher_deinit:
    uapi_cipher_symc_deinit();
sys_deinit:
    uapi_sys_deinit();
    return ret;
}


static td_s32 priv_klad_attr(const td_handle_klad *klad_handle)
{
    td_s32 ret;
    uapi_klad_attr attr_klad = {0};

    attr_klad.klad_cfg.owner_id = 0x01;
    attr_klad.klad_cfg.klad_type = UAPI_KLAD_TYPE_R2R;
    attr_klad.key_cfg.decrypt_support = 1;
    attr_klad.key_cfg.encrypt_support = 1;
    attr_klad.key_cfg.engine = UAPI_CRYPTO_ALG_RAW_AES;
    ret = uapi_klad_set_attr(klad_handle->handle_klad, &attr_klad);
    if (ret != TD_SUCCESS) {
        klad_print_error_func(uapi_klad_set_attr, ret);
    }

    return ret;
}

static td_s32 priv_klad_mem_addr_vir(td_mem_handle *src_handle, td_mem_handle *dst_handle,
                                     td_u8 **input_addr_vir, td_u8 **output_addr_vir)
{
    src_handle->addr_offset = 0;
    src_handle->mem_handle = uapi_mem_new("cipher_buf_in", strlen("cipher_buf_in"), DATA_LEN, 0);
    if (src_handle->mem_handle == 0) {
        klad_err_print_info("error: get phyaddr for input failed!\n");
        goto out;
    }

    *input_addr_vir = uapi_mem_map(src_handle->mem_handle, DATA_LEN);

    dst_handle->addr_offset = 0;
    dst_handle->mem_handle = uapi_mem_new("cipher_buf_out", strlen("cipher_buf_out"), DATA_LEN, 0);
    if (dst_handle->mem_handle == 0) {
        klad_err_print_info("error: get phyaddr for out_put failed!\n");
        goto klad_umem;
    }
    *output_addr_vir = uapi_mem_map(dst_handle->mem_handle, DATA_LEN);

    return TD_SUCCESS;
klad_umem:
    if (src_handle->mem_handle > 0) {
        uapi_mem_unmap(*input_addr_vir, DATA_LEN);
        uapi_mem_delete(src_handle->mem_handle);
    }
out:
    return TD_FAILURE;
}

static td_void priv_klad_unmem_addr_vir(const td_mem_handle *src_handle, const td_mem_handle *dst_handle,
                                        td_u8 *input_addr_vir, td_u8 *output_addr_vir);

static td_s32 priv_klad_attach(td_handle_klad *klad_handle, td_mem_handle *src_handle, td_mem_handle *dst_handle,
                               td_u8 **input_addr_vir, td_u8 **output_addr_vir)
{
    td_s32 ret;
    td_u8 iv[KEY_LEN] = {0};

    ret = priv_klad_mem_addr_vir(src_handle, dst_handle, input_addr_vir, output_addr_vir);
    if (ret != TD_SUCCESS) {
        klad_print_error_func(priv_klad_mem_addr_vir, ret);
        goto out;
    }

    ret = priv_klad_attr(klad_handle);
    if (ret != TD_SUCCESS) {
        klad_print_error_func(priv_klad_attr, ret);
        goto klad_umem;
    }

    ret = uapi_cipher_symc_get_keyslot_handle(klad_handle->handle_cipher, &klad_handle->handle_ks);
    if (ret != TD_SUCCESS) {
        klad_print_error_func(uapi_cipher_symc_get_keyslot_handle, ret);
        goto klad_umem;
    }

    ret = uapi_klad_attach(klad_handle->handle_klad, klad_handle->handle_ks);
    if (ret != TD_SUCCESS) {
        klad_print_error_func(uapi_klad_attach, ret);
        goto klad_umem;
    }

    ret = cfg_cipher(klad_handle->handle_cipher, UAPI_CIPHER_SYMC_ALG_AES, UAPI_CIPHER_SYMC_WORK_MODE_CBC, iv, KEY_LEN);
    if (ret != TD_SUCCESS) {
        klad_print_error_func(cfg_cipher, ret);
        goto klad_detach;
    }

    ret = set_key(klad_handle->handle_klad);
    if (ret != TD_SUCCESS) {
        klad_print_error_func(set_key, ret);
        goto klad_detach;
    }

    return TD_SUCCESS;
klad_detach:
    uapi_klad_detach(klad_handle->handle_klad, klad_handle->handle_ks);
klad_umem:
    priv_klad_unmem_addr_vir(src_handle, dst_handle, *input_addr_vir, *output_addr_vir);
out:
    return ret;
}

static td_s32 priv_klad_encrypt(const td_handle_klad *klad_handle, const td_mem_handle *src_handle,
                                const td_mem_handle *dst_handle, td_u8 *input_addr_vir, td_u8 *output_addr_vir)
{
    errno_t errno;
    td_s32 ret = TD_FAILURE;
    const td_mem_size_t data_len = DATA_LEN;
    uapi_cipher_buf_attr src_buf = {0};
    uapi_cipher_buf_attr dest_buf = {0};

    errno = memset_s(input_addr_vir, data_len, 0x0, data_len * sizeof(td_u8));
    if (errno != EOK) {
        klad_print_error_func(memset_s, errno);
        goto out;
    }
    errno = memcpy_s(input_addr_vir, data_len, g_data_in, data_len * sizeof(td_u8));
    if (errno != EOK) {
        klad_print_error_func(memcpy_s, errno);
        goto out;
    }
    errno  = memset_s(output_addr_vir, data_len, 0x0, data_len * sizeof(td_u8));
    if (errno != EOK) {
        klad_print_error_func(memset_s, errno);
        goto out;
    }
    src_buf.uapi_mem_handle = src_handle->mem_handle;
    dest_buf.uapi_mem_handle = dst_handle->mem_handle;
    ret = uapi_cipher_symc_encrypt(klad_handle->handle_cipher, &src_buf, &dest_buf, data_len);
    if (ret != TD_SUCCESS) {
        klad_print_error_func(uapi_cipher_encrypt, ret);
        goto out;
    }

    print_buffer("clear data:", g_data_in, data_len);
    print_buffer("encrypted data:", output_addr_vir, data_len);

out:
    return ret;
}

static td_void priv_klad_unmem_addr_vir(const td_mem_handle *src_handle, const td_mem_handle *dst_handle,
                                        td_u8 *input_addr_vir, td_u8 *output_addr_vir)
{
    if (src_handle->mem_handle > 0) {
        uapi_mem_unmap(input_addr_vir, DATA_LEN);
        uapi_mem_delete(src_handle->mem_handle);
    }
    if (dst_handle->mem_handle > 0) {
        uapi_mem_unmap(output_addr_vir, DATA_LEN);
        uapi_mem_delete(dst_handle->mem_handle);
    }

    return;
}

static td_void priv_klad_detach(const td_handle_klad *klad_handle,
                                const td_mem_handle *src_handle, const td_mem_handle *dst_handle,
                                td_u8 *input_addr_vir, td_u8 *output_addr_vir)
{
    td_s32 ret;

    ret = uapi_klad_detach(klad_handle->handle_klad, klad_handle->handle_ks);
    if (ret != TD_SUCCESS) {
        klad_print_error_func(uapi_klad_detach, ret);
    }

    priv_klad_unmem_addr_vir(src_handle, dst_handle, input_addr_vir, output_addr_vir);
    return;
}

static td_void priv_klad_deinit(const td_handle_klad *klad_handle)
{
    uapi_cipher_symc_destroy(klad_handle->handle_cipher);
    uapi_klad_destroy(klad_handle->handle_klad);

    uapi_klad_deinit();
    uapi_cipher_symc_deinit();
    uapi_sys_deinit();
    return;
}

static td_s32 priv_data_init(const td_char *cfg_path)
{
    td_s32 ret;
    td_s32 value_num = 0;

    ret = parse_config_file(cfg_path, &value_num);
    if (ret != TD_SUCCESS || value_num == 0 || value_num > MAX_VAR_NUM) {
        klad_print_error_func(parse_config_file, ret);
        goto out;
    }

    ret = get_key_value(SESSIONKEY_TAG, g_session_key, KEY_LEN);
    if (ret != TD_SUCCESS) {
        klad_print_error_func(get_key_value, ret);
        goto out;
    }

    ret = get_key_value(CONTENTKEY_TAG, g_content_key, KEY_LEN);
    if (ret != TD_SUCCESS) {
        klad_print_error_func(get_key_value, ret);
        goto out;
    }

    ret = get_key_value(CIPHER_INPUT_DATA, g_data_in, DATA_LEN);
    if (ret != TD_SUCCESS) {
        klad_print_error_func(get_key_value, ret);
        goto out;
    }
out:
    return ret;
}

td_s32 main(td_s32 argc, td_char *argv[])
{
    td_s32 ret;
    td_handle_klad klad_handle = {0};
    td_u8 *input_addr_vir = TD_NULL;
    td_u8 *output_addr_vir = TD_NULL;
    td_mem_handle src_handle;
    td_mem_handle dst_handle;
    (void)argc;

    if (argv == TD_NULL || argv[0x1] == TD_NULL) {
        printf("Please enter the path of the klad_data_cfg.ini file.!!!\n");
        return TD_FAILURE;
    }

    ret = priv_data_init(argv[0x1]);
    if (ret != TD_SUCCESS) {
        klad_print_error_func(priv_data_init, ret);
        return ret;
    }

    ret = priv_klad_init(&klad_handle);
    if (ret != TD_SUCCESS) {
        return ret;
    }

    ret = priv_klad_attach(&klad_handle, &src_handle, &dst_handle, &input_addr_vir, &output_addr_vir);
    if (ret != TD_SUCCESS) {
        klad_print_error_func(priv_klad_attach, ret);
        goto out;
    }

    ret = priv_klad_encrypt(&klad_handle, &src_handle, &dst_handle, input_addr_vir, output_addr_vir);
    if (ret != TD_SUCCESS) {
        klad_print_error_func(priv_klad_encrypt, ret);
    }

    priv_klad_detach(&klad_handle, &src_handle, &dst_handle, input_addr_vir, output_addr_vir);
out:
    priv_klad_deinit(&klad_handle);
    return ret;
}