/**
 * @file test_wireguard_proto.c
 * @brief WireGuard Protocol Integration Test
 * 
 * Tests the complete WireGuard protocol stack:
 * - X25519 key exchange
 * - Key derivation
 * - Packet encryption/decryption
 */

#include "se050_x25519_sw.h"
#include "se050_chacha20_poly1305.h"
#include "se050_blake2s.h"
#include "se050_hmac_blake2s.h"
#include "se050_crypto_utils.h"
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdbool.h>

static int passed = 0;
static int failed = 0;

#define TEST_ASSERT(cond, msg) do { \
    if (cond) { \
        printf("[PASS] %s\n", msg); \
        passed++; \
    } else { \
        printf("[FAIL] %s\n", msg); \
        failed++; \
    } \
} while(0)

/* RFC 7748 Test Vector 1 */
static const uint8_t RFC7748_SCALAR[32] = {
    0xa5,0x46,0xe3,0x6b,0xf0,0x52,0x7c,0x9d,
    0x3b,0x16,0x15,0x4b,0x82,0x46,0x5e,0xdd,
    0x62,0x14,0x4c,0x0a,0xc1,0xfc,0x5a,0x18,
    0x50,0x6a,0x22,0x44,0xba,0x44,0x9a,0xc4
};

static const uint8_t RFC7748_POINT[32] = {
    0xe6,0xdb,0x68,0x67,0x58,0x30,0x30,0xdb,
    0x35,0x94,0xc1,0xa4,0x24,0xb1,0x5f,0x7c,
    0x72,0x66,0x24,0xec,0x26,0xb3,0x35,0x3b,
    0x10,0xa9,0x03,0xa6,0xd0,0xab,0x1c,0x4c
};

static const uint8_t RFC7748_EXPECTED_SS[32] = {
    0xc3,0xda,0x55,0x37,0x9d,0xe9,0xc6,0x90,
    0x8e,0x94,0xea,0x4d,0xf2,0x8d,0x08,0x4f,
    0x32,0xec,0xcf,0x03,0x49,0x1c,0x71,0xf7,
    0x54,0xb4,0x07,0x55,0x77,0xa2,0x85,0x52
};

/* Test X25519 against RFC 7748 */
static void test_x25519_rfc7748(void)
{
    printf("\n--- Test X25519 RFC 7748 ---\n");
    
    uint8_t shared_secret[32];
    int ret = x25519_sw(shared_secret, RFC7748_SCALAR, RFC7748_POINT);
    
    TEST_ASSERT(ret == 0, "X25519 computation returns 0");
    TEST_ASSERT(memcmp(shared_secret, RFC7748_EXPECTED_SS, 32) == 0,
                "X25519 matches RFC 7748 test vector");
}

/* Test ChaCha20-Poly1305 AEAD */
static void test_chacha20_poly1305(void)
{
    printf("\n--- Test ChaCha20-Poly1305 ---\n");
    
    uint8_t key[32] = {0};
    for (int i = 0; i < 32; i++) key[i] = i;
    uint8_t nonce[12] = {0};
    uint8_t aad[16] = {0};
    uint8_t plaintext[32] = "Hello, WireGuard!";
    uint8_t ciphertext[32 + 16];
    uint8_t tag[16];
    uint8_t decrypted[32];
    
    /* Setup context */
    se050_chacha20_poly1305_ctx_t ctx;
    memcpy(ctx.key, key, 32);
    
    /* Encrypt */
    int ret = se050_chacha20_poly1305_encrypt(
        &ctx, nonce, plaintext, 17, aad, 16, ciphertext, tag);
    TEST_ASSERT(ret == 0, "Encryption returns 0");
    
    /* Decrypt */
    ret = se050_chacha20_poly1305_decrypt(
        &ctx, nonce, ciphertext, 17, aad, 16, tag, decrypted);
    TEST_ASSERT(ret == 0, "Decryption returns 0");
    TEST_ASSERT(memcmp(decrypted, plaintext, 17) == 0,
                "Decrypted content matches");
}

/* Test BLAKE2s */
static void test_blake2s(void)
{
    printf("\n--- Test BLAKE2s ---\n");
    
    uint8_t key[32] = {0};
    uint8_t data[32] = "test data";
    uint8_t mac[32];
    
    int ret = se050_hmac_blake2s(mac, key, 32, data, 9);
    TEST_ASSERT(ret == 0, "HMAC-BLAKE2s returns 0");
    
    /* Verify MAC is non-zero */
    uint8_t all_zero = 1;
    for (int i = 0; i < 32; i++) {
        if (mac[i] != 0) all_zero = 0;
    }
    TEST_ASSERT(all_zero == 0, "MAC is non-zero");
}

/* Test key derivation (simplified HKDF) */
static void test_key_derivation(void)
{
    printf("\n--- Test Key Derivation ---\n");
    
    /* Simulate WireGuard key derivation */
    uint8_t shared_secret[32] = {0};
    for (int i = 0; i < 32; i++) shared_secret[i] = i;
    
    uint8_t info[] = "WireGuard v1 zx2c4 IPsec v1";
    uint8_t output[64];
    
    /* Simple expansion: HMAC(shared_secret, info) */
    se050_hmac_blake2s(output, shared_secret, 32, info, sizeof(info) - 1);
    
    uint8_t all_zero = 1;
    for (int i = 0; i < 64; i++) {
        if (output[i] != 0) all_zero = 0;
    }
    TEST_ASSERT(all_zero == 0, "Derived keys are non-zero");
}

