/* BEGIN_HEADER */
#include <stdint.h>
#include "psa/crypto.h"
#include "psa_crypto_storage.h"
#include "mbedtls/md.h"

#define PSA_KEY_STORAGE_MAGIC_HEADER "PSA\0KEY"
#define PSA_KEY_STORAGE_MAGIC_HEADER_LENGTH ( sizeof( PSA_KEY_STORAGE_MAGIC_HEADER ) )

typedef struct {
    uint8_t magic[PSA_KEY_STORAGE_MAGIC_HEADER_LENGTH];
    uint8_t version[4];
    uint8_t type[sizeof( psa_key_type_t )];
    uint8_t policy[sizeof( psa_key_policy_t )];
    uint8_t data_len[4];
    uint8_t key_data[];
} psa_persistent_key_storage_format;

/* END_HEADER */

/* BEGIN_DEPENDENCIES
 * depends_on:MBEDTLS_PSA_CRYPTO_C:MBEDTLS_PSA_CRYPTO_STORAGE_C
 * END_DEPENDENCIES
 */

/* BEGIN_CASE */
void format_storage_data_check( data_t *key_data,
                                data_t *expected_file_data,
                                int key_type, int key_usage, int key_alg )
{
    uint8_t *file_data;
    size_t file_data_length;
    psa_key_policy_t key_policy;

    key_policy.usage = (psa_key_usage_t) key_usage;
    key_policy.alg = (psa_algorithm_t) key_alg;

    file_data_length = key_data->len + sizeof( psa_persistent_key_storage_format );
    file_data = mbedtls_calloc( 1, file_data_length );
    psa_format_key_data_for_storage( key_data->x, key_data->len,
                                     (psa_key_type_t) key_type, &key_policy,
                                     file_data );

    ASSERT_COMPARE( expected_file_data->x, expected_file_data->len,
                    file_data, file_data_length );
    mbedtls_free( file_data );
}
/* END_CASE */

/* BEGIN_CASE */
void parse_storage_data_check( data_t *file_data,
                               data_t *expected_key_data,
                               int expected_key_type,
                               int expected_key_usage,
                               int expected_key_alg,
                               int expected_status )
{
    uint8_t *key_data = NULL;
    size_t key_data_length = 0;
    psa_key_type_t key_type = 0;
    psa_key_policy_t key_policy;
    psa_status_t status;

    status = psa_parse_key_data_from_storage( file_data->x, file_data->len,
                                              &key_data, &key_data_length,
                                              &key_type, &key_policy );

    TEST_EQUAL( status, expected_status );
    if( status != PSA_SUCCESS )
        goto exit;

    TEST_EQUAL( key_type, (psa_key_type_t) expected_key_type );
    TEST_EQUAL( key_policy.usage, (uint32_t) expected_key_usage );
    TEST_EQUAL( key_policy.alg, (uint32_t) expected_key_alg );
    ASSERT_COMPARE( expected_key_data->x, expected_key_data->len,
                    key_data, key_data_length );

exit:
    mbedtls_free( key_data );
}
/* END_CASE */

/* BEGIN_CASE */
void save_large_persistent_key( int data_too_large, int expected_status )
{
    psa_key_id_t key_id = 42;
    psa_key_handle_t handle = 0;
    uint8_t *data = NULL;
    size_t data_length = PSA_CRYPTO_MAX_STORAGE_SIZE;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;

    if( data_too_large )
        data_length += 1;

    ASSERT_ALLOC( data, data_length );

    PSA_ASSERT( psa_crypto_init() );

    psa_make_key_persistent( &attributes, key_id, PSA_KEY_LIFETIME_PERSISTENT );
    psa_set_key_type( &attributes, PSA_KEY_TYPE_RAW_DATA );

    TEST_EQUAL( psa_import_key( &attributes, &handle,
                                data, data_length ),
                expected_status );

exit:
    mbedtls_free( data );
    mbedtls_psa_crypto_free();
    psa_destroy_persistent_key( key_id );
}
/* END_CASE */

/* BEGIN_CASE */
void persistent_key_destroy( int key_id_arg, int restart,
                             int first_type_arg, data_t *first_data,
                             int second_type_arg, data_t *second_data )
{
    psa_key_id_t key_id = key_id_arg;
    psa_key_handle_t handle = 0;
    psa_key_type_t first_type = (psa_key_type_t) first_type_arg;
    psa_key_type_t second_type = (psa_key_type_t) second_type_arg;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;

    PSA_ASSERT( psa_crypto_init() );

    psa_make_key_persistent( &attributes, key_id, PSA_KEY_LIFETIME_PERSISTENT );
    psa_set_key_type( &attributes, first_type );

    PSA_ASSERT( psa_import_key( &attributes, &handle,
                                first_data->x, first_data->len ) );

    if( restart )
    {
        psa_close_key( handle );
        mbedtls_psa_crypto_free();
        PSA_ASSERT( psa_crypto_init() );
        PSA_ASSERT( psa_open_key( PSA_KEY_LIFETIME_PERSISTENT, key_id,
                                  &handle ) );
    }
    TEST_EQUAL( psa_is_key_present_in_storage( key_id ), 1 );

    /* Destroy the key */
    PSA_ASSERT( psa_destroy_key( handle ) );

    /* Check key slot storage is removed */
    TEST_EQUAL( psa_is_key_present_in_storage( key_id ), 0 );
    TEST_EQUAL( psa_open_key( PSA_KEY_LIFETIME_PERSISTENT, key_id, &handle ),
                PSA_ERROR_DOES_NOT_EXIST );
    TEST_EQUAL( handle, 0 );

    /* Shutdown and restart */
    mbedtls_psa_crypto_free();
    PSA_ASSERT( psa_crypto_init() );

    /* Create another key in the same slot */
    psa_make_key_persistent( &attributes, key_id, PSA_KEY_LIFETIME_PERSISTENT );
    psa_set_key_type( &attributes, second_type );
    PSA_ASSERT( psa_import_key( &attributes, &handle,
                                second_data->x, second_data->len ) );

exit:
    mbedtls_psa_crypto_free();
    psa_destroy_persistent_key( key_id );
}
/* END_CASE */

