#include "crypto_module.h"

#include <string.h>
#include <tee_internal_api.h>
#include <tees_kdf.h>

#ifndef NULL
#define NULL 0
#endif

#if AP_ID_SIZE * 2 != SHA256_DIGEST_SIZE
#   error Bad apId derivation algorithm
#endif

CRYPTO_STATUS  crypto_gen_random( uint8_t *out, uint32_t size ) {
    uint32_t desired;

    while (size != 0) {
        desired = size;
        TEE_GenerateRandom( out, desired );

        out += desired;
        size -= desired;
    }

    return CRYPTO_STATUS_SUCCESS;
}


CRYPTO_STATUS  crypto_get_tz_encryption_key( uint8_t out[AES_256_KEY_SIZE] )
{
    TEE_Result res;
    TEE_ObjectHandle obj_hndl = TEE_HANDLE_NULL;
    uint8_t key_label[] = "OTP encryption key";
    uint8_t key_context[] = "SSPM TZ KDF";
    uint32_t byte_size = AES_256_KEY_SIZE;
    const uint32_t bit_size = byte_size * 8;

    // TODO check may be TEE_TYPE_AES type is needed below
    res = TEE_AllocateTransientObject( TEE_TYPE_GENERIC_SECRET, bit_size, &obj_hndl );
    if ( TEE_SUCCESS != res )
    {
        LOGE( "TEE_AllocateTransientObject failed with code %d", res );
        return CRYPTO_STATUS_FAILED;
    }

    res = TEES_DeriveKeyKDF( key_label, sizeof(key_label) - 1,
                             key_context, sizeof(key_context) - 1,
                             byte_size, obj_hndl );
    if( TEE_SUCCESS != res )
    {
        LOGE( "TEES_DeriveKeyKDF failed with code %d", res );
        return CRYPTO_STATUS_FAILED;
    }

    // obj_hndl contains a key => need to copy to the output buffer
    // TODO (may be another constant than TEE_ATTR_SECRET_VALUE)
    res = TEE_GetObjectBufferAttribute( obj_hndl, TEE_ATTR_SECRET_VALUE, out, &byte_size );
    if( TEE_SUCCESS != res )
    {
        LOGE( "TEE_GetObjectBufferAttribute failed with code %d", res );
        return CRYPTO_STATUS_FAILED;
    }

    TEE_CloseObject( obj_hndl );
    if( TEE_SUCCESS != res )
    {
        LOGE( "TEE_CloseObject failed with code %d", res );
        return CRYPTO_STATUS_FAILED;
    }

    return CRYPTO_STATUS_SUCCESS;
}

CRYPTO_STATUS  crypto_get_apId( uint8_t apId[AP_ID_SIZE] )
{
    CRYPTO_STATUS   status;
    TEE_Result     res;
    uint32_t    i;
    uint8_t     hashed[SHA256_DIGEST_SIZE];
    TEE_UUID    suid;

    res = TEE_GetPropertyAsUUID(TEE_PROPSET_CURRENT_TA,"gpd.ta.appID", &suid );
    if( TEE_SUCCESS != res )
    {
        LOGE( "TEE_GetPropertyAsUUID failed with code %d", res );
        return CRYPTO_STATUS_FAILED;
    }

    status = crypto_sha256( hashed, &suid, sizeof(suid), NULL );
    if( CRYPTO_STATUS_SUCCESS != status )
    {
        LOGE( "crypto_sha256 failed (error 0x%08X)", status );
        return status;
    }

    for ( i = 0; i < AP_ID_SIZE; ++i )
        apId[i] = hashed[i] ^ hashed[AP_ID_SIZE + i];

    return CRYPTO_STATUS_SUCCESS;
}