#include "kg_lock_cmd.h"

uint32_t KG_lock(tz_lock_payload_t *sendmsg, tz_lock_payload_t *respmsg) {
    KG_LOG("KG lock\n");
    uint32_t ret = KG_SUCCESS;
    kg_rpmb_info_t* info = NULL;
    uint8_t* lock_object = NULL;
    uint32_t lock_object_len = KG_LOCK_OBJECT_MAX;

    if (sendmsg->payload.cmd.action_name_len > KG_ACTION_NAME_MAX ||
        sendmsg->payload.cmd.lock_object_len > KG_LOCK_OBJECT_MAX) {
        KG_LOG("Received lock screen action name or lock object size overflow\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

#ifdef CONFIG_QSEE
    if (KG_SUCCESS != kg_rpmb_init()) {
        KG_LOG("RPMB is not available when locking\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;
    }

    info->kg_state = KG_STATE_LOCK;

    TEE_MemFill(info->action_name, KG_ACTION_NAME_MAX, 0);
    if (sendmsg->payload.cmd.action_name_len > KG_ACTION_NAME_MAX) {
        sendmsg->payload.cmd.action_name_len = KG_ACTION_NAME_MAX;
    }

    info->action_name_len = sendmsg->payload.cmd.action_name_len;
    TEE_MemMove(info->action_name, sendmsg->payload.cmd.action_name_buf, sendmsg->payload.cmd.action_name_len);

    lock_object = TEE_Malloc(lock_object_len, 0);
    if (NULL == lock_object) {
        KG_LOG("failed to alloc lock_object\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }
    
    if (sendmsg->payload.cmd.lock_object_len > KG_LOCK_OBJECT_MAX) {
        sendmsg->payload.cmd.lock_object_len = KG_LOCK_OBJECT_MAX;
    }

    info->lock_object_len = sendmsg->payload.cmd.lock_object_len;
    TEE_MemMove(lock_object, sendmsg->payload.cmd.lock_object_buf, sendmsg->payload.cmd.lock_object_len);

    if (KG_SUCCESS != (ret = write_lock_object(lock_object_len, lock_object))) {
        KG_LOG("failed to write lock object\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 (lock_object != NULL) {
        TEE_Free(lock_object);
        lock_object = NULL;
    }
    return ret;
}

uint32_t KG_unlock(tz_common_payload_t *sendmsg, tz_common_payload_t *respmsg) {
    KG_LOG("KG unlock\n");
    uint32_t ret = KG_SUCCESS;
    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;

#ifdef CONFIG_QSEE
    if (KG_SUCCESS != kg_rpmb_init()) {
        KG_LOG("RPMB is not available when unlocking\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;
    secure_data->kg_metadata.hotp_retry_counter = 0;
    secure_data->kg_metadata.hotp_iteration++;

    info->kg_state = KG_STATE_ACTIVE;

    // 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 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;
    }
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;
    }
    return ret;
}


uint32_t KG_get_lock_action(tz_common_payload_t *sendmsg, tz_common_payload_t *respmsg) {
    uint32_t ret = KG_SUCCESS;
    kg_rpmb_info_t* info = NULL;

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

#ifdef CONFIG_QSEE
    if (KG_SUCCESS != kg_rpmb_init()) {
        KG_LOG("RPMB is not available when getting lock action\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->action_name_len > KG_MAX_PAYLOAD_LEN){
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }
    respmsg->payload.resp.data_len = info->action_name_len;
    TEE_MemMove(respmsg->payload.resp.data_buf, info->action_name, info->action_name_len);

exit:
    if (info != NULL) {
        TEE_Free(info);
        info = NULL;
    }
    return ret;
}

uint32_t KG_get_lock_object(tz_common_payload_t *sendmsg, tz_common_payload_t *respmsg) {
    uint32_t ret = KG_SUCCESS;
    uint8_t* lock_object = NULL;
    uint32_t lock_object_len = 0;

#ifdef CONFIG_QSEE
    if (KG_SUCCESS != kg_rpmb_init()) {
        KG_LOG("RPMB is not available when getting lock object\n");
        ret = KG_RPMB_UNAVAILABLE;
        goto exit;
    }
#endif

    if (KG_SUCCESS != (ret = read_lock_object(&lock_object_len, &lock_object))
        || NULL == lock_object) {
        KG_LOG("failed to read lock object\n");
        goto exit;
    }
    respmsg->payload.resp.data_len = lock_object_len;
    if (lock_object_len > KG_MAX_PAYLOAD_LEN) {
        KG_LOG_DBG("Buffer length check failed\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }
    TEE_MemMove(respmsg->payload.resp.data_buf, lock_object, lock_object_len);
exit:
    if (NULL != lock_object) {
        TEE_Free(lock_object);
        lock_object = NULL;
    }
    return ret;
}

uint32_t KG_set_client_data(tz_common_payload_t *sendmsg, tz_common_payload_t *respmsg) {
    uint32_t ret = KG_SUCCESS;
    kg_rpmb_info_t* info = NULL;
    uint8_t* write_buffer = NULL;
    uint32_t write_buffer_len = KG_CLIENT_DATA_MAX;

    if (sendmsg->payload.cmd.data_len >= KG_CLIENT_DATA_MAX) {
        KG_LOG("Received client data buffer size overflow\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

#ifdef CONFIG_QSEE
    if (KG_SUCCESS != kg_rpmb_init()) {
        KG_LOG("RPMB is not available when setting client 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;
    }

    write_buffer = TEE_Malloc(write_buffer_len, 0);
    if (NULL == write_buffer) {
        KG_LOG("failed to alloc client write buffer\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }
    TEE_MemMove(write_buffer, sendmsg->payload.cmd.data_buf, sendmsg->payload.cmd.data_len);

    if (KG_SUCCESS != (ret = write_client_data(write_buffer_len, write_buffer))) {
        KG_LOG("failed to write client data\n");
        goto exit;
    }

    // copy max size client data which input data with zero padding
    info->client_data_len = sendmsg->payload.cmd.data_len;
    if (KG_SUCCESS != (ret = write_info_object(info))) {
        KG_LOG("failed to write info object\n");
        goto exit;
    }

exit:
    if (NULL != info) {
        TEE_Free(info);
        info = NULL;
    }
    if (NULL != write_buffer) {
        TEE_Free(write_buffer);
        write_buffer = NULL;
    }
    return ret;
}

uint32_t KG_get_client_data(tz_common_payload_t *sendmsg, tz_common_payload_t *respmsg) {
    uint32_t ret = KG_SUCCESS;
    uint8_t* client_data = NULL;
    uint32_t client_data_len = 0;

#ifdef CONFIG_QSEE
    if (KG_SUCCESS != kg_rpmb_init()) {
        KG_LOG("RPMB is not available when getting client data\n");
        ret = KG_RPMB_UNAVAILABLE;
        goto exit;
    }
#endif

    if (KG_SUCCESS != (ret = read_client_data(&client_data_len, &client_data))
        || NULL == client_data) {
        KG_LOG("failed to read client data\n");
        goto exit;
    }

    if(client_data_len > KG_MAX_PAYLOAD_LEN){
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }
    respmsg->payload.resp.data_len = client_data_len;
    TEE_MemMove(respmsg->payload.resp.data_buf, client_data, client_data_len);

exit:
    if (NULL != client_data) {
        TEE_Free(client_data);
        client_data = NULL;
    }
    return ret;
}