/* SPDX-License-Identifier: GPL-3.0-only */
/* UART0 SSH lifecycle, sessions, counters, and host-key recovery commands. */

#include "ssh_console.h"

#include <errno.h>
#include <inttypes.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "esp_console.h"
#include "mbedtls/base64.h"
#include "secure_random.h"
#include "ssh_security.h"
#include "ssh_transport.h"
#include "web_security.h"

static void print_usage(void)
{
    printf("Usage:\n");
    printf("  ssh status|start|stop|sessions\n");
    printf("  ssh disconnect <session-id>\n");
    printf("  ssh counters|clear-counters\n");
    printf("  ssh host-key info\n");
    printf("  ssh host-key rotate --force\n");
    printf("  ssh reset --force\n");
}

static const char *state_name(ssh_transport_session_state_t state)
{
    switch (state) {
        case SSH_TRANSPORT_SESSION_FREE:
            return "free";
        case SSH_TRANSPORT_SESSION_HANDSHAKE:
            return "handshake";
        case SSH_TRANSPORT_SESSION_ACTIVE:
            return "active";
        case SSH_TRANSPORT_SESSION_CLOSING:
            return "closing";
        default:
            return "unknown";
    }
}

static int print_sessions(const ssh_transport_snapshot_t *snapshot)
{
    printf("SSH sessions: active=%" PRIu32 "/%u\n",
           snapshot->active_sessions, SSH_TRANSPORT_MAX_SESSIONS);
    for (size_t index = 0U; index < SSH_TRANSPORT_MAX_SESSIONS; ++index) {
        const ssh_transport_session_snapshot_t *session = &snapshot->sessions[index];
        if (!session->active) {
            continue;
        }
        printf("  id=%" PRIu32 " slot=%u peer=%s state=%s auth=%s broker=%" PRIu32
               " role=%s rx-pending=%s tx-pending=%s closing=%s\n",
               session->session_id, (unsigned int)index, session->peer,
               state_name(session->state), session->authenticated ? "yes" : "no",
               session->broker_client_id,
               session->broker_client_id == SESSION_BROKER_NO_CLIENT
                   ? "unattached"
                   : (session->writer ? "writer" : "observer"),
               session->rx_pending ? "yes" : "no",
               session->tx_pending ? "yes" : "no",
               session->close_requested ? "yes" : "no");
    }
    return 0;
}

static int show_status(bool sessions_only)
{
    ssh_transport_snapshot_t snapshot;
    esp_err_t error = ssh_transport_get_snapshot(&snapshot);
    if (error != ESP_OK) {
        printf("SSH runtime unavailable: %s\n", esp_err_to_name(error));
        return 1;
    }
    if (!sessions_only) {
        char username[WEB_SECURITY_USERNAME_CAPACITY + 1U] = {0};
        size_t username_length = 0U;
        esp_err_t username_error = web_security_copy_username(
            username, sizeof(username), &username_length);
        printf("SSH: initialized=%s running=%s transitioning=%s port=%u last-error=%s\n",
               snapshot.initialized ? "yes" : "no",
               snapshot.running ? "yes" : "no",
               snapshot.transitioning ? "yes" : "no",
               (unsigned int)snapshot.port,
               esp_err_to_name(snapshot.last_error));
        if (username_error == ESP_OK) {
            printf("Authentication: SSH password, username=%.*s, shared with HTTPS\n",
                   (int)username_length, username);
        } else {
            printf("Administrative credentials unavailable: %s\n",
                   esp_err_to_name(username_error));
        }
        printf("Admission: shell/PTY only; exec, subsystem, forwarding, SCP, and SFTP disabled\n");
        printf("Owner task: core=%" PRId32 " stack=%" PRIu32
               " minimum-free=%" PRIu32 " bytes\n",
               snapshot.task_core_id, snapshot.task_stack_size,
               snapshot.task_stack_free_minimum);
    }
    return print_sessions(&snapshot);
}

