/* SPDX-License-Identifier: GPL-3.0-only */
/* Versioned NVS storage and validation for the SSH ECDSA P-256 host key. */

#include "ssh_security.h"

#include <limits.h>
#include <stdbool.h>
#include <stddef.h>
#include <string.h>

#include "freertos/FreeRTOS.h"
#include "freertos/semphr.h"
#include "freertos/task.h"
#include "mbedtls/ecp.h"
#include "mbedtls/pk.h"
#include "mbedtls/sha256.h"
#include "nvs.h"
#include "secure_random.h"

#define SSH_SECURITY_SCHEMA_VERSION 1U
#define SSH_SECURITY_BLOB_SIZE 312U
#define SSH_PUBLIC_POINT_LENGTH 65U
#define SSH_PUBLIC_BLOB_LENGTH                                                   \
    (4U + (sizeof(SSH_SECURITY_KEY_TYPE) - 1U) +                                \
     4U + (sizeof(SSH_SECURITY_CURVE_NAME) - 1U) +                              \
     4U + SSH_PUBLIC_POINT_LENGTH)

typedef struct {
    uint32_t schema_version;
    uint16_t blob_size;
    uint16_t private_key_length;
    uint32_t generation;
    uint8_t private_key_der[SSH_SECURITY_PRIVATE_KEY_DER_CAPACITY];
    uint8_t sha256_fingerprint[SSH_SECURITY_SHA256_LENGTH];
    uint8_t reserved[12];
} ssh_security_blob_t;

_Static_assert(offsetof(ssh_security_blob_t, private_key_der) == 12U,
               "SSH security schema offsets changed");
_Static_assert(offsetof(ssh_security_blob_t, sha256_fingerprint) == 268U,
               "SSH fingerprint offset changed");
_Static_assert(sizeof(ssh_security_blob_t) == SSH_SECURITY_BLOB_SIZE,
               "SSH security schema size changed");

static SemaphoreHandle_t s_security_mutex;
static portMUX_TYPE s_mutex_init_lock = portMUX_INITIALIZER_UNLOCKED;
static bool s_mutex_creating;
static ssh_security_blob_t s_material;
static bool s_material_ready;
static ssh_security_load_result_t s_load_result;
static uint32_t s_identity_token, s_next_identity_token;
static TaskHandle_t s_identity_owner;
static bool s_identity_used;

static bool bytes_are_zero(const uint8_t *data, size_t size)
{
    for (size_t index = 0U; index < size; ++index) {
        if (data[index] != 0U) {
            return false;
        }
    }
    return true;
}

static bool constant_time_equal(const uint8_t *left, const uint8_t *right,
                                size_t size)
{
    uint8_t difference = 0U;
    for (size_t index = 0U; index < size; ++index) {
        difference |= left[index] ^ right[index];
    }
    return difference == 0U;
}

static esp_err_t ensure_mutex(void)
{
    for (;;) {
        bool create = false;
        taskENTER_CRITICAL(&s_mutex_init_lock);
        if (s_security_mutex != NULL) {
            taskEXIT_CRITICAL(&s_mutex_init_lock);
            return ESP_OK;
        }
        if (!s_mutex_creating) {
            s_mutex_creating = true;
            create = true;
        }
        taskEXIT_CRITICAL(&s_mutex_init_lock);

        if (create) {
            SemaphoreHandle_t mutex = xSemaphoreCreateMutex();
            taskENTER_CRITICAL(&s_mutex_init_lock);
            s_security_mutex = mutex;
            s_mutex_creating = false;
            taskEXIT_CRITICAL(&s_mutex_init_lock);
            return mutex != NULL ? ESP_OK : ESP_ERR_NO_MEM;
        }
        vTaskDelay(1U);
    }
}

static void write_u32_be(uint8_t output[4], uint32_t value)
{
    output[0] = (uint8_t)(value >> 24U);
    output[1] = (uint8_t)(value >> 16U);
    output[2] = (uint8_t)(value >> 8U);
    output[3] = (uint8_t)value;
}