/* BEGIN_CASE */
void persistent_key_import( int key_id_arg, int type_arg, data_t *data,
                            int restart, int expected_status )
{
    psa_key_id_t key_id = (psa_key_id_t) key_id_arg;
    psa_key_type_t type = (psa_key_type_t) type_arg;
    psa_key_handle_t handle = 0;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;

    PSA_ASSERT( psa_crypto_init() );

    psa_make_key_persistent( &attributes, key_id, PSA_KEY_LIFETIME_PERSISTENT );
    psa_set_key_type( &attributes, type );
    TEST_EQUAL( psa_import_key( &attributes, &handle, data->x, data->len ),
                expected_status );

    if( expected_status != PSA_SUCCESS )
    {
        TEST_EQUAL( psa_is_key_present_in_storage( key_id ), 0 );
        goto exit;
    }

    if( restart )
    {
        psa_close_key( handle );
        mbedtls_psa_crypto_free();
        PSA_ASSERT( psa_crypto_init() );
        PSA_ASSERT( psa_open_key( PSA_KEY_LIFETIME_PERSISTENT, key_id,
                                  &handle ) );
    }

    psa_reset_key_attributes( &attributes );
    PSA_ASSERT( psa_get_key_attributes( handle, &attributes ) );
    TEST_EQUAL( psa_get_key_id( &attributes ), key_id );
    TEST_EQUAL( psa_get_key_lifetime( &attributes ),
                PSA_KEY_LIFETIME_PERSISTENT );
    TEST_EQUAL( psa_get_key_type( &attributes ), type );
    TEST_EQUAL( psa_get_key_usage_flags( &attributes ), 0 );
    TEST_EQUAL( psa_get_key_algorithm( &attributes ), 0 );

exit:
    psa_reset_key_attributes( &attributes );
    psa_destroy_persistent_key( key_id );
    mbedtls_psa_crypto_free();
}
/* END_CASE */

/* BEGIN_CASE */
void import_export_persistent_key( data_t *data, int type_arg,
                                   int expected_bits,
                                   int restart, int key_not_exist )
{
    psa_key_id_t key_id = 42;
    psa_key_type_t type = (psa_key_type_t) type_arg;
    psa_key_handle_t handle = 0;
    unsigned char *exported = NULL;
    size_t export_size = data->len;
    size_t exported_length;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;

    ASSERT_ALLOC( exported, export_size );

    PSA_ASSERT( psa_crypto_init( ) );

    psa_make_key_persistent( &attributes, key_id, PSA_KEY_LIFETIME_PERSISTENT );
    psa_set_key_type( &attributes, type );
    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_EXPORT );

    /* Import the key */
    PSA_ASSERT( psa_import_key( &attributes, &handle,
                                data->x, data->len ) );


    if( restart )
    {
        psa_close_key( handle );
        mbedtls_psa_crypto_free();
        PSA_ASSERT( psa_crypto_init() );
        PSA_ASSERT( psa_open_key( PSA_KEY_LIFETIME_PERSISTENT, key_id,
                                  &handle ) );
    }

    /* Test the key information */
    psa_reset_key_attributes( &attributes );
    PSA_ASSERT( psa_get_key_attributes( handle, &attributes ) );
    TEST_EQUAL( psa_get_key_id( &attributes ), key_id );
    TEST_EQUAL( psa_get_key_lifetime( &attributes ),
                PSA_KEY_LIFETIME_PERSISTENT );
    TEST_EQUAL( psa_get_key_type( &attributes ), type );
    TEST_EQUAL( psa_get_key_bits( &attributes ), (size_t) expected_bits );
    TEST_EQUAL( psa_get_key_usage_flags( &attributes ), PSA_KEY_USAGE_EXPORT );
    TEST_EQUAL( psa_get_key_algorithm( &attributes ), 0 );

    TEST_EQUAL( psa_is_key_present_in_storage( key_id ), 1 );

    if( key_not_exist )
    {
        psa_destroy_persistent_key( key_id );
    }
    /* Export the key */
    PSA_ASSERT( psa_export_key( handle, exported, export_size,
                                &exported_length ) );

    ASSERT_COMPARE( data->x, data->len, exported, exported_length );

    /* Destroy the key */
    PSA_ASSERT( psa_destroy_key( handle ) );
    TEST_EQUAL( psa_is_key_present_in_storage( key_id ), 0 );

exit:
    psa_reset_key_attributes( &attributes );
    mbedtls_free( exported );
    mbedtls_psa_crypto_free( );
    psa_destroy_persistent_key( key_id );
}
/* END_CASE */