#include <tees_secure_object.h>

#include "knoxai_io_datatypes.h"
#include "knoxai_logger.h"

uint32_t knoxai_unwrap(TEE_UUID target_uuid, uint8_t *blob_ptr, uint32_t blob_len, uint8_t *dest_ptr, uint32_t* dest_len){
    tz_knoxai_return_type   ret = KNOXAI_FAILURE;
    TEE_Result              unwrap_ret;

    if ( blob_ptr == NULL || dest_ptr == NULL || blob_len < SO_LOCAL_HEADERSIZE + 1 || dest_len == NULL) {    
        KNOXAI_DEBUG_LOG("ptr must be NOT NULL");
        ret = KNOXAI_FAILURE;
        goto exit;
    }
#if defined(SEC_SDK30) || defined(SEC_SDK40)
    SO_AccessControlInfoType ac_info;

    TEE_UUID creator_uuid = target_uuid;
    TEE_MemFill(&ac_info, 0, sizeof ac_info);

    TEE_MemMove(&ac_info.ta_id, &creator_uuid, sizeof(TEE_UUID));
    TEE_MemMove(&ac_info.auth_id, TA_SNAP_TA_AUTH_CRYPTOSUITE, strnlen(TA_SNAP_TA_AUTH_CRYPTOSUITE, 10));
    ac_info.access_flags = DELEGATED_TA_ID_AC;
    
    unwrap_ret = TEES_CheckSecureObjectCreator(blob_ptr, blob_len, &ac_info); 
    if (unwrap_ret != TEE_SUCCESS) {
      KNOXAI_DEBUG_LOG(":TEES_CheckSecureObjectCreator failed with ret=0x%08X, exit", unwrap_ret);
      ret = KNOXAI_FAILURE;
      goto exit;
    }
#endif
    unwrap_ret = TEES_UnwrapSecureObject(blob_ptr, blob_len, dest_ptr, dest_len);
    if (unwrap_ret != TEE_SUCCESS) {
        KNOXAI_DEBUG_LOG("knoxai_unwrap unwrap failed with ret=0x%08X, exit", unwrap_ret);
        ret = KNOXAI_FAILURE;
        goto exit;
    }/*
    if (*dest_len < KNOXAI_FAC_KEY_LEN || *dest_len > MAX_WRAPPED_KEY_LEN) {
        KNOXAI_DEBUG_LOG("knoxai_unwrap: Key Bigger than space");
        ret = KNOXAI_FAILURE;
        goto exit;
    }*/
    ret = KNOXAI_SUCCESS;
exit:    
    return ret;
}

uint32_t knoxai_wrap(TEE_UUID target_uuid, uint8_t *blob_ptr, uint32_t blob_len, uint8_t *dest_ptr, uint32_t* dest_len){
    tz_knoxai_return_type  ret = KNOXAI_FAILURE;
    TEE_Result             wrap_ret;

    if ( blob_ptr == NULL || dest_ptr == NULL || blob_len == 0 || blob_len > *dest_len) {
        KNOXAI_DEBUG_LOG("ptr must be NOT NULL");
        ret = KNOXAI_FAILURE;
        goto exit;
    }
#if defined(SEC_SDK30) || defined(SEC_SDK40)
    SO_AccessControlInfoType ac_info;

    TEE_UUID creator_uuid = target_uuid;
    TEE_MemFill(&ac_info, 0, sizeof ac_info);

    TEE_MemMove(&ac_info.ta_id, &creator_uuid, sizeof(TEE_UUID));
    TEE_MemMove(&ac_info.auth_id, TA_SNAP_TA_AUTH_CRYPTOSUITE, strnlen(TA_SNAP_TA_AUTH_CRYPTOSUITE, 10));
    ac_info.access_flags = DELEGATED_TA_ID_AC;
#endif
    if (*dest_len < SO_OUT_BUF_SIZE(blob_len, true)) {
        KNOXAI_DEBUG_LOG("knoxai_wrap failed - dest_len size error, exit");
        ret = KNOXAI_FAILURE;
        goto exit;
    }
#if defined(SEC_SDK30) || defined(SEC_SDK40)
    wrap_ret = TEES_WrapSecureObject(blob_ptr, blob_len, dest_ptr, dest_len, &ac_info);
#else
    wrap_ret = TEES_WrapSecureObject(blob_ptr, blob_len, dest_ptr, dest_len, NULL);
#endif
    if (wrap_ret != TEE_SUCCESS) {
        KNOXAI_DEBUG_LOG("knoxai_wrap wrap failed with ret=0x%08X, exit", wrap_ret);
        ret = wrap_ret;
        goto exit;
    }
    /*if (*dest_len < KNOXAI_FAC_KEY_LEN || *dest_len > MAX_WRAPPED_KEY_LEN) {
        KNOXAI_DEBUG_LOG("knoxai_wrap: Key Bigger than space");
        ret = KNOXAI_FAILURE;
        goto exit;
    }*/
    ret = KNOXAI_SUCCESS;
exit:    
    return ret;
}