/**
 * @file se050_wireguard.c
 * @brief WireGuard VPN Protocol Implementation (Clean-room)
 * 
 * Based on RFC 9153 - WireGuard
 * License: MIT (Clean-room implementation)
 * 
 * WireGuard is a modern, fast, and secure VPN protocol that uses:
 * - X25519 for key exchange
 * - ChaCha20 for encryption
 * - Poly1305 for authentication
 * - BLAKE2s for hashing
 * - TAI64N for timestamping
 */

#include "se050_wireguard.h"
#include "se050_x25519_sw.h"
#include "se050_chacha20_poly1305.h"
#include "se050_blake2s.h"
#include "se050_hmac_blake2s.h"
#include "se050_tai64n.h"
#include "se050_crypto_utils.h"
#include <string.h>
#include <stdint.h>
#include <stdlib.h>

/* =========================================================================
 * WireGuard Protocol Constants
 * ========================================================================= */

#define WG_NONCE_LEN 12
#define WG_MAX_PACKET_SIZE 65535
#define WG_MAC1_SIZE 16
#define WG_MAC2_SIZE 16

/* WireGuard packet types (RFC 9153) */
#define WG_TYPE_HANDSHAKE_INIT  1
#define WG_TYPE_HANDSHAKE_RESP  2
#define WG_TYPE_COOKIE_REPLY    3
#define WG_TYPE_DATA            4

/* Cookie magic */
static const uint8_t WG_COOKIE_MAGIC[16] = {
    0x12, 0x04, 0x21, 0x00, 0x00, 0x00, 0x00, 0x02,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00
};


/* =========================================================================
 * Helper Functions
 * ========================================================================= */

/* Constant-time comparison */
static bool constant_time_eq(const uint8_t *a, const uint8_t *b, size_t len)
{
    volatile uint8_t result = 0;
    for (size_t i = 0; i < len; i++) {
        result |= a[i] ^ b[i];
    }
    return result == 0;
}

/* HKDF for WireGuard - always uses 32-byte PRK
 * T(1) = HMAC(PRK, 0x01)
 * T(2) = HMAC(PRK, T(1) || 0x02)
 * T(3) = HMAC(PRK, T(2) || 0x03)
 */
static void wg_hkdf_2(const uint8_t *prk,
                     uint8_t *out1, uint8_t *out2)
{
    /* T(1) = HMAC(PRK, 0x01) */
    uint8_t c1 = 0x01;
    se050_hmac_blake2s(out1, prk, 32, &c1, 1);
    
    /* T(2) = HMAC(PRK, T(1) || 0x02) */
    uint8_t t2_input[33];
    memcpy(t2_input, out1, 32);
    t2_input[32] = 0x02;
    se050_hmac_blake2s(out2, prk, 32, t2_input, 33);
    
    memzero_explicit(t2_input, 33);
}

/* HKDF-3 (three outputs) - WireGuard style */
static void wg_hkdf_3(const uint8_t *prk,
                     uint8_t *out1, uint8_t *out2, uint8_t *out3)
{
    /* T(1) = HMAC(PRK, 0x01) */
    uint8_t c1 = 0x01;
    se050_hmac_blake2s(out1, prk, 32, &c1, 1);
    
    /* T(2) = HMAC(PRK, T(1) || 0x02) */
    uint8_t t2_input[33];
    memcpy(t2_input, out1, 32);
    t2_input[32] = 0x02;
    se050_hmac_blake2s(out2, prk, 32, t2_input, 33);
    
    /* T(3) = HMAC(PRK, T(2) || 0x03) */
    uint8_t t3_input[33];
    memcpy(t3_input, out2, 32);
    t3_input[32] = 0x03;
    se050_hmac_blake2s(out3, prk, 32, t3_input, 33);
    
    memzero_explicit(t2_input, 33);
    memzero_explicit(t3_input, 33);
}

/* =========================================================================
 * Session Management
 * ========================================================================= */

int se050_wireguard_session_init(se050_wireguard_session_t *session,
                                 const uint8_t *private_key,
                                 const uint8_t *peer_public_key)
{
    if (!session || !private_key || !peer_public_key) {
        return -1;
    }
    
    memset(session, 0, sizeof(*session));
    
    /* Copy keys */
    memcpy(session->private_key, private_key, WG_KEY_LEN);
    memcpy(session->peer_public_key, peer_public_key, WG_KEY_LEN);
    
    /* Derive public key from private key */
    if (se050_x25519_sw_derive_public_key(session->public_key, private_key) < 0) {
        return -1;
    }
    
    /* Initialize chain key with peer public key */
    memcpy(session->chain_key, peer_public_key, WG_KEY_LEN);
    
    /* Initialize cookie state */
    /* WireGuard uses keyed BLAKE2s, not HMAC */
    se050_blake2s_keyed(session->cookie_secret, 32, 
                        WG_COOKIE_MAGIC, sizeof(WG_COOKIE_MAGIC),
                        private_key, WG_KEY_LEN);
    
    return 0;
}

