/* BEGIN_HEADER */

#include "tinycrypt/ecc.h"
#include "tinycrypt/ecc_dh.h"
#include "tinycrypt/ecc_dsa.h"

#include "mbedtls/sha256.h"

/* END_HEADER */

/* BEGIN_DEPENDENCIES
 * depends_on:MBEDTLS_USE_TINYCRYPT
 * END_DEPENDENCIES
 */

/* BEGIN_CASE depends_on:MBEDTLS_USE_TINYCRYPT */
void test_ecdh()
{
    uint8_t private1[NUM_ECC_BYTES] = {0};
    uint8_t private2[NUM_ECC_BYTES] = {0};
    uint8_t public1[2*NUM_ECC_BYTES] = {0};
    uint8_t public2[2*NUM_ECC_BYTES] = {0};
    uint8_t secret1[NUM_ECC_BYTES] = {0};
    uint8_t secret2[NUM_ECC_BYTES] = {0};

    uECC_set_rng( &uecc_rng_wrapper );

    TEST_ASSERT( uECC_make_key( public1, private1 ) == UECC_SUCCESS );

    TEST_ASSERT( uECC_make_key( public2, private2 ) == UECC_SUCCESS );

    TEST_ASSERT( uECC_shared_secret( public2, private1, secret1 ) == UECC_SUCCESS );

    TEST_ASSERT( uECC_shared_secret( public1, private2, secret2 ) == UECC_SUCCESS );

    TEST_ASSERT( memcmp( secret1, secret2, sizeof( secret1 ) ) == 0 );
}
/* END_CASE */

/* BEGIN_CASE depends_on:MBEDTLS_USE_TINYCRYPT */
void test_ecdsa()
{
    uint8_t private[NUM_ECC_BYTES] = {0};
    uint8_t public[2*NUM_ECC_BYTES] = {0};
    uint8_t hash[NUM_ECC_BYTES] = {0};
    uint8_t sig[2*NUM_ECC_BYTES] = {0};

    uECC_set_rng( &uecc_rng_wrapper );

    TEST_ASSERT( rnd_std_rand( NULL, hash, NUM_ECC_BYTES ) == 0 );

    TEST_ASSERT( uECC_make_key( public, private ) == UECC_SUCCESS );

    TEST_ASSERT( uECC_sign( private, hash, sizeof( hash ), sig ) == UECC_SUCCESS );

    TEST_ASSERT( uECC_verify( public, hash, sizeof( hash ), sig ) == UECC_SUCCESS );
}
/* END_CASE */

/* BEGIN_CASE depends_on:MBEDTLS_USE_TINYCRYPT */
void ecdh_primitive_vector( data_t * qx_str, data_t * qy_str,
                            data_t * d_str, data_t * z_str )
{
    uint8_t public[2*NUM_ECC_BYTES] = {0};
    uint8_t private[NUM_ECC_BYTES] = {0};
    uint8_t secret[NUM_ECC_BYTES] = {0};

    memcpy( public, qx_str->x, qx_str->len );
    memcpy( public + NUM_ECC_BYTES, qy_str->x, qy_str->len );

    memcpy( private, d_str->x, d_str->len );

    // Compute shared secret and compare to test vector secret
    TEST_ASSERT( uECC_shared_secret( public, d_str->x, secret ) == UECC_SUCCESS );

    TEST_ASSERT( memcmp( secret, z_str->x, sizeof( secret ) ) == 0 );
}
/* END_CASE */

/* BEGIN_CASE depends_on:MBEDTLS_USE_TINYCRYPT */
void ecdh_primitive_testvec( data_t * private1, data_t * xA_str,
                             data_t * yA_str, data_t * private2,
                             data_t * xB_str, data_t * yB_str, data_t * z_str )
{
    uint8_t public1[2*NUM_ECC_BYTES] = {0};
    uint8_t public2[2*NUM_ECC_BYTES] = {0};
    uint8_t secret1[NUM_ECC_BYTES] = {0};
    uint8_t secret2[NUM_ECC_BYTES] = {0};

    memcpy( public1, xA_str->x, xA_str->len );
    memcpy( public1 + NUM_ECC_BYTES, yA_str->x, yA_str->len );
    memcpy( public2, xB_str->x, xB_str->len );
    memcpy( public2 + NUM_ECC_BYTES, yB_str->x, yB_str->len );

    // Compute shared secrets and compare to test vector secret
    TEST_ASSERT( uECC_shared_secret( public2, private1->x, secret1 ) == UECC_SUCCESS );

    TEST_ASSERT( uECC_shared_secret( public1, private2->x, secret2 ) == UECC_SUCCESS );

    TEST_ASSERT( memcmp( secret1, secret2, sizeof( secret1 ) ) == 0 );
    TEST_ASSERT( memcmp( secret1, z_str->x, sizeof( secret1 ) ) == 0 );
    TEST_ASSERT( memcmp( secret2, z_str->x, sizeof( secret2 ) ) == 0 );
}
/* END_CASE */

