/* BEGIN_HEADER */
#include <string.h>

#include <pkcs11.h>

#include "mbedtls/pkcs11_client.h"

#if defined(MBEDTLS_PK_C)

#include "mbedtls/oid.h"
#include "mbedtls/asn1write.h"
#include "mbedtls/bignum.h"
#include "mbedtls/rsa.h"
#include "mbedtls/pk.h"

#define ARRAY_LENGTH( a ) ( sizeof( a ) / sizeof( *( a ) ) )

#define CK_ASSERT( expr )                                               \
    do {                                                                \
        CK_RV CK_ASSERT_rv = ( expr );                                  \
        char CK_ASSERT_msg[sizeof( #expr ) + 20] = #expr;               \
        if( CK_ASSERT_rv != CKR_OK ) {                                  \
            sprintf( CK_ASSERT_msg + strlen( CK_ASSERT_msg ),           \
                     " -> 0x%x", (unsigned) CK_ASSERT_rv );             \
            test_fail( CK_ASSERT_msg, __LINE__, __FILE__ );             \
            goto exit;                                                  \
        }                                                               \
    } while( 0 )

#define RSA_KEY_SIZE_BITS 1024
#define RSA_KEY_SIZE_BYTES ( ( RSA_KEY_SIZE_BITS + 7 ) / 8 )
#define ECP_GROUP_ID MBEDTLS_ECP_DP_SECP256R1
#define ECP_GROUP_NAME( id ) #id

//static CK_BBOOL ck_false = CK_FALSE;
static CK_BBOOL ck_true = CK_TRUE;

static int pkcs11_token_label_is( const CK_TOKEN_INFO *info, const char *label )
{
    size_t n = strlen( label );
    if( n > sizeof( info->label ) )
        return( 0 );
    if( memcmp( info->label, label, n ) )
        return( 0 );
    for( ; n < sizeof( info->label ); n++ )
    {
        if( info->label[n] != ' ' )
            return( 0 );
    }
    return( 1 );
}

static int pkcs11_get_slot_id( const char *label, CK_SLOT_ID *slot )
{
    CK_SLOT_ID *slots = NULL;
    CK_ULONG count;
    CK_ULONG i;
    CK_TOKEN_INFO info;
    int found = 0;
    CK_ASSERT( C_GetSlotList( CK_TRUE, NULL_PTR, &count ) );
    slots = mbedtls_calloc( sizeof( *slots ), count );
    TEST_ASSERT( slots != NULL );
    CK_ASSERT( C_GetSlotList( CK_TRUE, slots, &count ) );
    for( i = 0; i < count; i++ )
    {
        CK_ASSERT( C_GetTokenInfo( slots[i], &info ) );
        if( pkcs11_token_label_is( &info, label ) )
        {
            *slot = slots[i];
            found = 1;
            break;
        }
    }
    if( !found )
        mbedtls_fprintf( stdout, "No token found with label %s\n", label );
exit:
    mbedtls_free( slots );
    return( found );
}

static CK_OBJECT_HANDLE pkcs11_init( void )
{
    CK_SESSION_HANDLE hSession;
    CK_SLOT_ID slot;
    unsigned char user_pin[4] = "0000";
    CK_ASSERT( C_Initialize( NULL_PTR ) );
    TEST_ASSERT( pkcs11_get_slot_id( "scratch", &slot ) );
    CK_ASSERT( C_OpenSession( slot,
                              CKF_RW_SESSION | CKF_SERIAL_SESSION,
                              NULL_PTR, NULL_PTR,
                              &hSession ) );
    CK_ASSERT( C_Login( hSession, CKU_USER, user_pin, sizeof( user_pin ) ) );
    return( hSession );
exit:
    return( CK_INVALID_HANDLE );
}

static CK_RV pkcs11_generate_key( mbedtls_pk_type_t key_type,
                                  CK_SESSION_HANDLE hSession,
                                  CK_OBJECT_HANDLE *phPublicKey,
                                  CK_OBJECT_HANDLE *phPrivateKey )
{
    CK_MECHANISM mechanism = {0, NULL_PTR, 0};
    CK_ATTRIBUTE public_attributes[] = {
        {0, 0, 0},
        {CKA_ENCRYPT, &ck_true, sizeof( ck_true )},
        {CKA_VERIFY, &ck_true, sizeof( ck_true )},
    };
    CK_ATTRIBUTE private_attributes[] = {
        {CKA_DECRYPT, &ck_true, sizeof( ck_true )},
        {CKA_SIGN, &ck_true, sizeof( ck_true )},
    };
    CK_ULONG ck_rsa_key_size = RSA_KEY_SIZE_BITS;

    switch( key_type )
    {
#if defined(MBEDTLS_RSA_C)
    case MBEDTLS_PK_RSA:
        mechanism.mechanism = CKM_RSA_PKCS_KEY_PAIR_GEN;
        public_attributes[0].type = CKA_MODULUS_BITS;
        public_attributes[0].pValue = &ck_rsa_key_size;
        public_attributes[0].ulValueLen = sizeof( ck_rsa_key_size );
        break;
#endif /* MBEDTLS_RSA_C */
    default:
        test_fail( "Unsupported key type in test data", __LINE__, __FILE__ );
        break;
    }

    return( C_GenerateKeyPair( hSession,
                               &mechanism,
                               public_attributes,
                               ARRAY_LENGTH( public_attributes ),
                               private_attributes,
                               ARRAY_LENGTH( private_attributes ),
                               phPublicKey, phPrivateKey ) );
}


#endif /* MBEDTLS_PK_C */

/* END_HEADER */

/* BEGIN_DEPENDENCIES
 * depends_on:MBEDTLS_PKCS11_CLIENT_C
 * END_DEPENDENCIES
 */

/* BEGIN_CASE depends_on:MBEDTLS_PK_C:MBEDTLS_SHA256_C */
void pk_generate_sign( int key_type )
{
    mbedtls_pk_context pkcs11_ctx;
    mbedtls_pk_context transparent_ctx;
    CK_SESSION_HANDLE hSession = CK_INVALID_HANDLE;
    CK_OBJECT_HANDLE hPublicKey = CK_INVALID_HANDLE;
    CK_OBJECT_HANDLE hPrivateKey = CK_INVALID_HANDLE;
    unsigned char hash_value[32] = "Fake hash, it doesn't matter....";
    unsigned char sig_buffer[RSA_KEY_SIZE_BYTES];
    size_t sig_length = sizeof( sig_buffer );

    mbedtls_pk_init( &pkcs11_ctx );
    mbedtls_pk_init( &transparent_ctx );

    /* Initialize cryptoki and generate a key in the token */
    hSession = pkcs11_init( );
    TEST_ASSERT( hSession != CK_INVALID_HANDLE );

    CK_ASSERT( pkcs11_generate_key( key_type,
                                    hSession,
                                    &hPublicKey, &hPrivateKey ) );
    TEST_ASSERT( hPublicKey != CK_INVALID_HANDLE );
    TEST_ASSERT( hPrivateKey != CK_INVALID_HANDLE );

    /* Prepare the mbed TLS contexts */
    TEST_ASSERT( mbedtls_pk_setup( &transparent_ctx,
                                   mbedtls_pk_info_from_type( key_type ) ) == 0 );
    TEST_ASSERT( mbedtls_pk_setup_pkcs11( &pkcs11_ctx,
                                          hSession,
                                          hPublicKey,
                                          hPrivateKey ) == 0 );

    /* Retrieve the public key from the token */
    switch( key_type )
    {
#if defined(MBEDTLS_RSA_C)
    case MBEDTLS_PK_RSA:
        {
            unsigned char n_buffer[RSA_KEY_SIZE_BYTES];
            unsigned char e_buffer[RSA_KEY_SIZE_BYTES];
            CK_ATTRIBUTE public_attributes[] = {
                {CKA_MODULUS, n_buffer, sizeof( n_buffer )},
                {CKA_PUBLIC_EXPONENT, e_buffer, sizeof( e_buffer )},
            };
            CK_ULONG *n_length = &public_attributes[0].ulValueLen;
            CK_ULONG *e_length = &public_attributes[1].ulValueLen;
            mbedtls_rsa_context *rsa_ctx = mbedtls_pk_rsa( transparent_ctx );

            CK_ASSERT( C_GetAttributeValue( hSession, hPublicKey,
                                            public_attributes, ARRAY_LENGTH( public_attributes ) ) );
            TEST_ASSERT( mbedtls_mpi_read_binary( &rsa_ctx->N,
                                                  n_buffer, *n_length ) == 0 );
            TEST_ASSERT( mbedtls_mpi_read_binary( &rsa_ctx->E,
                                                  e_buffer, *e_length ) == 0 );
            rsa_ctx->len = mbedtls_mpi_size( &rsa_ctx->N );
        }
        break;
#endif /* MBEDTLS_RSA_C */

    default:
        TEST_ASSERT( !"Unsupported key type in test data" );
        break;
    }

    /* Sign with the token and verify in software */
    TEST_ASSERT( mbedtls_pk_sign( &pkcs11_ctx, MBEDTLS_MD_SHA256,
                                  hash_value, 32,
                                  sig_buffer, &sig_length,
                                  NULL, NULL ) == 0 );
    TEST_ASSERT( mbedtls_pk_verify( &transparent_ctx, MBEDTLS_MD_SHA256,
                                    hash_value, 32,
                                    sig_buffer, sig_length ) == 0 );

exit:
    if( hPublicKey != CK_INVALID_HANDLE )
        C_DestroyObject( hSession, hPublicKey );
    if( hPrivateKey != CK_INVALID_HANDLE )
        C_DestroyObject( hSession, hPrivateKey );
    C_CloseSession( hSession );
    C_Finalize( NULL_PTR );
    mbedtls_pk_free( &pkcs11_ctx );
    mbedtls_pk_free( &transparent_ctx );
}
/* END_CASE */

/* BEGIN_CASE depends_on:MBEDTLS_PK_C:MBEDTLS_SHA256_C */
void pk_import_sign( char *file )
{
    mbedtls_pk_context pkcs11_ctx;
    mbedtls_pk_context transparent_ctx;
    CK_SESSION_HANDLE hSession = CK_INVALID_HANDLE;
    CK_OBJECT_HANDLE hPublicKey = CK_INVALID_HANDLE;
    CK_OBJECT_HANDLE hPrivateKey = CK_INVALID_HANDLE;
    unsigned char hash_value[32] = "Fake hash, it doesn't matter....";
    unsigned char sig_buffer[4096];
    size_t sig_length = sizeof( sig_buffer );

    mbedtls_pk_init( &pkcs11_ctx );
    mbedtls_pk_init( &transparent_ctx );

    /* Read a transparent key */
    TEST_ASSERT( mbedtls_pk_parse_keyfile( &transparent_ctx, file, NULL ) == 0 );

    /* Initialize cryptoki and import the key into the token */
    hSession = pkcs11_init( );
    TEST_ASSERT( hSession != CK_INVALID_HANDLE );

    TEST_ASSERT( mbedtls_pk_import_to_pkcs11( &transparent_ctx,
                                              MBEDTLS_PK_FLAG_SIGN |
                                              MBEDTLS_PK_FLAG_VERIFY,
                                              hSession,
                                              &hPublicKey,
                                              &hPrivateKey ) == 0 );
    TEST_ASSERT( hPublicKey != CK_INVALID_HANDLE );
    TEST_ASSERT( hPrivateKey != CK_INVALID_HANDLE );
    TEST_ASSERT( mbedtls_pk_setup_pkcs11( &pkcs11_ctx,
                                          hSession,
                                          hPublicKey,
                                          hPrivateKey ) == 0 );

    /* Sign with the token and verify in software */
    TEST_ASSERT( sizeof( sig_buffer ) >= mbedtls_pk_signature_size( &pkcs11_ctx ) );
    TEST_ASSERT( mbedtls_pk_sign( &pkcs11_ctx, MBEDTLS_MD_SHA256,
                                  hash_value, 32,
                                  sig_buffer, &sig_length,
                                  NULL, NULL ) == 0 );
    TEST_ASSERT( mbedtls_pk_verify( &transparent_ctx, MBEDTLS_MD_SHA256,
                                    hash_value, 32,
                                    sig_buffer, sig_length ) == 0 );

exit:
    if( hPublicKey != CK_INVALID_HANDLE )
        C_DestroyObject( hSession, hPublicKey );
    if( hPrivateKey != CK_INVALID_HANDLE )
        C_DestroyObject( hSession, hPrivateKey );
    C_CloseSession( hSession );
    C_Finalize( NULL_PTR );
    mbedtls_pk_free( &pkcs11_ctx );
    mbedtls_pk_free( &transparent_ctx );
}
/* END_CASE */

/* BEGIN_CASE depends_on:MBEDTLS_PK_C:MBEDTLS_SHA256_C */
void pk_import_sign_verify( char *file )
 {
    /* Sign with cryptoki, convert to Mbed TLS format and save,
       verify by cryptoki with a conversion to a raw, concatenated
       format by the engine. */
    mbedtls_pk_context pkcs11_ctx;
    mbedtls_pk_context transparent_ctx;
    CK_SESSION_HANDLE hSession = CK_INVALID_HANDLE;
    CK_OBJECT_HANDLE hPublicKey = CK_INVALID_HANDLE;
    CK_OBJECT_HANDLE hPrivateKey = CK_INVALID_HANDLE;
    unsigned char hash_value[32] = "Fake hash, it doesn't matter....";
    unsigned char sig_buffer[4096];
    size_t sig_length = sizeof( sig_buffer );

    mbedtls_pk_init( &pkcs11_ctx );
    mbedtls_pk_init( &transparent_ctx );

    /* Read a transparent key */
    TEST_ASSERT( mbedtls_pk_parse_keyfile( &transparent_ctx, file, NULL ) == 0 );

    /* Initialize cryptoki and import the key into the token */
    hSession = pkcs11_init( );
    TEST_ASSERT( hSession != CK_INVALID_HANDLE );

    TEST_ASSERT( mbedtls_pk_import_to_pkcs11( &transparent_ctx,
                                              MBEDTLS_PK_FLAG_SIGN |
                                              MBEDTLS_PK_FLAG_VERIFY,
                                              hSession,
                                              &hPublicKey,
                                              &hPrivateKey ) == 0 );
    TEST_ASSERT( hPublicKey != CK_INVALID_HANDLE );
    TEST_ASSERT( hPrivateKey != CK_INVALID_HANDLE );
    TEST_ASSERT( mbedtls_pk_setup_pkcs11( &pkcs11_ctx,
                                          hSession,
                                          hPublicKey,
                                          hPrivateKey ) == 0 );

    /* Sign with the token and verify with cryptoki */
    TEST_ASSERT( sizeof( sig_buffer ) >= mbedtls_pk_signature_size( &pkcs11_ctx ) );
    TEST_ASSERT( mbedtls_pk_sign( &pkcs11_ctx, MBEDTLS_MD_SHA256,
                                  hash_value, 32,
                                  sig_buffer, &sig_length,
                                  NULL, NULL ) == 0 );
    TEST_ASSERT( mbedtls_pk_verify( &pkcs11_ctx, MBEDTLS_MD_SHA256,
                                    hash_value, 32,
                                    sig_buffer, sig_length ) == 0 );

exit:
    mbedtls_pk_free( &pkcs11_ctx );
    mbedtls_pk_free( &transparent_ctx );
    if( hPublicKey != CK_INVALID_HANDLE )
        C_DestroyObject( hSession, hPublicKey );
    if( hPrivateKey != CK_INVALID_HANDLE )
        C_DestroyObject( hSession, hPrivateKey );
    C_CloseSession( hSession );
    C_Finalize( NULL_PTR );
}
/* END_CASE */

/* BEGIN_CASE depends_on:MBEDTLS_PK_C:MBEDTLS_SHA256_C */
void pk_import_verify_signed( char *file )
{
    /* Sign with mbedTLS, verify by cryptoki with a conversion
       to a raw, concatenated format by the engine. */
    mbedtls_pk_context pkcs11_ctx;
    mbedtls_pk_context transparent_ctx;
    CK_SESSION_HANDLE hSession = CK_INVALID_HANDLE;
    CK_OBJECT_HANDLE hPublicKey = CK_INVALID_HANDLE;
    CK_OBJECT_HANDLE hPrivateKey = CK_INVALID_HANDLE;
    unsigned char hash_value[32] = "Fake hash, it doesn't matter....";
    unsigned char sig_buffer[4096];
    size_t sig_length = sizeof( sig_buffer );

    mbedtls_pk_init( &pkcs11_ctx );
    mbedtls_pk_init( &transparent_ctx );

    /* Read a transparent key */
    TEST_ASSERT( mbedtls_pk_parse_keyfile( &transparent_ctx, file, NULL ) == 0 );

    /* Initialize cryptoki and import the key into the token */
    hSession = pkcs11_init( );
    TEST_ASSERT( hSession != CK_INVALID_HANDLE );

    TEST_ASSERT( mbedtls_pk_import_to_pkcs11( &transparent_ctx,
                                              MBEDTLS_PK_FLAG_VERIFY,
                                              hSession,
                                              &hPublicKey,
                                              NULL ) == 0 );
    TEST_ASSERT( hPublicKey != CK_INVALID_HANDLE );
    TEST_ASSERT( mbedtls_pk_setup_pkcs11( &pkcs11_ctx,
                                          hSession,
                                          hPublicKey,
                                          CK_INVALID_HANDLE ) == 0 );

    /* Sign with the token and verify with cryptoki */
    TEST_ASSERT( sizeof( sig_buffer ) >= mbedtls_pk_signature_size( &pkcs11_ctx ) );
    TEST_ASSERT( mbedtls_pk_sign( &transparent_ctx, MBEDTLS_MD_SHA256,
                                  hash_value, 32,
                                  sig_buffer, &sig_length,
                                  NULL, NULL ) == 0 );
    TEST_ASSERT( mbedtls_pk_verify( &pkcs11_ctx, MBEDTLS_MD_SHA256,
                                    hash_value, 32,
                                    sig_buffer, sig_length ) == 0 );

exit:
    if( hPublicKey != CK_INVALID_HANDLE )
        C_DestroyObject( hSession, hPublicKey );
    if( hPrivateKey != CK_INVALID_HANDLE )
        C_DestroyObject( hSession, hPrivateKey );
    C_CloseSession( hSession );
    C_Finalize( NULL_PTR );
    mbedtls_pk_free( &pkcs11_ctx );
    mbedtls_pk_free( &transparent_ctx );
}
/* END_CASE */

/* BEGIN_CASE depends_on:MBEDTLS_PK_C:MBEDTLS_RSA_C */
void pk_rsa_hardcoded_verify( char *message_hex_string, int digest,
                       int mod, int radix_N, char *input_N, int radix_E,
                       char *input_E, char *result_hex_str, int result )
{
    unsigned char message_str[1000];
    unsigned char hash_result[1000];
    unsigned char result_str[1000];
    mbedtls_rsa_context *rsa;
    mbedtls_pk_context transparent_ctx;
    int msg_len;

    mbedtls_pk_context pkcs11_ctx;
    CK_SESSION_HANDLE hSession = CK_INVALID_HANDLE;
    CK_OBJECT_HANDLE hPublicKey = CK_INVALID_HANDLE;
    CK_OBJECT_HANDLE hPrivateKey = CK_INVALID_HANDLE;

    mbedtls_pk_init( &transparent_ctx );

    memset( message_str, 0x00, 1000 );
    memset( hash_result, 0x00, 1000 );
    memset( result_str, 0x00, 1000 );

    TEST_ASSERT( mbedtls_pk_setup( &transparent_ctx, mbedtls_pk_info_from_type( MBEDTLS_PK_RSA ) ) == 0 );
    rsa = mbedtls_pk_rsa( transparent_ctx );

    rsa->len = mod / 8;
    TEST_ASSERT( mbedtls_mpi_read_string( &rsa->N, radix_N, input_N ) == 0 );
    TEST_ASSERT( mbedtls_mpi_read_string( &rsa->E, radix_E, input_E ) == 0 );

    msg_len = unhexify( message_str, message_hex_string );
    unhexify( result_str, result_hex_str );

    if( mbedtls_md_info_from_type( digest ) != NULL )
        TEST_ASSERT( mbedtls_md( mbedtls_md_info_from_type( digest ), message_str, msg_len, hash_result ) == 0 );

    // PKCS11 part
    mbedtls_pk_init( &pkcs11_ctx );

    /* Initialize cryptoki and import the key into the token */
    hSession = pkcs11_init( );
    TEST_ASSERT( hSession != CK_INVALID_HANDLE );
    TEST_ASSERT( mbedtls_pk_import_to_pkcs11( &transparent_ctx,
                                              MBEDTLS_PK_FLAG_SIGN |
                                              MBEDTLS_PK_FLAG_VERIFY,
                                              hSession,
                                              &hPublicKey,
                                              NULL ) == 0 );
    TEST_ASSERT( hPublicKey != CK_INVALID_HANDLE );
    TEST_ASSERT( mbedtls_pk_setup_pkcs11( &pkcs11_ctx,
                                          hSession,
                                          hPublicKey,
                                          CK_INVALID_HANDLE ) == 0 );

    TEST_ASSERT( mbedtls_pk_verify( &pkcs11_ctx, digest, hash_result, 0,
                            result_str, mbedtls_pk_get_len( &transparent_ctx ) ) == result );

exit:
    if( hPublicKey != CK_INVALID_HANDLE )
        C_DestroyObject( hSession, hPublicKey );
    if( hPrivateKey != CK_INVALID_HANDLE )
        C_DestroyObject( hSession, hPrivateKey );
    C_CloseSession( hSession );
    C_Finalize( NULL_PTR );
    mbedtls_pk_free( &pkcs11_ctx );
    mbedtls_pk_free( &transparent_ctx );
}
/* END_CASE */