/* 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 512
#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;
    unsigned char ecParams[16];
    size_t ecParams_length;

    switch( key_type )
    {
#if defined(MBEDTLS_ECDSA_C)
    case MBEDTLS_PK_ECDSA:
        {
            const char *oid;
            size_t oid_length;
            unsigned char *p;
            TEST_ASSERT( mbedtls_oid_get_oid_by_ec_grp( ECP_GROUP_ID,
                                                        &oid, &oid_length ) == 0 );
            ecParams_length = 2 + oid_length;
            p = ecParams + ecParams_length;
            TEST_ASSERT( mbedtls_asn1_write_oid( &p, ecParams, oid, oid_length ) > 0 );
        }
        mechanism.mechanism = CKM_EC_KEY_PAIR_GEN;
        public_attributes[0].type = CKA_EC_PARAMS;
        public_attributes[0].pValue = ecParams;
        public_attributes[0].ulValueLen = ecParams_length;
        break;
#endif /* MBEDTLS_ECDSA_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 ) );
exit:
    /* Shouldn't happen except if there's a test error (e.g. trying to
       use a curve that isn't compiled in). */
    return( -1 );
}


#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_ECDSA_C)
    case MBEDTLS_PK_ECDSA:
        {
            unsigned char ecParams[16];
            unsigned char ecPoint[128];
            CK_ATTRIBUTE public_attributes[] = {
                {CKA_EC_PARAMS, ecParams, sizeof( ecParams )},
                {CKA_EC_POINT, ecPoint, sizeof( ecPoint )},
            };
            mbedtls_ecp_keypair *ecp_ctx = mbedtls_pk_ec( transparent_ctx );

            CK_ASSERT( C_GetAttributeValue( hSession, hPublicKey,
                                            public_attributes, ARRAY_LENGTH( public_attributes ) ) );
            // TODO: lift out a function or two from pkparse.c
            // * pk_get_ecparams followed by pk_use_ecparams for ecParams?
            // * Some code from pk_group_from_specified to read an octet string for ecPoint?
            {
                mbedtls_asn1_buf params_asn1;
                CK_ULONG ecParams_length = public_attributes[0].ulValueLen;
                mbedtls_ecp_group_id grp_id;
                params_asn1.tag = ecParams[0];
                params_asn1.len = ecParams[1];
                params_asn1.p = ecParams + 2;
                TEST_ASSERT( ecParams_length == 2 + params_asn1.len );
                TEST_ASSERT( mbedtls_oid_get_ec_grp( &params_asn1, &grp_id ) == 0 );
                TEST_ASSERT( mbedtls_ecp_group_load( &ecp_ctx->grp, grp_id ) == 0 );
            }
            {
                unsigned char *p = ecPoint;
                size_t len;
                CK_ULONG ecPoint_length = public_attributes[1].ulValueLen;
                TEST_ASSERT( mbedtls_asn1_get_tag( &p,
                                                   ecPoint + ecPoint_length,
                                                   &len,
                                                   MBEDTLS_ASN1_OCTET_STRING ) == 0 );
                TEST_ASSERT( mbedtls_ecp_point_read_binary( &ecp_ctx->grp,
                                                            &ecp_ctx->Q,
                                                            p, len ) == 0 );
            }
        }
        break;
#endif /* MBEDTLS_ECDSA_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 mbedTLS 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:
    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_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_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 );

    /* 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_ECDSA_C */
void pk_ecdsa_hardcoded_verify( int type, int id, char *key_str,
                     char *hash_str, char * sig_str, int ret )
{
    mbedtls_pk_context transparent_ctx;
    mbedtls_ecp_keypair *eckey;
    unsigned char hash[100], sig[500], key[500];
    size_t hash_len, sig_len, key_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( hash, 0, sizeof( hash ) );  hash_len = unhexify(hash, hash_str);
    memset( sig, 0, sizeof( sig ) );    sig_len = unhexify(sig, sig_str);
    memset( key, 0, sizeof( key ) );    key_len = unhexify(key, key_str);

    TEST_ASSERT( mbedtls_pk_setup( &transparent_ctx, mbedtls_pk_info_from_type( type ) ) == 0 );

    TEST_ASSERT( mbedtls_pk_can_do( &transparent_ctx, MBEDTLS_PK_ECDSA ) );
    eckey = mbedtls_pk_ec( transparent_ctx );

    TEST_ASSERT( mbedtls_ecp_group_load( &eckey->grp, id ) == 0 );
    TEST_ASSERT( mbedtls_ecp_point_read_binary( &eckey->grp, &eckey->Q,
                                        key, key_len ) == 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, MBEDTLS_MD_NONE,
                            hash, hash_len, sig, sig_len ) == ret );

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 */