#include "kg_dh_cmd.h"

uint32_t KG_generate_dh_request(tz_dh_req_payload_t *sendmsg, tz_dh_req_payload_t *respmsg) {

    KG_LOG("KG generate dh request\n");
    uint32_t ret = KG_DH_GEN_DH_FAIL;
    uint8_t *dh_pub_key = NULL;
    int dh_pub_key_len = KG_DH_PUB_KEY_LEN;
    uint8_t *dh_priv_key = NULL;
    int dh_priv_key_len = KG_DH_PRIV_KEY_LEN;

    uint8_t *aes_key = NULL;
    uint32_t aes_key_len = KG_AES_KEY_LEN;
    uint8_t *aes_iv = NULL;
    uint32_t aes_iv_len = KG_AES_IV_LEN;
    uint8_t *aes_tag = NULL;
    uint32_t aes_tag_len = KG_AES_TAG_LEN;

    uint8_t *enc_dh_pub_key = NULL;
    uint32_t enc_dh_pub_key_len = KG_DH_PUB_KEY_LEN;
    uint8_t *enc_key_buffer = NULL;
    uint32_t enc_key_buffer_len = KG_BUF_LEN;

    uint8_t *dh_buffer = NULL;
    uint32_t dh_buffer_len = KG_BUF_LEN;
    uint8_t *dh_buffer_b64 = NULL;
    uint32_t dh_buffer_b64_len = KG_B64_BUF_LEN;

    kg_rpmb_info_t* info = NULL;
    uint8_t* wrap_data = NULL;
    uint8_t* unwrap_data = NULL;
    uint32_t unwrap_data_len = KG_BUF_LEN;
    uint8_t* rewrap_data = NULL;
    uint32_t rewrap_data_len = KG_SECURE_DATA_LEN;
    
    EVP_PKEY *pkey = NULL;
    RSA *rsakey = NULL;

    if (sendmsg->payload.resp.data_len > KG_MAX_PAYLOAD_LEN) {
        KG_LOG("Recieved invalid input buffer when generating dh request\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

    aes_key = TEE_Malloc(aes_key_len, 0);
    if (NULL == aes_key) {
        KG_LOG("Failed to alloc buffer for aes key\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    aes_iv = TEE_Malloc(aes_iv_len, 0);
    if (NULL == aes_iv) {
        KG_LOG("Failed to alloc buffer for aes iv\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    aes_tag = TEE_Malloc(aes_tag_len, 0);
    if (NULL == aes_tag) {
        KG_LOG("Failed to alloc buffer for aes tag\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    enc_dh_pub_key = TEE_Malloc(enc_dh_pub_key_len, 0);
    if (NULL == enc_dh_pub_key) {
        KG_LOG("Failed to alloc buffer for enc dh pub key\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    enc_key_buffer = TEE_Malloc(enc_key_buffer_len, 0);
    if (NULL == enc_key_buffer) {
        KG_LOG("Failed to alloc buffer for enc key\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    dh_buffer = TEE_Malloc(dh_buffer_len, 0);
    if (NULL == dh_buffer) {
        KG_LOG("Failed to alloc buffer for dh buffer\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    dh_buffer_b64 = TEE_Malloc(dh_buffer_b64_len, 0);
    if (NULL == dh_buffer_b64) {
        KG_LOG("Failed to alloc buffer for dh buffer b64\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

#ifdef CONFIG_QSEE
    if (KG_SUCCESS != (ret = kg_rpmb_init())) {
        KG_LOG("KG rpmb is unavailable when verifying policy\n");
        goto exit;
    }
#endif    

    if (KG_SUCCESS != (ret = read_info_object(&info))
        || NULL == info) {
        KG_LOG("failed to read info object\n");
        goto exit;
    }

    if(info->kg_wrap_data_len > KG_SECURE_DATA_LEN){
        KG_LOG("Received invalid size for wrap_data\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

#ifdef CONFIG_QSEE
/* error: taking address of packed member 'kg_wrap_data_len' of class or structure 'kg_rpmb_info' may result in an unaligned pointer value [-Werror,-Waddress-of-packed-member]*/ 
    uint32_t temp_wrap_data_len = info->kg_wrap_data_len;
    if (KG_SUCCESS != (ret = read_wrap_data(&temp_wrap_data_len, &wrap_data))
        || NULL == wrap_data) {
        KG_LOG("failed to read wrap data\n");
        goto exit;
    }
    info->kg_wrap_data_len = temp_wrap_data_len;
#else
    if (KG_SUCCESS != (ret = read_wrap_data(&(info->kg_wrap_data_len), &wrap_data))
        || NULL == wrap_data) {
        KG_LOG("failed to read wrap data\n");
        goto exit;
    }
#endif
    unwrap_data = TEE_Malloc(unwrap_data_len, 0);
    if (NULL == unwrap_data) {
        KG_LOG("failed to alloc unwrap data\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

#ifdef CONFIG_QSEE
    if (TZ_API_OK != TZ_unwrap_persist_data((uint8_t *)KG_NAME, strlen(KG_NAME), 
        wrap_data, info->kg_wrap_data_len, unwrap_data, &unwrap_data_len)) {
        KG_LOG("Failed to unwrap data\n");
        ret = KG_TZ_API_FAIL;
        goto exit;
    }
#else
    if (TZ_API_OK != TZ_unwrap_data_with_derived_key((uint8_t *)KG_NAME, strlen(KG_NAME), 
        wrap_data, info->kg_wrap_data_len, unwrap_data, &unwrap_data_len)) {
        KG_LOG("failed to unwrap data\n")
        ret = KG_TZ_API_FAIL;
        goto exit;
    }
#endif

    if (unwrap_data_len != sizeof(kg_secure_data_t)) {
        KG_LOG("KG TA recovering unwraped secure data size check failed\n");
        ret = KG_RPMB_UNWRAP_FAIL;
        goto exit;
    }

    kg_secure_data_t* secure_data = (kg_secure_data_t*)unwrap_data;
    
    if (KG_SUCCESS != KG_dh_get_keypair(&dh_pub_key, &dh_pub_key_len, &dh_priv_key, &dh_priv_key_len)) {
        KG_LOG_DBG("Failed to generate dh keypair\n");
        ret = KG_DH_GEN_KEYPAIR_FAIL;
        goto exit;
    }
    
    KG_DUMP_DBG("dh pub key: \n", dh_pub_key, dh_pub_key_len);
    KG_DUMP_DBG("dh priv key: \n", dh_priv_key, dh_priv_key_len);

    secure_data->kg_dh_data.dh_state = DH_STATE_NEGO;
    TEE_MemMove(secure_data->kg_dh_data.dh_key, dh_priv_key, dh_priv_key_len);

    // store priv key
    rewrap_data = TEE_Malloc(rewrap_data_len, 0);
    if (NULL == rewrap_data) {
        KG_LOG("KG TA failed to alloc buffer to hold wrap data\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

#ifdef CONFIG_QSEE
    if (TZ_API_OK != TZ_wrap_persist_data((uint8_t *)KG_NAME, strlen(KG_NAME), 
        unwrap_data, sizeof(kg_secure_data_t), rewrap_data, &rewrap_data_len)) {
        KG_LOG("Failed to wrap kg secure data structure in QC\n");
        ret = KG_TZ_API_FAIL;
        goto exit;
    }
#else
    if (TZ_API_OK != TZ_wrap_data_with_derived_key((uint8_t *)KG_NAME, strlen(KG_NAME), 
        unwrap_data, sizeof(kg_secure_data_t), rewrap_data, &rewrap_data_len)) {
        KG_LOG("Failed to wrap kg secure data structure in QC\n");
        ret = KG_TZ_API_FAIL;
        goto exit;
    }
#endif

    if (rewrap_data_len > KG_SECURE_DATA_LEN) {
        KG_LOG("Wraped kg secure data size overflow\n");
        ret = KG_RPMB_WRAP_FAIL;
        goto exit;
    }

    info->kg_wrap_data_len = rewrap_data_len;
    if (KG_SUCCESS != (ret = write_wrap_data(info->kg_wrap_data_len, rewrap_data))) {
        KG_LOG("failed to write wrap data\n");
        goto exit;
    }

    if (KG_SUCCESS != (ret = write_info_object(info))) {
        KG_LOG("failed to write info object\n");
        goto exit;
    }

    if (TZ_API_OK != TZ_gen_rand_data(aes_key, &aes_key_len)) {
        KG_LOG_DBG("Failed to generate aes key");
        ret = KG_DH_GEN_AES_FAIL;
        goto exit;
    }
    KG_DUMP_DBG("aes key: \n", aes_key, aes_key_len);
    
    if (TZ_API_OK != TZ_gen_rand_data(aes_iv, &aes_iv_len)) {
        KG_LOG_DBG("Failed to generate aes iv");
        ret = KG_DH_GEN_IV_FAIL;
        goto exit;
    }
    KG_DUMP_DBG("aes iv: \n", aes_iv, aes_iv_len);

    if (KG_SUCCESS != KG_gcm_encrypt(dh_pub_key, dh_pub_key_len, aes_key, aes_iv, aes_iv_len, enc_dh_pub_key, (int *)&enc_dh_pub_key_len, aes_tag)) {
        KG_LOG_DBG("gcm encryption failed\n");
        ret = KG_CRYPTO_AES_GCM_FAIL;
        goto exit;
    }
    KG_DUMP_DBG("enc_dh_pub_key: \n", enc_dh_pub_key, enc_dh_pub_key_len);
    KG_DUMP_DBG("aes tag: \n", aes_tag, aes_tag_len);

    TEE_MemFill(dh_buffer, 0x0, dh_buffer_len);
    if (aes_iv_len + aes_key_len + aes_tag_len > dh_buffer_len) {
        KG_LOG("dh buffer overflow\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

    dh_buffer_len = 0;
    TEE_MemMove(dh_buffer, (void *)aes_key, aes_key_len);
    dh_buffer_len += aes_key_len;
    TEE_MemMove(dh_buffer + dh_buffer_len, (void *)aes_iv, aes_iv_len);
    dh_buffer_len += aes_iv_len;
    TEE_MemMove(dh_buffer + dh_buffer_len, (void *)aes_tag, aes_tag_len);
    dh_buffer_len += aes_tag_len;

    KG_DUMP_DBG("dh buffer: \n", dh_buffer, dh_buffer_len);

    if (secure_data->kg_metadata.reg_info == REGION_EU) {
		KG_LOG_DBG("reg_info is EU\n");
        pkey = get_public_key(hotp_provision_cert_EU);
    } else if (secure_data->kg_metadata.reg_info == REGION_US) {
		KG_LOG_DBG("reg_info is US\n");
        pkey = get_public_key(hotp_provision_cert_US);
    } else {
		KG_LOG_DBG("reg_info is not EU/US\n");
	}

    if (NULL == pkey) {
        KG_LOG_DBG("Invalid pem cert to extract the public key\n");
        goto exit;
    }

    rsakey = EVP_PKEY_get1_RSA(pkey);
    if (!rsakey) {
        KG_LOG_DBG("Failed to create RSA key\n");
        goto exit;
    }

    if (KG_SUCCESS != kg_public_encrypt(dh_buffer, dh_buffer_len, rsakey, enc_key_buffer, (int *)&enc_key_buffer_len)) {
        KG_LOG_DBG("Failed to encrypted client key to generate key info\n");
        ret = KG_CRYPTO_RSA_ENC_FAIL;
        goto exit;
    }

    KG_DUMP_DBG("client encryption key info: \n", enc_key_buffer, enc_key_buffer_len);

    if (enc_key_buffer_len + enc_dh_pub_key_len > KG_BUF_LEN) {
        KG_LOG("KG generated dh response size check failed\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

    TEE_MemFill(dh_buffer, 0x0, KG_BUF_LEN);
    dh_buffer_len = 0;

    TEE_MemMove(dh_buffer, (void *)enc_dh_pub_key, enc_dh_pub_key_len);
    dh_buffer_len += enc_dh_pub_key_len;
    TEE_MemMove(dh_buffer + dh_buffer_len, (void *)enc_key_buffer, enc_key_buffer_len);
    dh_buffer_len += enc_key_buffer_len;

    //KG_DUMP_DBG("result bytes: \n", dh_buffer, dh_buffer_len);

    if (BASE64_OK != base64_encode(dh_buffer, dh_buffer_len, dh_buffer_b64, &dh_buffer_b64_len)) {
        KG_LOG("KG base64 encoding failed when generating dh request\n");
        ret = KG_BASE64_ENCODE_FAIL;
        goto exit;
    }

    if (dh_buffer_b64_len > KG_B64_BUF_LEN) {
        KG_LOG("Encoded string length check failed\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

    TEE_MemFill(respmsg->payload.resp.data_buf, 0x0, KG_MAX_PAYLOAD_LEN);
    if(dh_buffer_b64_len > KG_MAX_PAYLOAD_LEN){
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }
    respmsg->payload.resp.data_len = dh_buffer_b64_len;
    TEE_MemMove(respmsg->payload.resp.data_buf, (void *)dh_buffer_b64, dh_buffer_b64_len);
    KG_LOG_DBG("Encoded result string is %s\n", dh_buffer_b64);

    ret = KG_SUCCESS;
exit:
    if (dh_pub_key != NULL) {
        TEE_MemFill(dh_pub_key, 0, dh_pub_key_len);
        TEE_Free(dh_pub_key);
        dh_pub_key = NULL;
    }
    if (dh_priv_key != NULL) {
        TEE_MemFill(dh_priv_key, 0, dh_priv_key_len);
        TEE_Free(dh_priv_key);
        dh_priv_key = NULL;
    }
    if (aes_key != NULL) {
        TEE_MemFill(aes_key, 0, aes_key_len);
        TEE_Free(aes_key);
        aes_key = NULL;
    }
    if (aes_iv != NULL) {
        TEE_Free(aes_iv);
        aes_iv = NULL;
    }
    if (aes_tag != NULL) {
        TEE_Free(aes_tag);
        aes_tag = NULL;
    }
    if (enc_dh_pub_key != NULL) {
        TEE_Free(enc_dh_pub_key);
        enc_dh_pub_key = NULL;
    }
    if (enc_key_buffer != NULL) {
        TEE_Free(enc_key_buffer);
        enc_key_buffer = NULL;
    }
    if (dh_buffer != NULL) {
        TEE_Free(dh_buffer);
        dh_buffer = NULL;
    }
    if (dh_buffer_b64 != NULL) {
        TEE_Free(dh_buffer_b64);
        dh_buffer_b64 = NULL;
    }
    if (pkey != NULL) {
        EVP_PKEY_free(pkey);
    }
    if (rsakey != NULL) {
        RSA_free(rsakey);
    }
    if (wrap_data != NULL) {
        TEE_Free(wrap_data);
        wrap_data = NULL;
    }
    if (unwrap_data != NULL) {
        TEE_MemFill(unwrap_data, 0, unwrap_data_len);
        TEE_Free(unwrap_data);
        unwrap_data = NULL;
    }
    if (rewrap_data != NULL) {
        TEE_Free(rewrap_data);
        rewrap_data = NULL;
    }
    if (info != NULL) {
        TEE_Free(info);
        info = NULL;
    }
    return ret; 
}

uint32_t KG_verify_dh_response(tz_dh_verify_payload_t *sendmsg, tz_dh_verify_payload_t *respmsg) {
    KG_LOG("KG verify dh response");
    uint32_t ret = KG_SUCCESS;
    uint8_t *svr_pub_buffer = NULL;
    uint32_t svr_pub_buffer_len = KG_BUF_LEN;
    uint8_t *svr_pub_sig_buffer = NULL;
    uint32_t svr_pub_sig_buffer_len = KG_BUF_LEN;
    uint8_t *pub_mod = NULL;
    uint32_t pub_mod_len = KG_BUF_LEN;
    uint8_t *pub_exp = NULL;
    uint32_t pub_exp_len = KG_BUF_LEN;
    uint8_t *dh_secret = NULL;
    uint32_t dh_secret_len = KG_BUF_LEN;
    uint8_t *hotp_secret_key = NULL;
    uint32_t hotp_secret_key_len = KG_KEY_LEN;
    uint8_t *cha_check = NULL;
    uint32_t cha_check_len = KG_HOTP_LEN;

    kg_rpmb_info_t* info = NULL;
    uint8_t* wrap_data = NULL;
    uint8_t* unwrap_data = NULL;
    uint32_t unwrap_data_len = KG_BUF_LEN;
    uint8_t* rewrap_data = NULL;
    uint32_t rewrap_data_len = KG_SECURE_DATA_LEN;


    uint8_t seed[1] = {0x30};

    if (sendmsg->payload.cmd.dh_pub_len > KG_SERVER_DH_PUB_MAX ||
        sendmsg->payload.cmd.dh_pub_sign_len > KG_SERVER_DH_PUB_SIGN_MAX ||
        sendmsg->payload.cmd.dh_challenge_len > KG_DH_CHALLENGE_MAX) {
        KG_LOG("Recieved invalid input buffer when verifying dh response\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

    svr_pub_buffer = TEE_Malloc(svr_pub_buffer_len, 0);
    if (NULL == svr_pub_buffer) {
        KG_LOG("Failed to alloc buffer for svr public key\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    svr_pub_sig_buffer = TEE_Malloc(svr_pub_sig_buffer_len, 0);
    if (NULL == svr_pub_sig_buffer) {
        KG_LOG("Failed to alloc buffer for svr public key signature\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    pub_mod = TEE_Malloc(pub_mod_len, 0);
    if (NULL == pub_mod) {
        KG_LOG("Failed to alloc buffer for pub_mod\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    pub_exp = TEE_Malloc(pub_exp_len, 0);
    if (NULL == pub_exp) {
        KG_LOG("Failed to alloc buffer for pub_mod\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    dh_secret = TEE_Malloc(dh_secret_len, 0);
    if (NULL == dh_secret) {
        KG_LOG("Failed to alloc buffer for dh_secret\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    hotp_secret_key = TEE_Malloc(hotp_secret_key_len, 0);
    if (NULL == hotp_secret_key) {
        KG_LOG("Failed to alloc key buffer for hotp secret key\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    cha_check = TEE_Malloc(cha_check_len, 0);
    if (NULL == cha_check) {
        KG_LOG("Failed to alloc buffer for challenge check\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    KG_LOG_DBG("input server dh_pub b64 is %s\n", sendmsg->payload.cmd.dh_pub);
    KG_LOG_DBG("input server dh_pub b64 len is %d\n", sendmsg->payload.cmd.dh_pub_len);
    if (BASE64_OK != base64_decode(sendmsg->payload.cmd.dh_pub,
            sendmsg->payload.cmd.dh_pub_len, svr_pub_buffer, &svr_pub_buffer_len)) {
        KG_LOG("Failed to decode base64 encoded input svr dh public key\n");
        ret = KG_BASE64_DECODE_FAIL;
        goto exit;
    }
    if (svr_pub_buffer_len != KG_DH_PUB_KEY_LEN) {
        KG_LOG("Decoded svr public key length check failed\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }
    if (BASE64_OK != base64_decode(sendmsg->payload.cmd.dh_pub_sign,
            sendmsg->payload.cmd.dh_pub_sign_len, svr_pub_sig_buffer, &svr_pub_sig_buffer_len)) {
        KG_LOG("Failed to decode base64 encoded input svr dh public key signature\n");
        ret = KG_BASE64_DECODE_FAIL;
        goto exit;
    }
    if (svr_pub_sig_buffer_len != KG_SIG_LEN) {
        KG_LOG("Decoded svr dh public key signature buffer length check failed\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

    if (KG_SUCCESS != kg_rpmb_init()) {
        KG_LOG("RPMB is not available when initing KG secure data\n");
        ret = KG_RPMB_UNAVAILABLE;
        goto exit;
    }

    if (KG_SUCCESS != (ret = read_info_object(&info))
        || NULL == info) {
        KG_LOG("failed to read info object\n");
        goto exit;
    }

    if (info->kg_state == KG_STATE_LOCK) {
        KG_LOG("KG verify DH response state check failed\n");
        ret = KG_DH_STATE_CHECK_FAIL;
        goto exit;
    }

    if(info->kg_wrap_data_len > KG_SECURE_DATA_LEN){
        KG_LOG("Received invalid size for wrap_data\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

#ifdef CONFIG_QSEE
/* error: taking address of packed member 'kg_wrap_data_len' of class or structure 'kg_rpmb_info' may result in an unaligned pointer value [-Werror,-Waddress-of-packed-member]*/ 
    uint32_t temp_wrap_data_len = info->kg_wrap_data_len;
    if (KG_SUCCESS != (ret = read_wrap_data(&temp_wrap_data_len, &wrap_data))
        || NULL == wrap_data) {
        KG_LOG("failed to read wrap data\n");
        goto exit;
    }
    info->kg_wrap_data_len = temp_wrap_data_len;
#else
    if (KG_SUCCESS != (ret = read_wrap_data(&(info->kg_wrap_data_len), &wrap_data))
        || NULL == wrap_data) {
        KG_LOG("failed to read wrap data\n");
        goto exit;
    }
#endif    
    unwrap_data = TEE_Malloc(unwrap_data_len, 0);
    if (NULL == unwrap_data) {
        KG_LOG("failed to alloc unwrap data\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

#ifdef CONFIG_QSEE
    if (TZ_API_OK != TZ_unwrap_persist_data((uint8_t *)KG_NAME, strlen(KG_NAME), 
        wrap_data, info->kg_wrap_data_len, unwrap_data, &unwrap_data_len)) {
        KG_LOG("Failed to unwrap data\n");
        ret = KG_TZ_API_FAIL;
        goto exit;
    }
#else
    if (TZ_API_OK != TZ_unwrap_data_with_derived_key((uint8_t *)KG_NAME, strlen(KG_NAME), 
        wrap_data, info->kg_wrap_data_len, unwrap_data, &unwrap_data_len)) {
        KG_LOG("failed to unwrap data\n")
        ret = KG_TZ_API_FAIL;
        goto exit;
    }
#endif    

    kg_secure_data_t* secure_data = (kg_secure_data_t*)unwrap_data;
    if (secure_data->kg_metadata.reg_info == REGION_EU) {
        if (KG_SUCCESS != extract_public_keybytes(hotp_provision_cert_EU, pub_mod, &pub_mod_len, pub_exp, &pub_exp_len)) {
            KG_LOG_DBG("Failed to extract publick key bytes from hotp provisioning cert EU\n");
            ret = KG_CRYPTO_PKEY_PARSE_FAIL;
            goto exit;
        }
    } else if (secure_data->kg_metadata.reg_info == REGION_US) {
        if (KG_SUCCESS != extract_public_keybytes(hotp_provision_cert_US, pub_mod, &pub_mod_len, pub_exp, &pub_exp_len)) {
            KG_LOG_DBG("Failed to extract publick key bytes from hotp provisioning cert US\n");
            ret = KG_CRYPTO_PKEY_PARSE_FAIL;
            goto exit;
        }
    } else {
        KG_LOG("target region is unknown %d\n", secure_data->kg_metadata.reg_info);
        ret = KG_REGION_NOT_CONFIG;
        goto exit;
    }

    KG_DUMP_DBG("dump pubkey n:\n", pub_mod, pub_mod_len);
    KG_DUMP_DBG("dump pubkey e:\n", pub_exp, pub_exp_len);
    bool verified = false;
    if (TZ_API_OK != TZ_verify_CKM_SHA256_RSA_PKCS(pub_mod, pub_mod_len, pub_exp, pub_exp_len,
            svr_pub_buffer, svr_pub_buffer_len, svr_pub_sig_buffer, KG_SIG_LEN, &verified)) {
        KG_LOG_DBG("TZ API failed when verifying svr dh public key signature\n");
        ret = KG_TZ_API_FAIL;
        goto exit;
    }
    if (false == verified) {
        KG_LOG_DBG("Failed to verify svr dh public key signature\n");
        ret = KG_DH_VERIFY_SIG_FAIL;
        goto exit;
    }
    KG_LOG("Svr dh public signature has been verified successfully\n");

    if (KG_SUCCESS != KG_dh_gen_shared_key(svr_pub_buffer, KG_DH_PUB_KEY_LEN, dh_secret, &dh_secret_len, &(secure_data->kg_dh_data))) {
        KG_LOG_DBG("Failed to generate shared key from svr pub\n");
        ret = KG_DH_GEN_SECRET_FAIL;
        goto exit;
    }

    KG_DUMP_DBG("Generated shared dh secret: \n", dh_secret, dh_secret_len);

    if (TZ_API_OK != TZ_digest_SHA256(dh_secret, dh_secret_len, hotp_secret_key, &hotp_secret_key_len)) {
        KG_LOG("Failed to generate hotp secret key\n");
        ret = KG_DH_GEN_HOTP_FAIL;
        goto exit;
    }

    if (hotp_secret_key_len != KG_KEY_LEN) {
        KG_LOG("Generated hotp secret key length check failed\n");
        ret = KG_DH_GEN_HOTP_FAIL;
        goto exit;
    }

    KG_DUMP_DBG("hotp secret key: \n", hotp_secret_key, hotp_secret_key_len);

    if (KG_SUCCESS != hotp_challenge(hotp_secret_key, hotp_secret_key_len, seed, 1, cha_check, &cha_check_len)) {
        KG_LOG("Faile to generate hotp challenge to verify server response\n");
        ret = KG_HOTP_CHALL_GEN_FAIL;
        goto exit;
    }

    if (cha_check_len != KG_OTP_LEN) {
        KG_LOG("Generated hotp challenge check failed\n");
        ret = KG_HOTP_CHALL_GEN_FAIL;
        goto exit;
    }

    KG_DUMP_DBG("hotp challenge: \n", cha_check, cha_check_len);

    if (KG_OTP_LEN != sendmsg->payload.cmd.dh_challenge_len) {
        KG_LOG("Hotp provisioning recived invalid challenge length from server\n");
        ret = KG_HOTP_CHALL_CHECK_FAIL;
        goto exit;
    }

    uint32_t idx = 0;
    bool check_flag = true;
    while (idx < KG_OTP_LEN) {
        if (cha_check[idx] != sendmsg->payload.cmd.dh_challenge[idx]) {
            check_flag = false;
        }
        idx++;
    }

    if (false == check_flag) {
        KG_LOG("Hotp provisioning on 1st iteration mismatch\n");
        ret = KG_HOTP_CHALL_CHECK_FAIL;
        goto exit;
    }

    secure_data->kg_metadata.hotp_iteration++;
    TEE_MemMove(secure_data->kg_metadata.hotp_key, (void *)hotp_secret_key, KG_KEY_LEN);

    rewrap_data = TEE_Malloc(rewrap_data_len, 0);
    if (NULL == rewrap_data) {
        KG_LOG("KG TA failed to alloc buffer to hold wrap data\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }
#ifdef CONFIG_QSEE
    if (TZ_API_OK != TZ_wrap_persist_data((uint8_t *)KG_NAME, strlen(KG_NAME), 
        unwrap_data, sizeof(kg_secure_data_t), rewrap_data, &rewrap_data_len)) {
        KG_LOG("Failed to wrap kg secure data structure in QC\n");
        ret = KG_TZ_API_FAIL;
        goto exit;
    }
#else
    if (TZ_API_OK != TZ_wrap_data_with_derived_key((uint8_t *)KG_NAME, strlen(KG_NAME), 
        unwrap_data, sizeof(kg_secure_data_t), rewrap_data, &rewrap_data_len)) {
        KG_LOG("Failed to wrap kg secure data structure in QC\n");
        ret = KG_TZ_API_FAIL;
        goto exit;
    }
#endif    
    if (rewrap_data_len > KG_SECURE_DATA_LEN) {
        KG_LOG("Wraped kg secure data size overflow\n");
        ret = KG_RPMB_WRAP_FAIL;
        goto exit;
    }
    info->kg_wrap_data_len = rewrap_data_len;
    if (KG_SUCCESS != (ret = write_wrap_data(info->kg_wrap_data_len, rewrap_data))) {
        KG_LOG("failed to write wrap data\n");
        goto exit;
    }

    info->kg_state = KG_STATE_ACTIVE;
    if (KG_SUCCESS != (ret = write_info_object(info))) {
        KG_LOG("failed to write info object\n");
        goto exit;
    }

exit:
    if (svr_pub_buffer != NULL) {
        TEE_Free(svr_pub_buffer);
        svr_pub_buffer = NULL;
    }
    if (svr_pub_sig_buffer != NULL) {
        TEE_Free(svr_pub_sig_buffer);
        svr_pub_sig_buffer = NULL;
    }
    if (pub_mod != NULL) {
        TEE_Free(pub_mod);
        pub_mod = NULL;
    }
    if (pub_exp != NULL) {
        TEE_Free(pub_exp);
        pub_exp = NULL;
    }
    if (dh_secret != NULL) {
        TEE_MemFill(dh_secret, 0x0, dh_secret_len);
        TEE_Free(dh_secret);
        dh_secret = NULL;
    }
    if (hotp_secret_key != NULL) {
        TEE_MemFill(hotp_secret_key, 0x0, hotp_secret_key_len);
        TEE_Free(hotp_secret_key);
        hotp_secret_key = NULL;
    }
    if (cha_check != NULL) {
        TEE_Free(cha_check);
        cha_check = NULL;
    }
    if (NULL != info) {
        TEE_Free(info);
        info = NULL;
    }
    if (NULL != wrap_data) {
        TEE_Free(wrap_data);
        wrap_data = NULL;
    }
    if (NULL != unwrap_data) {
        TEE_MemFill(unwrap_data, 0, unwrap_data_len);
        TEE_Free(unwrap_data);
        unwrap_data = NULL;
    }
    if (NULL != rewrap_data) {
        TEE_Free(rewrap_data);
        rewrap_data = NULL;
    }
    return ret;
}