void se050_wireguard_session_cleanup(se050_wireguard_session_t *session)
{
    if (!session) return;
    
    /* Zeroize all sensitive data */
    memzero_explicit(session->private_key, WG_KEY_LEN);
    memzero_explicit(session->public_key, WG_KEY_LEN);
    memzero_explicit(session->peer_public_key, WG_KEY_LEN);
    memzero_explicit(session->handshake_secret, 32);
    memzero_explicit(session->chain_key, 32);
    memzero_explicit(session->sending_key, 32);
    memzero_explicit(session->receiving_key, 32);
    memzero_explicit(session->cookie_secret, 32);
    memzero_explicit(session->last_mac1, WG_MAC1_SIZE);
    
    memset(session, 0, sizeof(*session));
}

/* =========================================================================
 * Handshake (simplified - no full handshake implementation)
 * ========================================================================= */

int se050_wireguard_derive_keys(se050_wireguard_session_t *session,
                               const uint8_t *shared_secret)
{
    if (!session || !shared_secret) {
        return -1;
    }
    
    /* Derive sending and receiving keys using HKDF
     * WireGuard uses simplified HKDF with 32-byte PRK
     * 
     * Key derivation differs for initiator vs responder:
     * - Initiator:  sending = T(1), receiving = T(2)
     * - Responder:  sending = T(2), receiving = T(1)
     * 
     * For now, using initiator mode (can be extended with is_initiator flag)
     */
    wg_hkdf_2(shared_secret, session->sending_key, session->receiving_key);
    
    /* Reset nonces */
    session->sending_nonce = 0;
    session->receiving_nonce = 0;
    
    session->handshake_complete = 1;
    
    return 0;
}

/* =========================================================================
 * Packet Encryption/Decryption
 * ========================================================================= */

int se050_wireguard_encrypt_packet(se050_wireguard_session_t *session,
                                   uint8_t *out, size_t *out_len,
                                   const uint8_t *plaintext, size_t plaintext_len)
{
    if (!session || !out || !out_len || !plaintext) {
        return -1;
    }
    
    if (!session->handshake_complete) {
        return -1;
    }
    
    if (plaintext_len > WG_MAX_PACKET_SIZE) {
        return -1;
    }
    
    /* Construct packet: [header (16)] [encrypted payload + MAC (16)] */
    /* Header: type (4) + reserved (4) + key index (4) + nonce (8) */
    
    uint8_t header[16];
    header[0] = WG_TYPE_DATA;  /* RFC 9153: Data packet */
    memset(header + 1, 0, 3);     /* Reserved */
    memset(header + 4, 0, 4);     /* Key index (not used) */
    
    /* Encode nonce (little-endian) */
    uint64_t nonce = session->sending_nonce;
    header[8] = nonce & 0xff;
    header[9] = (nonce >> 8) & 0xff;
    header[10] = (nonce >> 16) & 0xff;
    header[11] = (nonce >> 24) & 0xff;
    header[12] = (nonce >> 32) & 0xff;
    header[13] = (nonce >> 40) & 0xff;
    header[14] = (nonce >> 48) & 0xff;
    header[15] = (nonce >> 56) & 0xff;
    
    memcpy(out, header, 16);
    
     /* Encrypt payload with ChaCha20-Poly1305
    * Note: We encrypt directly into the output buffer to avoid large stack allocation
    * out = [header(16)][ciphertext][tag(16)]
    */
    uint8_t nonce_buf[WG_NONCE_LEN];
    memset(nonce_buf, 0, 4);
    memcpy(nonce_buf + 4, header + 8, 8);
    
    uint8_t tag[16];
    
    se050_chacha20_poly1305_ctx_t aead_ctx;
    se050_chacha20_poly1305_init(&aead_ctx, session->sending_key);
    
    int ret = se050_chacha20_poly1305_encrypt(
        &aead_ctx,                             /* ctx */
        nonce_buf,                             /* nonce */
        plaintext, plaintext_len,              /* plaintext */
        header, 16,                            /* aad */
        out + 16,                              /* ciphertext (direct write) */
        tag                                    /* tag */
    );
    
    se050_chacha20_poly1305_zeroize(&aead_ctx);
    
    if (ret < 0) {
        memzero_explicit(tag, 16);
        return -1;
    }
    
    memcpy(out + 16 + plaintext_len, tag, 16);
    *out_len = 16 + plaintext_len + 16;
    memzero_explicit(tag, 16);
    
    /* Increment nonce */
    session->sending_nonce++;
    
    return 0;
}