static int show_counters(void)
{
    ssh_transport_snapshot_t snapshot;
    esp_err_t error = ssh_transport_get_snapshot(&snapshot);
    if (error != ESP_OK) {
        printf("Could not read SSH counters: %s\n", esp_err_to_name(error));
        return 1;
    }
    const ssh_transport_counters_t *counter = &snapshot.counters;
    printf("Lifecycle: starts=%" PRIu64 " start-failures=%" PRIu64
           " stops=%" PRIu64 " tcp-connect=%" PRIu64
           " capacity-reject=%" PRIu64 "\n",
           counter->starts, counter->start_failures, counter->stops,
           counter->tcp_connections, counter->capacity_rejections);
    printf("Handshake: success=%" PRIu64 " failures=%" PRIu64
           " timeouts=%" PRIu64 " auth-attempts=%" PRIu64
           " auth-failures=%" PRIu64 " request-rejects=%" PRIu64 "\n",
           counter->handshake_successes, counter->handshake_failures,
           counter->handshake_timeouts, counter->authentication_attempts,
           counter->authentication_failures, counter->request_rejections);
    printf("Broker: connect=%" PRIu64 " failures=%" PRIu64
           " disconnect=%" PRIu64 " writer-requests=%" PRIu64
           " grants=%" PRIu64 " denials=%" PRIu64
           " revocations=%" PRIu64 "\n",
           counter->broker_connections, counter->broker_failures,
           counter->disconnections, counter->writer_requests,
           counter->writer_grants, counter->writer_denials,
           counter->writer_revocations);
    printf("Stream: rx=%" PRIu64 " accepted=%" PRIu64
           " rejected=%" PRIu64 " tx=%" PRIu64
           " io-failures=%" PRIu64 " session-revocations=%" PRIu64 "\n",
           counter->rx_bytes, counter->rx_accepted_bytes,
           counter->rx_rejected_bytes, counter->tx_bytes,
           counter->io_failures, counter->session_revocations);
    return 0;
}

static int show_host_key(void)
{
    ssh_security_metadata_t metadata;
    esp_err_t error = ssh_security_get_metadata(&metadata);
    if (error != ESP_OK) {
        printf("Could not read SSH host-key information: %s\n", esp_err_to_name(error));
        printf("Use 'ssh reset --force' to replace incompatible or corrupt material.\n");
        return 1;
    }

    unsigned char encoded[48] = {0};
    size_t encoded_length = 0U;
    int result = mbedtls_base64_encode(encoded, sizeof(encoded), &encoded_length,
                                       metadata.sha256_fingerprint,
                                       sizeof(metadata.sha256_fingerprint));
    if (result != 0 || encoded_length >= sizeof(encoded)) {
        secure_wipe(encoded, sizeof(encoded));
        printf("Could not encode SSH host-key fingerprint.\n");
        return 1;
    }
    while (encoded_length > 0U && encoded[encoded_length - 1U] == '=') {
        --encoded_length;
    }
    encoded[encoded_length] = '\0';
    printf("SSH host key: generation=%" PRIu32 " type=%s curve=%s\n",
           metadata.generation, SSH_SECURITY_KEY_TYPE, SSH_SECURITY_CURVE_NAME);
    printf("OpenSSH SHA-256 fingerprint: SHA256:%s\n", encoded);
    secure_wipe(encoded, sizeof(encoded));
    return 0;
}

static bool parse_session_id(const char *text, uint32_t *session_id)
{
    if (text == NULL || text[0] == '\0' || session_id == NULL) {
        return false;
    }
    errno = 0;
    char *end = NULL;
    unsigned long value = strtoul(text, &end, 10);
    if (errno != 0 || end == text || *end != '\0' || value == 0UL ||
        value > UINT32_MAX) {
        return false;
    }
    *session_id = (uint32_t)value;
    return true;
}

