/**
 * @file se050_keystore.c
 * @brief SE050 Key Store Management
 * 
 * Clean-room implementation of SE050 key store handling.
 * 
 * License: MIT (Clean-room implementation)
 */

#define _GNU_SOURCE  /* For MADV_DONTDUMP, MADV_WIPEONFORK */
#define _POSIX_C_SOURCE 200809L

#include "se050_wireguard.h"
#include "se050_crypto_utils.h"
#include "se050_keystore_internal.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>

/* Linux memory protection */
#ifdef __linux__
#include <sys/mman.h>
#endif

/* ============================================================================
 * Memory Protection (Linux-specific)
 * ============================================================================ */

/**
 * @brief Protect sensitive memory from being swapped or dumped
 */
#ifdef __linux__
static se050_status_t protect_sensitive_memory(void *ptr, size_t size)
{
    /* 1. Prevent swapping to disk (optional - may fail due to permissions) */
    if (mlock(ptr, size) != 0) {
        /* mlock may fail due to ulimit restrictions. Log warning but continue. */
        fprintf(stderr, "Warning: mlock failed (%s). Memory may be swapped.\n", strerror(errno));
        /* Continue without mlock - better than failing entirely */
    }
    
    /* 2. Exclude from core dumps */
    if (madvise(ptr, size, MADV_DONTDUMP) != 0) {
        perror("madvise MADV_DONTDUMP failed");
        /* Non-fatal, continue */
    }
    
    /* 3. Clear in child processes after fork() */
    if (madvise(ptr, size, MADV_WIPEONFORK) != 0) {
        perror("madvise MADV_WIPEONFORK failed");
        /* Non-fatal, continue */
    }
    
    return SE050_OK;
}

static void release_memory_protection(void *ptr, size_t size)
{
    munlock(ptr, size);
}
#else
static se050_status_t protect_sensitive_memory(void *ptr, size_t size)
{
    (void)ptr;
    (void)size;
    return SE050_OK;
}

static void release_memory_protection(void *ptr, size_t size)
{
    (void)ptr;
    (void)size;
}
#endif

/* ============================================================================
 * Key Store Management
 * ============================================================================ */

se050_status_t se050_keystore_init(se050_keystore_ctx_t **ctx, se050_session_ctx_t *session)
{
    se050_keystore_ctx_t *keystore;
    size_t ctx_size;
    
    if (!ctx || !session) {
        return SE050_ERR_INVALID_ARG;
    }
    
    /* Allocate key store context */
    keystore = (se050_keystore_ctx_t *)calloc(1, sizeof(*keystore));
    if (!keystore) {
        return SE050_ERR_FAIL;
    }
    
    ctx_size = sizeof(*keystore);
    
    /* Apply memory protection (Linux only) */
    if (protect_sensitive_memory(keystore, ctx_size) != SE050_OK) {
        free(keystore);
        return SE050_ERR_FAIL;
    }
    
    keystore->session = session;
    keystore->max_objects = 16;  /* Maximum 16 keys */
    keystore->objects = (key_object_t *)calloc(keystore->max_objects, sizeof(key_object_t));
    
    if (!keystore->objects) {
        release_memory_protection(keystore, ctx_size);
        free(keystore);
        return SE050_ERR_FAIL;
    }
    
    /* Protect key objects array */
    if (protect_sensitive_memory(keystore->objects, keystore->max_objects * sizeof(key_object_t)) != SE050_OK) {
        free(keystore->objects);
        release_memory_protection(keystore, ctx_size);
        free(keystore);
        return SE050_ERR_FAIL;
    }
    
    keystore->num_objects = 0;
    
    *ctx = keystore;
    return SE050_OK;
}