int se050_wireguard_decrypt_packet(se050_wireguard_session_t *session,
                                   uint8_t *plaintext, size_t *plaintext_len,
                                   const uint8_t *packet, size_t packet_len)
{
    if (!session || !plaintext || !plaintext_len || !packet) {
        return -1;
    }
    
    if (!session->handshake_complete) {
        return -1;
    }
    
    if (packet_len < 32) {  /* Minimum: header (16) + ciphertext (0) + tag (16) */
        return -1;
    }
    
    /* Parse header */
    const uint8_t *header = packet;
    uint8_t type = packet[0];
    if (type != WG_TYPE_DATA) {
        return -1;
    }
    
    /* Extract nonce from header */
    uint64_t nonce = 0;
    for (int i = 0; i < 8; i++) {
        nonce |= ((uint64_t)packet[8 + i]) << (8 * i);
    }
    
    /* Check replay - strictly reject nonce <= last received nonce */
    if (session->packets_received > 0 && nonce <= session->receiving_nonce) {
        return -1;  /* Replay detected */
    }
    
    /* Decrypt payload */
    uint8_t nonce_buf[WG_NONCE_LEN];
    memset(nonce_buf, 0, 4);
    memcpy(nonce_buf + 4, packet + 8, 8);
    
    size_t ciphertext_len = packet_len - 16 - 16;  /* Total - header - tag */
    uint8_t tag[16];
    memcpy(tag, packet + 16 + ciphertext_len, 16);
    
    se050_chacha20_poly1305_ctx_t aead_ctx;
    se050_chacha20_poly1305_init(&aead_ctx, session->receiving_key);
    
    int ret = se050_chacha20_poly1305_decrypt(
        &aead_ctx,                               /* ctx */
        nonce_buf,                               /* nonce */
        packet + 16, ciphertext_len,             /* ciphertext */
        header, 16,                              /* aad, aad_len */
        tag,                                     /* tag */
        plaintext                                /* plaintext (output) */
    );
    
    se050_chacha20_poly1305_zeroize(&aead_ctx);
    memzero_explicit(tag, 16);
    
    if (ret < 0) {
        return -1;
    }
    
    /* Update plaintext length and nonce only on success */
    *plaintext_len = ciphertext_len;
    session->receiving_nonce = nonce;
    session->packets_received++;
    
    return 0;
}

/* =========================================================================
 * Cookie Mechanism (DoS Protection)
 * ========================================================================= */

int se050_wireguard_compute_mac1(se050_wireguard_session_t *session,
                                 const uint8_t *packet, size_t packet_len,
                                 uint8_t *mac1)
{
    if (!session || !packet || !mac1) {
        return -1;
    }
    
    /* WireGuard uses keyed BLAKE2s for MAC1 */
    return se050_blake2s_keyed(mac1, 16, session->peer_public_key, WG_KEY_LEN,
                               packet, packet_len);
}

int se050_wireguard_compute_mac2(se050_wireguard_session_t *session,
                                 const uint8_t *mac1,
                                 const uint8_t *packet, size_t packet_len,
                                 uint8_t *mac2)
{
    if (!session || !mac1 || !packet || !mac2) {
        return -1;
    }
    
    /* MAC2 is only used during handshake (packets < 148 bytes)
     * Fixed buffer is sufficient and avoids malloc dependency
     * This is safe for u-boot and other embedded environments
     */
    uint8_t data[256];  /* Handshake packets are typically < 148 bytes */
    
    if (packet_len + WG_MAC1_SIZE > sizeof(data)) {
        return -1;  /* Should never happen for valid handshake packets */
    }
    
    memcpy(data, packet, packet_len);
    memcpy(data + packet_len, mac1, WG_MAC1_SIZE);
    
    /* WireGuard uses keyed BLAKE2s for MAC2 */
    int ret = se050_blake2s_keyed(mac2, 16, session->cookie_secret, 32,
                                  data, packet_len + WG_MAC1_SIZE);
    
    memzero_explicit(data, sizeof(data));
    return ret;
}

/* =========================================================================
 * Key Generation Utility
 * ========================================================================= */

/* Simple test RNG for development (NOT SECURE!) */
#ifdef X25519_SW_TEST
static int simple_rng(uint8_t *out, size_t len, void *ctx)
{
    static uint32_t seed = 12345;
    for (size_t i = 0; i < len; i++) {
        seed = seed * 1103515245 + 12345;
        out[i] = (seed >> 16) & 0xff;
    }
    return 0;
}
#endif

/* System RNG fallback (uses /dev/urandom on POSIX) */
#include <stdio.h>
#include <fcntl.h>
#include <unistd.h>

static int system_rng(uint8_t *out, size_t len, void *ctx)
{
    int fd = open("/dev/urandom", O_RDONLY);
    if (fd < 0) {
        return -1;
    }
    
    size_t total = 0;
    while (total < len) {
        ssize_t n = read(fd, out + total, len - total);
        if (n < 0) {
            close(fd);
            return -1;
        }
        total += n;
    }
    
    close(fd);
    return 0;
}

int se050_wireguard_generate_keypair(uint8_t *private_key, uint8_t *public_key)
{
    if (!private_key || !public_key) {
        return -1;
    }
    
    se050_x25519_sw_keypair_t keypair;
    
#ifdef X25519_SW_TEST
    /* Use simple RNG for testing */
    if (se050_x25519_sw_generate_keypair(&keypair, simple_rng, NULL) < 0) {
        return -1;
    }
#else
    /* Production: use system RNG (can be replaced with SE050 RNG) */
    if (se050_x25519_sw_generate_keypair(&keypair, system_rng, NULL) < 0) {
        return -1;
    }
#endif
    
    memcpy(private_key, keypair.private_key, WG_KEY_LEN);
    memcpy(public_key, keypair.public_key, WG_KEY_LEN);
    
    return 0;
}