/* test_ecc_encrypt.c - TinyCrypt test for client-server key exchange and encryption */

/*
 *  Copyright (C) 2017 by Intel Corporation, All Rights Reserved.
 *
 *  Redistribution and use in source and binary forms, with or without
 *  modification, are permitted provided that the following conditions are met:
 *
 *    - Redistributions of source code must retain the above copyright notice,
 *     this list of conditions and the following disclaimer.
 *
 *    - Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 *    - Neither the name of Intel Corporation nor the names of its contributors
 *    may be used to endorse or promote products derived from this software
 *    without specific prior written permission.
 *
 *  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 *  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 *  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 *  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
 *  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 *  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 *  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 *  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 *  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 *  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 *  POSSIBILITY OF SUCH DAMAGE.
 */

#include <tinycrypt/ecc.h>
#include <tinycrypt/ecc_dh.h>
#include <tinycrypt/aes.h>
#include <tinycrypt/ctr_mode.h>
#include <tinycrypt/constants.h>
#include <tinycrypt/ecc_platform_specific.h>

#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include "../lib/debug_config.h"

#define TC_PASS 0
#define TC_FAIL 1

#define TC_PRINT(...) DEBUG_INFO(DEBUG_CATEGORY_CRYPTO, __VA_ARGS__)
#define TC_ERROR(...) DEBUG_ERROR(DEBUG_CATEGORY_CRYPTO, __VA_ARGS__)
#define TC_START(name) DEBUG_INFO(DEBUG_CATEGORY_CRYPTO, "Starting test: %s", name)
#define TC_END_RESULT(result) DEBUG_INFO(DEBUG_CATEGORY_CRYPTO, "Result: %s", (result) == TC_PASS ? "PASS" : "FAIL")
#define TC_END_REPORT(result) DEBUG_INFO(DEBUG_CATEGORY_CRYPTO, "Test %s", (result) == TC_PASS ? "SUCCESSFUL" : "FAILED")

static void show_str(const char *label, const uint8_t *s, size_t len)
{
    char hex_buf[256];
    char *ptr = hex_buf;
    int remaining = sizeof(hex_buf);
    int written = snprintf(ptr, remaining, "%s = ", label);
    if (written < 0 || written >= remaining) return;
    ptr += written;
    remaining -= written;
    
    for (unsigned int i = 0; i < len && remaining > 2; ++i) {
        written = snprintf(ptr, remaining, "%02x", s[i]);
        if (written < 0 || written >= remaining) break;
        ptr += written;
        remaining -= written;
    }
    DEBUG_DEBUG(DEBUG_CATEGORY_CRYPTO, "%s", hex_buf);
}

static int test_key_exchange_and_encrypt(void)
{
    const struct uECC_Curve_t *curve = uECC_secp256r1();

    /* Client key pair */
    uint8_t client_private[NUM_ECC_BYTES];
    uint8_t client_public[2 * NUM_ECC_BYTES];

    /* Server key pair */
    uint8_t server_private[NUM_ECC_BYTES];
    uint8_t server_public[2 * NUM_ECC_BYTES];

    /* Shared secrets */
    uint8_t client_secret[NUM_ECC_BYTES];
    uint8_t server_secret[NUM_ECC_BYTES];

    TC_PRINT("Generating client key pair...\n");
    if (!uECC_make_key(client_public, client_private, curve)) {
        TC_ERROR("Failed to generate client key pair\n");
        return TC_FAIL;
    }

    TC_PRINT("Generating server key pair...\n");
    if (!uECC_make_key(server_public, server_private, curve)) {
        TC_ERROR("Failed to generate server key pair\n");
        return TC_FAIL;
    }

    /* Compute shared secret on client side */
    TC_PRINT("Computing shared secret on client side...\n");
    if (!uECC_shared_secret(server_public, client_private, client_secret, curve)) {
        TC_ERROR("Failed to compute client shared secret\n");
        return TC_FAIL;
    }

    /* Compute shared secret on server side */
    TC_PRINT("Computing shared secret on server side...\n");
    if (!uECC_shared_secret(client_public, server_private, server_secret, curve)) {
        TC_ERROR("Failed to compute server shared secret\n");
        return TC_FAIL;
    }

    /* Verify that both shared secrets match */
    if (memcmp(client_secret, server_secret, NUM_ECC_BYTES) != 0) {
        TC_ERROR("Shared secrets do not match\n");
        show_str("Client secret", client_secret, NUM_ECC_BYTES);
        show_str("Server secret", server_secret, NUM_ECC_BYTES);
        return TC_FAIL;
    }
    TC_PRINT("Shared secrets match.\n");

    /* Derive AES key from shared secret (first 16 bytes) */
    uint8_t aes_key[TC_AES_KEY_SIZE];
    memcpy(aes_key, client_secret, TC_AES_KEY_SIZE);

    /* Prepare test data */
    const char *plaintext = "Hello, secure world!";
    size_t plaintext_len = strlen(plaintext) + 1; /* include null terminator */
    uint8_t ciphertext[plaintext_len];
    uint8_t decrypted[plaintext_len];

    /* Initialize AES key schedule */
    struct tc_aes_key_sched_struct sched;
    if (tc_aes128_set_encrypt_key(&sched, aes_key) != TC_CRYPTO_SUCCESS) {
        TC_ERROR("Failed to set AES encryption key\n");
        return TC_FAIL;
    }

    /* Initialize counter (little-endian 32-bit) */
    uint8_t ctr[TC_AES_BLOCK_SIZE] = {0}; /* zero-initialized counter */

    TC_PRINT("Encrypting plaintext...\n");
    if (tc_ctr_mode(ciphertext, plaintext_len, (const uint8_t *)plaintext,
                    plaintext_len, ctr, &sched) != TC_CRYPTO_SUCCESS) {
        TC_ERROR("Encryption failed\n");
        return TC_FAIL;
    }

    /* Reset counter for decryption (CTR mode uses same counter sequence) */
    memset(ctr, 0, TC_AES_BLOCK_SIZE);

    TC_PRINT("Decrypting ciphertext...\n");
    if (tc_ctr_mode(decrypted, plaintext_len, ciphertext,
                    plaintext_len, ctr, &sched) != TC_CRYPTO_SUCCESS) {
        TC_ERROR("Decryption failed\n");
        return TC_FAIL;
    }

    /* Verify decrypted matches original */
    if (memcmp(plaintext, decrypted, plaintext_len) != 0) {
        TC_ERROR("Decrypted text does not match original\n");
        show_str("Original", (const uint8_t *)plaintext, plaintext_len);
        show_str("Decrypted", decrypted, plaintext_len);
        return TC_FAIL;
    }
    TC_PRINT("Decryption successful. Test passed.\n");

    return TC_PASS;
}

int main(void)
{
    debug_config_init();
    debug_set_level(DEBUG_LEVEL_TRACE);
    debug_set_categories(DEBUG_CATEGORY_ALL);
    
    unsigned int result = TC_PASS;

    DEBUG_INFO(DEBUG_CATEGORY_CRYPTO, "Performing client-server key exchange and encryption test:");

    /* Setup cryptographically secure PRNG */
    uECC_set_rng(&default_CSPRNG);

    result = test_key_exchange_and_encrypt();

    DEBUG_INFO(DEBUG_CATEGORY_CRYPTO, "Result: %s", result == TC_PASS ? "PASS" : "FAIL");
    DEBUG_INFO(DEBUG_CATEGORY_CRYPTO, "Test %s", result == TC_PASS ? "SUCCESSFUL" : "FAILED");

    return result;
}