/***************************************************************************
 * Copyright (C) 2026 Eclipse ThreadX contributors
 *
 * This program and the accompanying materials are made available under the
 * terms of the MIT License which is available at
 * https://opensource.org/licenses/MIT.
 *
 * AI Disclosure: This file was largely AI-generated by Codex gpt-5.5.
 * The AI-generated portions may be considered public domain (CC0-1.0)
 * and not subject to the project's licence. The human contributor has
 * reviewed and verified that the code is correct.
 *
 * SPDX-License-Identifier: MIT and CC0-1.0
 **************************************************************************/

#include <stdio.h>
#include <string.h>

#include "nx_secure_tls.h"
#include "nx_crypto.h"
#include "tls_test_utility.h"

extern void test_control_return(UINT status);

#ifdef NX_SECURE_ENABLE_PSK_CIPHERSUITES
static UCHAR psk_one[] = {0x11};
static UCHAR psk_two[] = {0x22, 0x33, 0x44};
static UCHAR psk_id_one[] = "first";
static UCHAR psk_id_two[] = "second";
static UCHAR psk_hint[] = "hint";
static UCHAR ecdhe_public_key[] = {0xAA, 0xBB, 0xCC};
static UCHAR ecdhe_shared_secret[] = {0x55, 0x66, 0x77, 0x88};

static NX_CRYPTO_METHOD crypto_method_test_null =
{
    NX_CRYPTO_ENCRYPTION_NULL,
    0,
    0,
    0,
    0,
    0,
    NX_CRYPTO_NULL,
    NX_CRYPTO_NULL,
    NX_CRYPTO_NULL
};

static NX_CRYPTO_METHOD crypto_method_test_psk =
{
    NX_CRYPTO_KEY_EXCHANGE_PSK,
    0,
    0,
    0,
    0,
    0,
    NX_CRYPTO_NULL,
    NX_CRYPTO_NULL,
    NX_CRYPTO_NULL
};

static NX_SECURE_TLS_CIPHERSUITE_INFO psk_ciphersuite =
{
    TLS_PSK_WITH_AES_128_CBC_SHA256,
    &crypto_method_test_null,
    &crypto_method_test_psk,
    &crypto_method_test_null,
    0,
    0,
    &crypto_method_test_null,
    0,
    &crypto_method_test_null
};

/* Store two PSKs so the tests can verify that the requested identity selects a non-zero store entry. */
static VOID psk_credentials_setup(NX_SECURE_TLS_CREDENTIALS *credentials)
{

    memset(credentials, 0, sizeof(NX_SECURE_TLS_CREDENTIALS));

    credentials -> nx_secure_tls_psk_count = 2;

    NX_SECURE_MEMCPY(credentials -> nx_secure_tls_psk_store[0].nx_secure_tls_psk_data, psk_one, sizeof(psk_one));
    credentials -> nx_secure_tls_psk_store[0].nx_secure_tls_psk_data_size = sizeof(psk_one);
    NX_SECURE_MEMCPY(credentials -> nx_secure_tls_psk_store[0].nx_secure_tls_psk_id, psk_id_one, sizeof(psk_id_one) - 1u);
    credentials -> nx_secure_tls_psk_store[0].nx_secure_tls_psk_id_size = sizeof(psk_id_one) - 1u;
    NX_SECURE_MEMCPY(credentials -> nx_secure_tls_psk_store[0].nx_secure_tls_psk_id_hint, psk_hint, sizeof(psk_hint) - 1u);
    credentials -> nx_secure_tls_psk_store[0].nx_secure_tls_psk_id_hint_size = sizeof(psk_hint) - 1u;

    NX_SECURE_MEMCPY(credentials -> nx_secure_tls_psk_store[1].nx_secure_tls_psk_data, psk_two, sizeof(psk_two));
    credentials -> nx_secure_tls_psk_store[1].nx_secure_tls_psk_data_size = sizeof(psk_two);
    NX_SECURE_MEMCPY(credentials -> nx_secure_tls_psk_store[1].nx_secure_tls_psk_id, psk_id_two, sizeof(psk_id_two) - 1u);
    credentials -> nx_secure_tls_psk_store[1].nx_secure_tls_psk_id_size = sizeof(psk_id_two) - 1u;
    NX_SECURE_MEMCPY(credentials -> nx_secure_tls_psk_store[1].nx_secure_tls_psk_id_hint, psk_hint, sizeof(psk_hint) - 1u);
    credentials -> nx_secure_tls_psk_store[1].nx_secure_tls_psk_id_hint_size = sizeof(psk_hint) - 1u;

    NX_SECURE_MEMCPY(credentials -> nx_secure_tls_remote_psk_id, psk_id_two, sizeof(psk_id_two) - 1u);
    credentials -> nx_secure_tls_remote_psk_id_size = sizeof(psk_id_two) - 1u;
}

