#include "kg_hotp_cmd.h"

uint32_t KG_generate_hotp_challenge(tz_common_payload_t *sendmsg, tz_common_payload_t *respmsg) {
    KG_LOG("KG generate hotp challenge\n");
    uint32_t ret = KG_SUCCESS;
    uint8_t *hotp_c = NULL;
    uint32_t hotp_c_len = KG_HOTP_LEN;
    uint8_t *hotp_s = NULL;
    uint32_t hotp_s_len = KG_HOTP_LEN;

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

    if (sendmsg->payload.cmd.data_len > KG_MAX_PAYLOAD_LEN) {
        KG_LOG("Received invalid input buffer when generating hotp challenge\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }
    
    hotp_c = TEE_Malloc(hotp_c_len, 0);
    if (NULL == hotp_c) {
        KG_LOG("Failed to alloc buffer for hotp challenge\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

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

#ifdef CONFIG_QSEE
    if (KG_SUCCESS != kg_rpmb_init()) {
        KG_LOG("RPMB is not available when initing KG secure data\n");
        ret = KG_RPMB_UNAVAILABLE;
        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;

    // TODO, we should do magic check for every read from RPMB
    KG_LOG_DBG("Current hotp iteration is %u\n", secure_data->kg_metadata.hotp_iteration);
    char *counterBytes = decimal_to_binary(secure_data->kg_metadata.hotp_iteration);
    if (NULL == counterBytes) {
        KG_LOG("Failed to convert hotp iteration to counterBytes\n");
        ret = KG_HOTP_CHALL_GEN_FAIL;
        goto exit;
    }
    uint32_t start_idx = 0;
    while (start_idx < strlen(counterBytes)) {
        if (counterBytes[start_idx] != 0x30) {
            break;
        }
        start_idx++;
    }
    KG_LOG_DBG("Converted hotp iteration string is %s\n", counterBytes + start_idx);
    KG_DUMP_DBG("hotp iteration byte array is \n", (uint8_t *)(counterBytes + start_idx), strlen(counterBytes) - start_idx);

    if (KG_SUCCESS != hotp_challenge(secure_data->kg_metadata.hotp_key, KG_KEY_LEN, 
        (uint8_t *)(counterBytes + start_idx), strlen(counterBytes) - start_idx, hotp_c, &hotp_c_len)) {
        KG_LOG_DBG("Failed to generated hotp challenge with given parameters\n");
        ret = KG_HOTP_CHALL_GEN_FAIL;
        goto exit;
    }

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

    bool flag_check = true;
    uint32_t check_idx = 0;
    while (check_idx < KG_OTP_LEN) {
        if (hotp_c[check_idx] != secure_data->kg_metadata.hotp_cha[check_idx]) {
            flag_check = false;
        }
        check_idx++;
    }

    if (true == flag_check) {
        if(hotp_c_len > KG_MAX_PAYLOAD_LEN){
            ret = KG_BUFFER_SIZE_FAIL;
            goto exit;
        }
        respmsg->payload.resp.data_len = hotp_c_len;
        TEE_MemMove(respmsg->payload.resp.data_buf, hotp_c, hotp_c_len);
        goto exit;
    }

    if (KG_SUCCESS != hotp_passcode(secure_data->kg_metadata.hotp_key, KG_KEY_LEN,
        (uint8_t *)(counterBytes + start_idx), strlen(counterBytes) - start_idx, hotp_s, &hotp_s_len)) {
        KG_LOG_DBG("Failed to generated hotp secret with given parameters\n");
        ret = KG_HOTP_CHALL_GEN_FAIL;
        goto exit;
    }

    if (hotp_s_len != KG_OTP_LEN) {
        KG_LOG_DBG("Generated hotp secret length check failed\n");
        ret = KG_HOTP_CHALL_GEN_FAIL;
        goto exit;
    }

    TEE_MemMove(secure_data->kg_metadata.hotp_cha, hotp_c, KG_OTP_LEN);
    TEE_MemMove(secure_data->kg_metadata.hotp_secret, hotp_s, KG_OTP_LEN);
    KG_DUMP_DBG("KG rpmb data dump: \n", (uint8_t *)&(secure_data->kg_metadata), sizeof(kg_metadata_t));

    // save kg_wrap_data
    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 data\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), 
        (uint8_t *)secure_data, sizeof(kg_secure_data_t), rewrap_data, &rewrap_data_len)) {
        KG_LOG("Failed to wrap kg secure data structure\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(hotp_c_len > KG_MAX_PAYLOAD_LEN){
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }
    respmsg->payload.resp.data_len = hotp_c_len;
    TEE_MemMove(respmsg->payload.resp.data_buf, hotp_c, hotp_c_len);

exit:
    if (hotp_c != NULL) {
        TEE_Free(hotp_c);
        hotp_c = NULL;
    }
    if (hotp_s != NULL) {
        TEE_Free(hotp_s);
        hotp_s = NULL;
    }
    if (info != NULL) {
        TEE_Free(info);
        info = NULL;
    }
    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;
    }
    return ret;
}

uint32_t KG_verify_hotp_secret(tz_common_payload_t *sendmsg, tz_common_payload_t *respmsg) {
    KG_LOG("KG verify hotp secret\n");
    uint32_t ret = KG_SUCCESS;
    uint32_t retry = 0;

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


    if (sendmsg->payload.cmd.data_len > KG_MAX_PAYLOAD_LEN) {
        KG_LOG("Received invalid input buffer when verifying hotp secret\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

#ifdef CONFIG_QSEE
    if (KG_SUCCESS != kg_rpmb_init()) {
        KG_LOG("RPMB is not available when initing KG secure data\n");
        ret = KG_RPMB_UNAVAILABLE;
        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;
    
    KG_DUMP_DBG("rpmb hotp secret: \n", secure_data->kg_metadata.hotp_secret, KG_OTP_LEN);

    KG_LOG_DBG("Received secret length is %d\n", sendmsg->payload.cmd.data_len);
    KG_LOG_DBG("Received secret is %s\n",  sendmsg->payload.cmd.data_buf);

    if (sendmsg->payload.cmd.data_len != KG_OTP_LEN) {
        KG_LOG_DBG("KG input buffer check for hotp verification failed\n");
		secure_data->kg_metadata.hotp_retry_counter++;
        retry = secure_data->kg_metadata.hotp_retry_counter;
        ret = KG_HOTP_SECRET_CHECK_FAIL;
		
        goto kg_rpmb_write_metadata;
    }

    uint32_t i = 0;
    bool check_flag = true;
    while (i < KG_OTP_LEN) {
        if (secure_data->kg_metadata.hotp_secret[i] != sendmsg->payload.cmd.data_buf[i]) {
            check_flag = false;
        }
        i++;
    }

    if (false == check_flag) {
        KG_LOG("HOTP verification failed\n");
		secure_data->kg_metadata.hotp_retry_counter++;
        retry = secure_data->kg_metadata.hotp_retry_counter;
        ret = KG_HOTP_SECRET_CHECK_FAIL;

        goto kg_rpmb_write_metadata;
    }

    retry = secure_data->kg_metadata.hotp_retry_counter = 0;
    secure_data->kg_metadata.hotp_iteration++;

    info->kg_state = KG_STATE_ACTIVE;

kg_rpmb_write_metadata:
    // save kg_wrap_data
    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\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;
    }
exit:
    if (info != NULL) {
        TEE_Free(info);
        info = NULL;
    }
    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;
    }
	KG_LOG("retry=%d\n",retry);	
    respmsg->payload.resp.result = retry;
	
    return ret;
}