static size_t append_ssh_string(uint8_t *output, size_t offset,
                                const uint8_t *value, size_t value_length)
{
    write_u32_be(output + offset, (uint32_t)value_length);
    offset += 4U;
    memcpy(output + offset, value, value_length);
    return offset + value_length;
}

static esp_err_t fingerprint_key(const mbedtls_pk_context *key,
                                 uint8_t fingerprint[SSH_SECURITY_SHA256_LENGTH])
{
    const mbedtls_ecp_keypair *ec = mbedtls_pk_ec(*key);
    if (ec == NULL ||
        mbedtls_ecp_keypair_get_group_id(ec) != MBEDTLS_ECP_DP_SECP256R1) {
        return ESP_ERR_INVALID_RESPONSE;
    }

    uint8_t point[SSH_PUBLIC_POINT_LENGTH] = {0};
    size_t point_length = 0U;
    int result = mbedtls_ecp_point_write_binary(
        &ec->MBEDTLS_PRIVATE(grp), &ec->MBEDTLS_PRIVATE(Q),
        MBEDTLS_ECP_PF_UNCOMPRESSED,
        &point_length, point, sizeof(point));
    if (result != 0 || point_length != sizeof(point)) {
        secure_wipe(point, sizeof(point));
        return ESP_ERR_INVALID_RESPONSE;
    }

    uint8_t public_blob[SSH_PUBLIC_BLOB_LENGTH] = {0};
    size_t offset = append_ssh_string(
        public_blob, 0U, (const uint8_t *)SSH_SECURITY_KEY_TYPE,
        sizeof(SSH_SECURITY_KEY_TYPE) - 1U);
    offset = append_ssh_string(
        public_blob, offset, (const uint8_t *)SSH_SECURITY_CURVE_NAME,
        sizeof(SSH_SECURITY_CURVE_NAME) - 1U);
    offset = append_ssh_string(public_blob, offset, point, point_length);

    result = offset == sizeof(public_blob)
                 ? mbedtls_sha256(public_blob, offset, fingerprint, 0)
                 : -1;
    secure_wipe(public_blob, sizeof(public_blob));
    secure_wipe(point, sizeof(point));
    return result == 0 ? ESP_OK : ESP_FAIL;
}

static esp_err_t parse_and_fingerprint(const ssh_security_blob_t *blob,
                                       uint8_t fingerprint[SSH_SECURITY_SHA256_LENGTH])
{
    mbedtls_pk_context key;
    mbedtls_pk_init(&key);
    int result = mbedtls_pk_parse_key(&key,
                                      blob->private_key_der,
                                      blob->private_key_length,
                                      NULL, 0U,
                                      secure_random_mbedtls, NULL);
    esp_err_t error = ESP_ERR_INVALID_RESPONSE;
    if (result == 0 && mbedtls_pk_get_type(&key) == MBEDTLS_PK_ECKEY) {
        const mbedtls_ecp_keypair *ec = mbedtls_pk_ec(key);
        if (ec != NULL &&
            mbedtls_ecp_keypair_get_group_id(ec) == MBEDTLS_ECP_DP_SECP256R1 &&
            mbedtls_ecp_check_privkey(&ec->MBEDTLS_PRIVATE(grp),
                                      &ec->MBEDTLS_PRIVATE(d)) == 0 &&
            mbedtls_ecp_check_pubkey(&ec->MBEDTLS_PRIVATE(grp),
                                     &ec->MBEDTLS_PRIVATE(Q)) == 0 &&
            mbedtls_pk_check_pair(&key, &key,
                                  secure_random_mbedtls, NULL) == 0) {
            error = fingerprint_key(&key, fingerprint);
        }
    }
    mbedtls_pk_free(&key);
    return error;
}