/* Validate that a plain PSK premaster uses the client-requested PSK identity. */
static VOID plain_psk_identity_test(VOID)
{
NX_SECURE_TLS_CREDENTIALS credentials;
NX_SECURE_TLS_KEY_MATERIAL key_material;
USHORT received_remote_credentials = NX_FALSE;
UINT status;

    psk_credentials_setup(&credentials);
    memset(&key_material, 0, sizeof(key_material));

    status = _nx_secure_generate_premaster_secret(&psk_ciphersuite, NX_SECURE_TLS_VERSION_TLS_1_2,
                                                  &key_material, &credentials,
                                                  NX_SECURE_TLS_SESSION_TYPE_SERVER,
                                                  &received_remote_credentials,
                                                  NX_NULL, 0, NX_NULL);
    EXPECT_EQ(NX_SECURE_TLS_SUCCESS, status);
    EXPECT_EQ(NX_TRUE, received_remote_credentials);
    EXPECT_EQ((UINT)(2u + sizeof(psk_two) + 2u + sizeof(psk_two)), key_material.nx_secure_tls_pre_master_secret_size);
    EXPECT_EQ(0, key_material.nx_secure_tls_pre_master_secret[0]);
    EXPECT_EQ(sizeof(psk_two), key_material.nx_secure_tls_pre_master_secret[1]);
    EXPECT_EQ(0, NX_SECURE_MEMCMP(&key_material.nx_secure_tls_pre_master_secret[2u + sizeof(psk_two) + 2u], psk_two, sizeof(psk_two)));

    NX_SECURE_MEMCPY(credentials.nx_secure_tls_remote_psk_id, psk_hint, sizeof(psk_hint) - 1u);
    credentials.nx_secure_tls_remote_psk_id_size = sizeof(psk_hint) - 1u;
    status = _nx_secure_generate_premaster_secret(&psk_ciphersuite, NX_SECURE_TLS_VERSION_TLS_1_2,
                                                  &key_material, &credentials,
                                                  NX_SECURE_TLS_SESSION_TYPE_SERVER,
                                                  &received_remote_credentials,
                                                  NX_NULL, 0, NX_NULL);
    EXPECT_EQ(NX_SECURE_TLS_NO_MATCHING_PSK, status);
}

/* Validate ClientKeyExchange PSK identity bounds checking for plain PSK. */
static VOID plain_psk_client_key_exchange_length_test(VOID)
{
NX_SECURE_TLS_CREDENTIALS credentials;
NX_SECURE_TLS_KEY_MATERIAL key_material;
USHORT received_remote_credentials = NX_FALSE;
UCHAR truncated_identity[] = {0x00, 0x06, 's', 'e'};
UINT status;

    psk_credentials_setup(&credentials);
    memset(&key_material, 0, sizeof(key_material));

    status = _nx_secure_process_client_key_exchange(&psk_ciphersuite, NX_SECURE_TLS_VERSION_TLS_1_2,
                                                    truncated_identity, sizeof(truncated_identity),
                                                    &received_remote_credentials, &key_material, &credentials,
                                                    NX_NULL, 0, NX_NULL, 0, NX_NULL);
    EXPECT_EQ(NX_SECURE_TLS_INCORRECT_MESSAGE_LENGTH, status);
}