void se050_keystore_free(se050_keystore_ctx_t *ctx)
{
    size_t i;
    
    if (!ctx) {
        return;
    }
    
    /* Securely zeroize all key objects */
    for (i = 0; i < ctx->num_objects; i++) {
        key_object_t *obj = &ctx->objects[i];
        
        if (obj->flags & KEY_FLAG_GENERATED) {
            memzero_explicit(obj->private_key, sizeof(obj->private_key));
            memzero_explicit(obj->public_key, sizeof(obj->public_key));
        }
    }
    
    /* Release memory protection for key objects array */
    if (ctx->objects) {
        release_memory_protection(ctx->objects, ctx->max_objects * sizeof(key_object_t));
        free(ctx->objects);
        ctx->objects = NULL;
    }
    
    ctx->num_objects = 0;
    ctx->max_objects = 0;
    
    /* Release memory protection for context */
    release_memory_protection(ctx, sizeof(*ctx));
    
    /* Free key store context */
    free(ctx);
}

/* ============================================================================
 * Key Object Management
 * ============================================================================ */

/* Internal functions - defined in se050_keystore_internal.h */

key_object_t *find_key_object(se050_keystore_ctx_t *keystore, uint32_t key_id)
{
    size_t i;
    
    for (i = 0; i < keystore->num_objects; i++) {
        if (keystore->objects[i].key_id == key_id) {
            return &keystore->objects[i];
        }
    }
    
    return NULL;
}

key_object_t *allocate_key_object(se050_keystore_ctx_t *keystore)
{
    key_object_t *obj;
    
    if (keystore->num_objects >= keystore->max_objects) {
        return NULL;
    }
    
    obj = &keystore->objects[keystore->num_objects];
    memset(obj, 0, sizeof(*obj));
    
    keystore->num_objects++;
    return obj;
}

se050_status_t se050_keystore_generate_key(se050_keystore_ctx_t *keystore,
                                            uint32_t key_id,
                                            cipher_type_t cipher_type,
                                            size_t key_size,
                                            uint8_t *private_key,
                                            uint8_t *public_key)
{
    key_object_t *obj;
    
    if (!keystore || !private_key || !public_key) {
        return SE050_ERR_INVALID_ARG;
    }
    
    /* Check if key already exists */
    if (find_key_object(keystore, key_id) != NULL) {
        return SE050_ERR_FAIL;
    }
    
    /* Allocate new key object */
    obj = allocate_key_object(keystore);
    if (!obj) {
        return SE050_ERR_FAIL;
    }
    
    /* Initialize key object */
    obj->key_id = key_id;
    obj->key_part = KEY_PART_PAIR;
    obj->cipher_type = cipher_type;
    obj->key_size = key_size;
    obj->flags = KEY_FLAG_GENERATED | KEY_FLAG_PERSISTENT;
    
    /* Copy key data securely */
    secure_memcpy(obj->private_key, private_key, key_size);
    secure_memcpy(obj->public_key, public_key, key_size);
    
    return SE050_OK;
}

se050_status_t se050_keystore_get_public_key(se050_keystore_ctx_t *keystore,
                                              uint32_t key_id,
                                              uint8_t *public_key,
                                              size_t *key_size)
{
    key_object_t *obj;
    
    if (!keystore || !public_key || !key_size) {
        return SE050_ERR_INVALID_ARG;
    }
    
    /* Find key object */
    obj = find_key_object(keystore, key_id);
    if (!obj) {
        return SE050_ERR_FAIL;
    }
    
    /* Check if public key is available */
    if (!(obj->flags & KEY_FLAG_GENERATED)) {
        return SE050_ERR_FAIL;
    }
    
    /* Copy public key securely */
    *key_size = obj->key_size;
    secure_memcpy(public_key, obj->public_key, obj->key_size);
    
    return SE050_OK;
}

se050_status_t se050_keystore_get_private_key(se050_keystore_ctx_t *keystore,
                                               uint32_t key_id,
                                               uint8_t *private_key,
                                               size_t *key_size)
{
    key_object_t *obj;
    
    if (!keystore || !private_key || !key_size) {
        return SE050_ERR_INVALID_ARG;
    }
    
    /* Find key object */
    obj = find_key_object(keystore, key_id);
    if (!obj) {
        return SE050_ERR_FAIL;
    }
    
    /* Check if private key is available */
    if (!(obj->flags & KEY_FLAG_GENERATED)) {
        return SE050_ERR_FAIL;
    }
    
    /* Copy private key securely */
    *key_size = obj->key_size;
    secure_memcpy(private_key, obj->private_key, obj->key_size);
    
    return SE050_OK;
}