static esp_err_t validate_blob(const ssh_security_blob_t *blob)
{
    if (blob == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    if (blob->schema_version != SSH_SECURITY_SCHEMA_VERSION ||
        blob->blob_size != SSH_SECURITY_BLOB_SIZE) {
        return ESP_ERR_INVALID_VERSION;
    }
    if (blob->generation == 0U || blob->private_key_length == 0U ||
        blob->private_key_length > sizeof(blob->private_key_der) ||
        !bytes_are_zero(blob->private_key_der + blob->private_key_length,
                        sizeof(blob->private_key_der) - blob->private_key_length) ||
        !bytes_are_zero(blob->reserved, sizeof(blob->reserved))) {
        return ESP_ERR_INVALID_RESPONSE;
    }

    uint8_t fingerprint[SSH_SECURITY_SHA256_LENGTH] = {0};
    esp_err_t error = parse_and_fingerprint(blob, fingerprint);
    if (error == ESP_OK &&
        !constant_time_equal(fingerprint, blob->sha256_fingerprint,
                             sizeof(fingerprint))) {
        error = ESP_ERR_INVALID_RESPONSE;
    }
    secure_wipe(fingerprint, sizeof(fingerprint));
    return error;
}

static esp_err_t generate_blob(ssh_security_blob_t *blob, uint32_t generation)
{
    memset(blob, 0, sizeof(*blob));
    blob->schema_version = SSH_SECURITY_SCHEMA_VERSION;
    blob->blob_size = SSH_SECURITY_BLOB_SIZE;
    blob->generation = generation;

    mbedtls_pk_context key;
    mbedtls_pk_init(&key);
    int result = mbedtls_pk_setup(&key, mbedtls_pk_info_from_type(MBEDTLS_PK_ECKEY));
    esp_err_t error = result == 0 ? ESP_OK : ESP_ERR_NO_MEM;
    if (error == ESP_OK) {
        result = mbedtls_ecp_gen_key(MBEDTLS_ECP_DP_SECP256R1,
                                     mbedtls_pk_ec(key),
                                     secure_random_mbedtls, NULL);
        error = result == 0 ? ESP_OK : ESP_FAIL;
    }
    if (error == ESP_OK) {
        result = mbedtls_pk_write_key_der(&key, blob->private_key_der,
                                          sizeof(blob->private_key_der));
        if (result <= 0 || (size_t)result > sizeof(blob->private_key_der)) {
            error = ESP_FAIL;
        } else {
            size_t length = (size_t)result;
            memmove(blob->private_key_der,
                    blob->private_key_der + sizeof(blob->private_key_der) - length,
                    length);
            memset(blob->private_key_der + length, 0,
                   sizeof(blob->private_key_der) - length);
            blob->private_key_length = (uint16_t)length;
        }
    }
    if (error == ESP_OK) {
        error = fingerprint_key(&key, blob->sha256_fingerprint);
    }
    mbedtls_pk_free(&key);
    if (error == ESP_OK) {
        error = validate_blob(blob);
    }
    return error;
}

static esp_err_t save_blob(const ssh_security_blob_t *blob)
{
    esp_err_t error = validate_blob(blob);
    if (error != ESP_OK) {
        return error;
    }

    nvs_handle_t handle;
    error = nvs_open(SSH_SECURITY_NVS_NAMESPACE, NVS_READWRITE, &handle);
    if (error != ESP_OK) {
        return error;
    }
    error = nvs_set_blob(handle, SSH_SECURITY_NVS_BLOB_KEY, blob, sizeof(*blob));
    if (error == ESP_OK) {
        error = nvs_commit(handle);
    }
    nvs_close(handle);
    return error;
}

static esp_err_t load_blob(ssh_security_blob_t *blob, bool *missing)
{
    *missing = false;
    nvs_handle_t handle;
    esp_err_t error = nvs_open(SSH_SECURITY_NVS_NAMESPACE, NVS_READONLY, &handle);
    if (error == ESP_ERR_NVS_NOT_FOUND) {
        *missing = true;
        return ESP_OK;
    }
    if (error != ESP_OK) {
        return error;
    }

    size_t size = 0U;
    error = nvs_get_blob(handle, SSH_SECURITY_NVS_BLOB_KEY, NULL, &size);
    if (error == ESP_ERR_NVS_NOT_FOUND) {
        *missing = true;
        nvs_close(handle);
        return ESP_OK;
    }
    if (error == ESP_ERR_NVS_TYPE_MISMATCH) {
        nvs_close(handle);
        return ESP_ERR_INVALID_RESPONSE;
    }
    if (error != ESP_OK) {
        nvs_close(handle);
        return error;
    }
    if (size != sizeof(*blob)) {
        nvs_close(handle);
        return ESP_ERR_INVALID_VERSION;
    }

    memset(blob, 0, sizeof(*blob));
    error = nvs_get_blob(handle, SSH_SECURITY_NVS_BLOB_KEY, blob, &size);
    nvs_close(handle);
    if (error == ESP_ERR_NVS_INVALID_LENGTH) {
        return ESP_ERR_INVALID_VERSION;
    }
    return error == ESP_OK ? validate_blob(blob) : error;
}

static void install_blob(const ssh_security_blob_t *candidate)
{
    secure_wipe(&s_material, sizeof(s_material));
    s_material = *candidate;
    s_material_ready = true;
    s_load_result = SSH_SECURITY_LOAD_STORED;
}

esp_err_t ssh_security_init(ssh_security_load_result_t *load_result)
{
    esp_err_t error = secure_random_init();
    if (error != ESP_OK) {
        return error;
    }
    error = ensure_mutex();
    if (error != ESP_OK) {
        return error;
    }

    xSemaphoreTake(s_security_mutex, portMAX_DELAY);
    if (s_material_ready) {
        if (load_result != NULL) {
            *load_result = s_load_result;
        }
        xSemaphoreGive(s_security_mutex);
        return ESP_OK;
    }

    if (s_identity_token) {
        xSemaphoreGive(s_security_mutex);
        return ESP_ERR_INVALID_STATE;
    }
    ssh_security_blob_t candidate;
    bool missing = false;
    error = load_blob(&candidate, &missing);
    if (error == ESP_OK && missing) {
        error = generate_blob(&candidate, 1U);
        if (error == ESP_OK) {
            error = save_blob(&candidate);
        }
    }
    if (error == ESP_OK) {
        s_material = candidate;
        s_material_ready = true;
        s_load_result = missing ? SSH_SECURITY_LOAD_GENERATED_MISSING
                                : SSH_SECURITY_LOAD_STORED;
        if (load_result != NULL) {
            *load_result = s_load_result;
        }
    }
    secure_wipe(&candidate, sizeof(candidate));
    xSemaphoreGive(s_security_mutex);
    return error;
}

esp_err_t ssh_security_copy_private_key(uint8_t *output, size_t capacity,
                                        size_t *output_length)
{
    if (output_length == NULL || (output == NULL && capacity != 0U)) {
        return ESP_ERR_INVALID_ARG;
    }
    if (s_security_mutex == NULL) {
        return ESP_ERR_INVALID_STATE;
    }

    xSemaphoreTake(s_security_mutex, portMAX_DELAY);
    esp_err_t error = ESP_ERR_INVALID_STATE;
    if (s_material_ready) {
        *output_length = s_material.private_key_length;
        if (output == NULL) {
            error = capacity == 0U ? ESP_OK : ESP_ERR_INVALID_ARG;
        } else if (capacity < s_material.private_key_length) {
            error = ESP_ERR_INVALID_SIZE;
        } else {
            memcpy(output, s_material.private_key_der,
                   s_material.private_key_length);
            error = ESP_OK;
        }
    }
    xSemaphoreGive(s_security_mutex);
    return error;
}

esp_err_t ssh_security_get_metadata(ssh_security_metadata_t *metadata)
{
    if (metadata == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    if (s_security_mutex == NULL) {
        return ESP_ERR_INVALID_STATE;
    }

    xSemaphoreTake(s_security_mutex, portMAX_DELAY);
    esp_err_t error = ESP_ERR_INVALID_STATE;
    if (s_material_ready) {
        memset(metadata, 0, sizeof(*metadata));
        metadata->generation = s_material.generation;
        memcpy(metadata->sha256_fingerprint, s_material.sha256_fingerprint,
               sizeof(metadata->sha256_fingerprint));
        error = ESP_OK;
    }
    xSemaphoreGive(s_security_mutex);
    return error;
}

esp_err_t ssh_security_get_identity_snapshot(ssh_security_identity_snapshot_t *snapshot)
{
    if (!snapshot) return ESP_ERR_INVALID_ARG;
    memset(snapshot, 0, sizeof(*snapshot));
    if (!s_security_mutex) return ESP_ERR_INVALID_STATE;
    if (xSemaphoreTake(s_security_mutex, 0U) != pdTRUE) return ESP_ERR_TIMEOUT;
    esp_err_t error = s_material_ready ? ESP_OK : ESP_ERR_INVALID_STATE;
    if (error == ESP_OK) {
        snapshot->metadata.generation = s_material.generation;
        memcpy(snapshot->metadata.sha256_fingerprint, s_material.sha256_fingerprint,
               sizeof(snapshot->metadata.sha256_fingerprint));
        snapshot->busy = s_identity_token != 0 || s_next_identity_token == UINT32_MAX;
    }
    xSemaphoreGive(s_security_mutex);
    return error;
}

esp_err_t ssh_security_reserve_identity(uint32_t generation, bool reset, uint32_t *token)
{
    if (!token || (reset && generation)) return ESP_ERR_INVALID_ARG;
    *token = 0;
    if (reset) {
        esp_err_t error = secure_random_init();
        if (error == ESP_OK) error = ensure_mutex();
        if (error != ESP_OK) return error;
    }
    if (!s_security_mutex) return ESP_ERR_INVALID_STATE;
    if (xSemaphoreTake(s_security_mutex, 0U) != pdTRUE) return ESP_ERR_TIMEOUT;
    if (s_identity_token || s_next_identity_token == UINT32_MAX ||
        (!s_material_ready && !reset) ||
        (s_material_ready && s_material.generation == UINT32_MAX) ||
        (generation && generation != s_material.generation)) {
        xSemaphoreGive(s_security_mutex);
        return ESP_ERR_INVALID_STATE;
    }
    *token = s_identity_token = ++s_next_identity_token;
    s_identity_owner = xTaskGetCurrentTaskHandle();
    s_identity_used = false;
    xSemaphoreGive(s_security_mutex);
    return ESP_OK;
}

esp_err_t ssh_security_replace_reserved(uint32_t token)
{
    if (!s_security_mutex || !token) return ESP_ERR_INVALID_STATE;
    xSemaphoreTake(s_security_mutex, portMAX_DELAY);
    if (s_identity_token != token || s_identity_used ||
        s_identity_owner != xTaskGetCurrentTaskHandle()) {
        xSemaphoreGive(s_security_mutex);
        return ESP_ERR_INVALID_STATE;
    }
    s_identity_used = true;
    uint32_t generation = s_material_ready ? s_material.generation + 1U : 1U;
    xSemaphoreGive(s_security_mutex);

    /* Reservation excludes writers while crypto and flash run outside locks. */
    ssh_security_blob_t candidate = {0};
    esp_err_t error = generate_blob(&candidate, generation);
    if (error == ESP_OK) error = save_blob(&candidate);
    xSemaphoreTake(s_security_mutex, portMAX_DELAY);
    if (error == ESP_OK) install_blob(&candidate);
    xSemaphoreGive(s_security_mutex);
    secure_wipe(&candidate, sizeof(candidate));
    return error;
}

void ssh_security_release_identity(uint32_t token)
{
    if (!s_security_mutex || !token) return;
    xSemaphoreTake(s_security_mutex, portMAX_DELAY);
    if (s_identity_token == token && s_identity_owner == xTaskGetCurrentTaskHandle()) {
        s_identity_token = 0;
        s_identity_owner = NULL;
    }
    xSemaphoreGive(s_security_mutex);
}

static esp_err_t replace_identity(bool reset)
{
    uint32_t token = 0;
    esp_err_t error = ssh_security_reserve_identity(0, reset, &token);
    if (error == ESP_OK) error = ssh_security_replace_reserved(token);
    ssh_security_release_identity(token);
    return error;
}

esp_err_t ssh_security_rotate(void) { return replace_identity(false); }
esp_err_t ssh_security_reset(void) { return replace_identity(true); }