#ifdef NX_SECURE_ENABLE_ECC_CIPHERSUITE
#ifndef NX_SECURE_DISABLE_X509
static UINT test_ecdhe_operation(UINT op, VOID *handler, NX_CRYPTO_METHOD *method,
                                 UCHAR *key, NX_CRYPTO_KEY_SIZE key_size_in_bits,
                                 UCHAR *input, ULONG input_length_in_byte,
                                 UCHAR *iv_ptr, UCHAR *output, ULONG output_length_in_byte,
                                 VOID *crypto_metadata, ULONG crypto_metadata_size,
                                 VOID *packet_ptr,
                                 VOID (*nx_crypto_hw_process_callback)(VOID *packet_ptr, UINT status))
{
NX_CRYPTO_EXTENDED_OUTPUT *extended_output;

    NX_PARAMETER_NOT_USED(handler);
    NX_PARAMETER_NOT_USED(method);
    NX_PARAMETER_NOT_USED(key);
    NX_PARAMETER_NOT_USED(key_size_in_bits);
    NX_PARAMETER_NOT_USED(iv_ptr);
    NX_PARAMETER_NOT_USED(output_length_in_byte);
    NX_PARAMETER_NOT_USED(crypto_metadata);
    NX_PARAMETER_NOT_USED(crypto_metadata_size);
    NX_PARAMETER_NOT_USED(packet_ptr);
    NX_PARAMETER_NOT_USED(nx_crypto_hw_process_callback);

    if ((op == NX_CRYPTO_EC_CURVE_SET) || (op == NX_CRYPTO_DH_KEY_PAIR_IMPORT))
    {
        return(NX_CRYPTO_SUCCESS);
    }

    if (op == NX_CRYPTO_DH_CALCULATE)
    {
        if ((input_length_in_byte != sizeof(ecdhe_public_key)) ||
            (NX_SECURE_MEMCMP(input, ecdhe_public_key, sizeof(ecdhe_public_key)) != 0))
        {
            return(NX_CRYPTO_NOT_SUCCESSFUL);
        }

        extended_output = (NX_CRYPTO_EXTENDED_OUTPUT *)output;
        if (extended_output -> nx_crypto_extended_output_length_in_byte < sizeof(ecdhe_shared_secret))
        {
            return(NX_CRYPTO_NOT_SUCCESSFUL);
        }

        NX_SECURE_MEMCPY(extended_output -> nx_crypto_extended_output_data, ecdhe_shared_secret, sizeof(ecdhe_shared_secret));
        extended_output -> nx_crypto_extended_output_actual_size = sizeof(ecdhe_shared_secret);
        return(NX_CRYPTO_SUCCESS);
    }

    return(NX_CRYPTO_NOT_SUCCESSFUL);
}

static NX_CRYPTO_METHOD crypto_method_test_ecdhe =
{
    NX_CRYPTO_KEY_EXCHANGE_ECDHE,
    0,
    0,
    0,
    0,
    0,
    NX_CRYPTO_NULL,
    NX_CRYPTO_NULL,
    test_ecdhe_operation
};

static NX_CRYPTO_METHOD crypto_method_test_curve =
{
    NX_CRYPTO_EC_SECP256R1,
    0,
    0,
    0,
    0,
    0,
    NX_CRYPTO_NULL,
    NX_CRYPTO_NULL,
    NX_CRYPTO_NULL
};

static NX_SECURE_TLS_CIPHERSUITE_INFO ecdhe_psk_ciphersuite =
{
    TLS_ECDHE_PSK_WITH_AES_128_CBC_SHA,
    &crypto_method_test_ecdhe,
    &crypto_method_test_psk,
    &crypto_method_test_null,
    0,
    0,
    &crypto_method_test_null,
    0,
    &crypto_method_test_null
};

/* Validate that ECDHE-PSK parses identity before public key and appends the selected PSK. */
static VOID ecdhe_psk_client_key_exchange_test(VOID)
{
NX_SECURE_TLS_CREDENTIALS credentials;
NX_SECURE_TLS_KEY_MATERIAL key_material;
NX_SECURE_TLS_ECDHE_HANDSHAKE_DATA *ecdhe_data;
NX_SECURE_TLS_ECC ecc;
const NX_CRYPTO_METHOD *curves[1];
USHORT groups[1];
USHORT received_remote_credentials = NX_FALSE;
UCHAR client_key_exchange[2u + sizeof(psk_id_two) - 1u + 1u + sizeof(ecdhe_public_key)];
UINT offset;
UINT status;

    psk_credentials_setup(&credentials);
    memset(&key_material, 0, sizeof(key_material));

    ecdhe_data = (NX_SECURE_TLS_ECDHE_HANDSHAKE_DATA *)key_material.nx_secure_tls_new_key_material_data;
    ecdhe_data -> nx_secure_tls_ecdhe_named_curve = NX_CRYPTO_EC_SECP256R1;
    ecdhe_data -> nx_secure_tls_ecdhe_private_key[0] = 1;
    ecdhe_data -> nx_secure_tls_ecdhe_private_key_length = 1;

    groups[0] = (USHORT)NX_CRYPTO_EC_SECP256R1;
    curves[0] = &crypto_method_test_curve;
    ecc.nx_secure_tls_ecc_supported_groups = groups;
    ecc.nx_secure_tls_ecc_supported_groups_count = 1;
    ecc.nx_secure_tls_ecc_curves = curves;

    offset = 0;
    client_key_exchange[offset++] = 0;
    client_key_exchange[offset++] = (UCHAR)(sizeof(psk_id_two) - 1u);
    NX_SECURE_MEMCPY(&client_key_exchange[offset], psk_id_two, sizeof(psk_id_two) - 1u);
    offset += sizeof(psk_id_two) - 1u;
    client_key_exchange[offset++] = sizeof(ecdhe_public_key);
    NX_SECURE_MEMCPY(&client_key_exchange[offset], ecdhe_public_key, sizeof(ecdhe_public_key));
    offset += sizeof(ecdhe_public_key);

    status = _nx_secure_process_client_key_exchange(&ecdhe_psk_ciphersuite, NX_SECURE_TLS_VERSION_TLS_1_2,
                                                    client_key_exchange, offset,
                                                    &received_remote_credentials, &key_material, &credentials,
                                                    NX_NULL, 0, NX_NULL, 0, &ecc);
    EXPECT_EQ(NX_SECURE_TLS_SUCCESS, status);
    EXPECT_EQ(NX_TRUE, received_remote_credentials);
    EXPECT_EQ((UINT)(2u + sizeof(ecdhe_shared_secret) + 2u + sizeof(psk_two)), key_material.nx_secure_tls_pre_master_secret_size);
    EXPECT_EQ(sizeof(ecdhe_shared_secret), key_material.nx_secure_tls_pre_master_secret[1]);
    EXPECT_EQ(0, NX_SECURE_MEMCMP(&key_material.nx_secure_tls_pre_master_secret[2], ecdhe_shared_secret, sizeof(ecdhe_shared_secret)));
    EXPECT_EQ(sizeof(psk_two), key_material.nx_secure_tls_pre_master_secret[2u + sizeof(ecdhe_shared_secret) + 1u]);
    EXPECT_EQ(0, NX_SECURE_MEMCMP(&key_material.nx_secure_tls_pre_master_secret[2u + sizeof(ecdhe_shared_secret) + 2u], psk_two, sizeof(psk_two)));
}