static int replace_host_key(bool reset)
{
    ssh_security_metadata_t before = {0};
    bool had_before = ssh_security_get_metadata(&before) == ESP_OK;
    esp_err_t error = ssh_transport_replace_host_key(reset);

    ssh_security_metadata_t after = {0};
    bool have_after = ssh_security_get_metadata(&after) == ESP_OK;
    bool replaced = have_after && (!had_before || after.generation != before.generation);
    if (error != ESP_OK) {
        if (replaced) {
            printf("SSH host key was persisted, but the transport could not complete its restart: %s\n",
                   esp_err_to_name(error));
        } else {
            printf("Could not %s SSH host key: %s\n",
                   reset ? "reset" : "rotate", esp_err_to_name(error));
        }
        return 1;
    }

    printf("SSH host key replaced and persisted; existing clients must verify the new fingerprint.\n");
    return show_host_key();
}

static bool force_is_present(int argc, char **argv, int expected_argc)
{
    return argc == expected_argc && strcmp(argv[expected_argc - 1], "--force") == 0;
}

static int command_ssh(int argc, char **argv)
{
    if (argc == 1 || (argc == 2 && strcmp(argv[1], "help") == 0)) {
        print_usage();
        return 0;
    }
    if (argc == 2 && strcmp(argv[1], "status") == 0) {
        return show_status(false);
    }
    if (argc == 2 && strcmp(argv[1], "sessions") == 0) {
        return show_status(true);
    }
    if (argc == 2 && strcmp(argv[1], "start") == 0) {
        esp_err_t error = ssh_transport_start();
        if (error != ESP_OK) {
            printf("Could not start SSH: %s\n", esp_err_to_name(error));
            return 1;
        }
        printf("SSH started on TCP port %u.\n", SSH_TRANSPORT_PORT);
        return 0;
    }
    if (argc == 2 && strcmp(argv[1], "stop") == 0) {
        esp_err_t error = ssh_transport_stop();
        if (error != ESP_OK) {
            printf("Could not stop SSH: %s\n", esp_err_to_name(error));
            return 1;
        }
        printf("SSH stopped.\n");
        return 0;
    }
    if (argc == 2 && strcmp(argv[1], "counters") == 0) {
        return show_counters();
    }
    if (argc == 2 && strcmp(argv[1], "clear-counters") == 0) {
        esp_err_t error = ssh_transport_clear_counters();
        if (error != ESP_OK) {
            printf("Could not clear SSH counters: %s\n", esp_err_to_name(error));
            return 1;
        }
        printf("SSH counters cleared.\n");
        return 0;
    }
    if (argc == 3 && strcmp(argv[1], "disconnect") == 0) {
        uint32_t session_id = 0U;
        if (!parse_session_id(argv[2], &session_id)) {
            printf("Session ID must be a nonzero decimal integer.\n");
            return 1;
        }
        esp_err_t error = ssh_transport_disconnect(session_id);
        if (error != ESP_OK) {
            printf("Could not disconnect SSH session: %s\n", esp_err_to_name(error));
            return 1;
        }
        printf("SSH session %" PRIu32 " scheduled for disconnect.\n", session_id);
        return 0;
    }
    if (argc == 3 && strcmp(argv[1], "host-key") == 0 &&
        strcmp(argv[2], "info") == 0) {
        return show_host_key();
    }
    if (argc >= 3 && strcmp(argv[1], "host-key") == 0 &&
        strcmp(argv[2], "rotate") == 0) {
        if (!force_is_present(argc, argv, 4)) {
            printf("Host-key rotation requires: ssh host-key rotate --force\n");
            return 1;
        }
        return replace_host_key(false);
    }
    if (strcmp(argv[1], "reset") == 0) {
        if (!force_is_present(argc, argv, 3)) {
            printf("Host-key recovery requires: ssh reset --force\n");
            return 1;
        }
        return replace_host_key(true);
    }

    print_usage();
    return 1;
}

esp_err_t ssh_console_register_commands(void)
{
    const esp_console_cmd_t command = {
        .command = "ssh",
        .help = "Manage authenticated SSH serial transport and host identity",
        .hint = NULL,
        .func = &command_ssh,
        .argtable = NULL,
    };
    return esp_console_cmd_register(&command);
}