/* Test full DH exchange simulation */
static void test_dh_exchange(void)
{
    printf("\n--- Test DH Exchange Simulation ---\n");
    
    /* Use RFC 7748 test vectors for Alice and Bob */
    const uint8_t ALICE_PRIVATE[32] = {
        0x77,0x07,0x6d,0x0a,0x73,0x18,0xa5,0x7d,
        0x3c,0x16,0xc1,0x72,0x51,0xb2,0x66,0x45,
        0xdf,0xbc,0x9f,0x03,0xc8,0xf2,0xbc,0x2b,
        0x4f,0x81,0xb0,0x9c,0x0c,0xcb,0x3e,0x9a
    };
    
    const uint8_t BOB_PRIVATE[32] = {
        0x5d,0xab,0x08,0x7e,0x62,0x4a,0x8a,0x4b,
        0x79,0xe1,0x7f,0x8b,0x83,0x80,0x0e,0xe6,
        0x6f,0x31,0x97,0xb6,0x93,0x44,0x41,0x76,
        0x61,0x6c,0x39,0x5c,0xbe,0x7a,0x4b,0x82
    };
    
    /* Derive public keys */
    uint8_t alice_public[32], bob_public[32];
    se050_x25519_sw_derive_public_key(alice_public, ALICE_PRIVATE);
    se050_x25519_sw_derive_public_key(bob_public, BOB_PRIVATE);
    
    /* Compute shared secrets */
    uint8_t alice_shared[32], bob_shared[32];
    
    int ret1 = x25519_sw(alice_shared, ALICE_PRIVATE, bob_public);
    int ret2 = x25519_sw(bob_shared, BOB_PRIVATE, alice_public);
    
    TEST_ASSERT(ret1 == 0, "Alice computes shared secret");
    TEST_ASSERT(ret2 == 0, "Bob computes shared secret");
    TEST_ASSERT(memcmp(alice_shared, bob_shared, 32) == 0,
                "Shared secrets match");
}

/* Test packet encryption/decryption flow */
static void test_packet_flow(void)
{
    printf("\n--- Test Packet Encryption Flow ---\n");
    
    /* Setup: Generate keypair and derive shared secret */
    se050_x25519_sw_keypair_t keypair;
    se050_x25519_sw_generate_keypair(&keypair, NULL, NULL);
    
    uint8_t peer_public[32];
    se050_x25519_sw_derive_public_key(peer_public, keypair.private_key);
    
    uint8_t shared_secret[32];
    x25519_sw(shared_secret, keypair.private_key, peer_public);
    
    /* Derive session key (simplified) */
    uint8_t session_key[32];
    uint8_t info[] = "WireGuard v1 zx2c4 IPsec v1";
    se050_hmac_blake2s(session_key, shared_secret, 32, info, sizeof(info) - 1);
    
    /* Encrypt packet */
    const char *plaintext = "Test WireGuard packet";
    uint8_t nonce[12] = {0};
    uint8_t header[16] = {1};  /* Packet type */
    
    uint8_t ciphertext[100];
    uint8_t tag[16];
    
    se050_chacha20_poly1305_ctx_t ctx;
    memcpy(ctx.key, session_key, 32);
    
    int ret = se050_chacha20_poly1305_encrypt(
        &ctx, nonce, (uint8_t*)plaintext, strlen(plaintext),
        header, 16, ciphertext, tag);
    TEST_ASSERT(ret == 0, "Packet encryption succeeds");
    
    /* Decrypt packet */
    uint8_t decrypted[100];
    ret = se050_chacha20_poly1305_decrypt(
        &ctx, nonce, ciphertext, strlen(plaintext),
        header, 16, tag, decrypted);
    TEST_ASSERT(ret == 0, "Packet decryption succeeds");
    TEST_ASSERT(memcmp(decrypted, plaintext, strlen(plaintext)) == 0,
                "Decrypted packet matches");
}

/* Test memory zeroizing */
static void test_memory_zeroizing(void)
{
    printf("\n--- Test Memory Zeroizing ---\n");
    
    uint8_t sensitive[32] = {0};
    for (int i = 0; i < 32; i++) sensitive[i] = i;
    
    /* Zeroize */
    memzero_explicit(sensitive, 32);
    
    /* Verify */
    uint8_t all_zero = 1;
    for (int i = 0; i < 32; i++) {
        if (sensitive[i] != 0) all_zero = 0;
    }
    TEST_ASSERT(all_zero == 1, "Memory zeroized correctly");
}

int main(void)
{
    printf("========================================\n");
    printf(" WireGuard Protocol Integration Test\n");
    printf("========================================\n");
    
    test_x25519_rfc7748();
    test_chacha20_poly1305();
    test_blake2s();
    test_key_derivation();
    test_dh_exchange();
    test_packet_flow();
    test_memory_zeroizing();
    
    printf("\n========================================\n");
    printf(" Results: %d passed, %d failed\n", passed, failed);
    printf("========================================\n");
    
    return failed == 0 ? EXIT_SUCCESS : EXIT_FAILURE;
}