/* Validate that ECDHE-PSK client messages include the PSK identity before the EC public key. */
static VOID ecdhe_psk_client_key_exchange_generate_test(VOID)
{
NX_SECURE_TLS_CREDENTIALS credentials;
NX_SECURE_TLS_KEY_MATERIAL key_material;
UCHAR output[32];
ULONG output_size = 0;
UINT status;

    memset(&credentials, 0, sizeof(credentials));
    memset(&key_material, 0, sizeof(key_material));
    NX_SECURE_MEMCPY(credentials.nx_secure_tls_client_psk.nx_secure_tls_psk_id, psk_id_two, sizeof(psk_id_two) - 1u);
    credentials.nx_secure_tls_client_psk.nx_secure_tls_psk_id_size = sizeof(psk_id_two) - 1u;
    key_material.nx_secure_tls_new_key_material_data[0] = sizeof(ecdhe_public_key);
    NX_SECURE_MEMCPY(&key_material.nx_secure_tls_new_key_material_data[1], ecdhe_public_key, sizeof(ecdhe_public_key));

    status = _nx_secure_generate_client_key_exchange(&ecdhe_psk_ciphersuite, &key_material, &credentials,
                                                     output, sizeof(output), &output_size,
                                                     NX_NULL, 0, NX_NULL, 0);
    EXPECT_EQ(NX_SECURE_TLS_SUCCESS, status);
    EXPECT_EQ((ULONG)(2u + sizeof(psk_id_two) - 1u + 1u + sizeof(ecdhe_public_key)), output_size);
    EXPECT_EQ(0, output[0]);
    EXPECT_EQ(sizeof(psk_id_two) - 1u, output[1]);
    EXPECT_EQ(0, NX_SECURE_MEMCMP(&output[2], psk_id_two, sizeof(psk_id_two) - 1u));
    EXPECT_EQ(sizeof(ecdhe_public_key), output[2u + sizeof(psk_id_two) - 1u]);
    EXPECT_EQ(0, NX_SECURE_MEMCMP(&output[2u + sizeof(psk_id_two)], ecdhe_public_key, sizeof(ecdhe_public_key)));
}
#endif /* NX_SECURE_DISABLE_X509 */
#endif /* NX_SECURE_ENABLE_ECC_CIPHERSUITE */
#endif /* NX_SECURE_ENABLE_PSK_CIPHERSUITES */

#ifdef CTEST
void test_application_define(void *first_unused_memory);
void test_application_define(void *first_unused_memory)
#else
void nx_secure_tls_psk_identity_test_application_define(void *first_unused_memory)
#endif
{

    NX_PARAMETER_NOT_USED(first_unused_memory);

    printf("NetX Secure Test:   TLS PSK Identity Test...............................");

#ifdef NX_SECURE_ENABLE_PSK_CIPHERSUITES
    nx_system_initialize();
    nx_secure_tls_initialize();
    plain_psk_identity_test();
    plain_psk_client_key_exchange_length_test();
#if defined(NX_SECURE_ENABLE_ECC_CIPHERSUITE) && !defined(NX_SECURE_DISABLE_X509)
    ecdhe_psk_client_key_exchange_test();
    ecdhe_psk_client_key_exchange_generate_test();
#endif
    printf("SUCCESS!\n");
#else
    printf("N/A\n");
#endif

    test_control_return(0);
}