/* BEGIN_CASE depends_on:MBEDTLS_USE_TINYCRYPT */
void ecdsa_verify_msg( data_t * message, data_t * xQ_str, data_t * yQ_str,
                       data_t * r_str, data_t * s_str,
                       int res )
{
    uint8_t pub_bytes[2*NUM_ECC_BYTES] = {0};
    uint8_t sig_bytes[2*NUM_ECC_BYTES] = {0};
    unsigned char hash[32] = {0};

    memcpy( pub_bytes, xQ_str->x, xQ_str->len );
    memcpy( pub_bytes + NUM_ECC_BYTES, yQ_str->x, yQ_str->len );
    memcpy( sig_bytes, r_str->x, r_str->len );
    memcpy( sig_bytes + NUM_ECC_BYTES, s_str->x, r_str->len );

    TEST_ASSERT( mbedtls_sha256_ret( message->x, message->len, hash, 0 ) == 0 );
    TEST_ASSERT( uECC_verify( pub_bytes, hash, sizeof(hash),
                              sig_bytes ) == res );
}
/* END_CASE */

/* BEGIN_CASE depends_on:MBEDTLS_USE_TINYCRYPT */
void ecdsa_sig( data_t * message, data_t * d_str, data_t * k_str,
                data_t * r_str, data_t * s_str )
{
    uint8_t sig_bytes[2*NUM_ECC_BYTES] = {0};
    uint8_t hash[32] = {0};
    uint8_t k[NUM_ECC_BYTES] = {0};
    uint8_t i = 0;
    /* Reverse k as expected by tinycrypt */
    for( i = 0; i < NUM_ECC_BYTES; i++ )
    {
        k[i] = k_str->x[NUM_ECC_BYTES - i - 1];
    }
    TEST_ASSERT( mbedtls_sha256_ret( message->x, message->len, hash, 0 ) == 0 );
    TEST_ASSERT( uECC_sign_with_k( d_str->x, hash, sizeof(hash), (uECC_word_t*) k,
                                   sig_bytes ) == UECC_SUCCESS );
    TEST_ASSERT( memcmp( sig_bytes, r_str->x, NUM_ECC_BYTES) == 0 );
    TEST_ASSERT( memcmp( sig_bytes + NUM_ECC_BYTES, s_str->x, NUM_ECC_BYTES ) == 0 );
}
/* END_CASE */

/* BEGIN_CASE depends_on:MBEDTLS_USE_TINYCRYPT */
void ecdsa_primitive_testvec( data_t * xQ_str, data_t * yQ_str,
                              data_t * hash, data_t * r_str, data_t * s_str )
{
    uint8_t pub_bytes[2*NUM_ECC_BYTES] = {0};
    uint8_t sig_bytes[2*NUM_ECC_BYTES] = {0};

    memcpy( pub_bytes, xQ_str->x, xQ_str->len );
    memcpy( pub_bytes + NUM_ECC_BYTES, yQ_str->x, yQ_str->len );
    memcpy( sig_bytes, r_str->x, r_str->len );
    memcpy( sig_bytes + NUM_ECC_BYTES, s_str->x, r_str->len );

    TEST_ASSERT( uECC_verify( pub_bytes, hash->x, hash->len,
                              sig_bytes ) == UECC_SUCCESS );

    // Alter the signature and check the verification fails
    for( int i = 0; i < 2*NUM_ECC_BYTES; i++ )
    {
        uint8_t temp = sig_bytes[i];
        sig_bytes[i] = ( sig_bytes[i] + 1 ) % 256;
        TEST_ASSERT( uECC_verify( pub_bytes, hash->x, hash->len,
                                  sig_bytes ) == UECC_FAILURE );
        sig_bytes[i] = temp;
    }

}
/* END_CASE */