/* server.c
 *
 * Copyright (C) 2014-2019 wolfSSL Inc.
 *
 * This file is part of wolfSSH.
 *
 * wolfSSH is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3 of the License, or
 * (at your option) any later version.
 *
 * wolfSSH is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with wolfSSH.  If not, see <http://www.gnu.org/licenses/>.
 */
#pragma GCC diagnostic error "-Wall"

#define WOLFSSH_TEST_SERVER
#define WOLFSSH_TEST_THREADING

#include <stdbool.h>
#include "FreeRTOS.h"
#include "task.h"
#include "semphr.h"
#include "wolfssh_test.h"
#include "cli.h"
#include "CLI_Commands.h"
#include "settings_api.h"

#ifdef WOLFSSL_USER_SETTINGS
    #include <wolfssl/wolfcrypt/settings.h>
#else
    #include <wolfssl/options.h>
#endif

#include <wolfssl/wolfcrypt/sha256.h>
#include <wolfssl/wolfcrypt/coding.h>
#include <wolfssh/ssh.h>
//#include <wolfssh/test.h>
#include <wolfssl/wolfcrypt/ecc.h>
#include "examples/server/server.h"

#ifdef NO_FILESYSTEM
    #include <wolfssh/certs_test.h>
#endif


bool fl_reinit_ssh = false; //флаг разрешение переинициализации ssh сервера 

static const char serverBanner[] = "BT-6709 command server\n";


typedef struct {
    WOLFSSH* ssh;
    SOCKET_T fd;
    word32 id;
} thread_ctx_t;

static const bool nonBlock = true;

#ifndef EXAMPLE_HIGHWATER_MARK
    #define EXAMPLE_HIGHWATER_MARK 0x3FFF8000 /* 1GB - 32kB */
#endif
#ifndef EXAMPLE_BUFFER_SZ
    #define EXAMPLE_BUFFER_SZ 256
#endif
#define SCRATCH_BUFFER_SZ 1200

/*
static int dump_stats(thread_ctx_t* ctx)
{
    char stats[1024];
    word32 statsSz;
    word32 txCount, rxCount, seq, peerSeq;

    wolfSSH_GetStats(ctx->ssh, &txCount, &rxCount, &seq, &peerSeq);

    WSNPRINTF(stats, sizeof(stats),
            "Statistics for Thread #%u:\r\n"
            "  txCount = %u\r\n  rxCount = %u\r\n"
            "  seq = %u\r\n  peerSeq = %u\r\n",
            ctx->id, txCount, rxCount, seq, peerSeq);
    statsSz = (word32)strlen(stats);

    fprintf(stderr, "%s", stats);
    return wolfSSH_stream_send(ctx->ssh, (byte*)stats, statsSz);
}
*/


static int NonBlockSSH_accept(WOLFSSH* ssh)
{
    int ret;
    int error;
    SOCKET_T sockfd;
    int select_ret = 0;

    ret = wolfSSH_accept(ssh);
    error = wolfSSH_get_error(ssh);
    sockfd = (SOCKET_T)wolfSSH_get_fd(ssh);

    while (ret != WS_SUCCESS &&
            (error == WS_WANT_READ || error == WS_WANT_WRITE))
    {
#if DEBUG
        if (error == WS_WANT_READ)
            printf("... client would read block\n");
        else if (error == WS_WANT_WRITE)
            printf("... client would write block\n");
#endif

        select_ret = tcp_select(sockfd, 1);
        if (select_ret == WS_SELECT_RECV_READY  ||
            select_ret == WS_SELECT_ERROR_READY ||
            error == WS_WANT_WRITE)
        {
            ret = wolfSSH_accept(ssh);
            error = wolfSSH_get_error(ssh);
        }
        else if (select_ret == WS_SELECT_TIMEOUT)
            error = WS_WANT_READ;
        else
            error = WS_FATAL_ERROR;
    }

    return ret;
}

static void cli_send(intptr_t fd, const char *str, unsigned len)
{
    wolfSSH_stream_send((WOLFSSH *)fd, str, len);
}

static void *server_worker(void* vArgs)
{
    int ret;
    thread_ctx_t* threadCtx = (thread_ctx_t*)vArgs;

    user_level_t user_id;
    wolfSSH_SetUserAuthCtx(threadCtx->ssh, &user_id);

    if (!nonBlock)
        ret = wolfSSH_accept(threadCtx->ssh);
    else
        ret = NonBlockSSH_accept(threadCtx->ssh);
    fl_reinit_ssh = true;
    cli_state_t *cli_state;
    if (!sSettings.sSSH.SSHEnable) {
        // kick the client out
    } else if (ret == WS_SUCCESS) {
        if ((cli_state = alloc_state())) {
            // create the new CLI contex
            cli_state->num_connect = threadCtx->ssh;
            cli_state->input_state = CLI_CMD;
            cli_state->send = cli_send;
            cli_state->user_id = user_id;
            cli_hello(cli_state);

            bool stop = false;
            do {
                uint8_t buf[EXAMPLE_BUFFER_SZ];
                int rxSz = 0;
                do {
                    if (nonBlock) {
                        SOCKET_T sockfd;
                        int select_ret = 0;

                        sockfd = (SOCKET_T)wolfSSH_get_fd(threadCtx->ssh);

                        select_ret = tcp_select(sockfd, 1);
                        if (!sSettings.sSSH.SSHEnable ||
                            (select_ret != WS_SELECT_RECV_READY &&
                            select_ret != WS_SELECT_ERROR_READY &&
                            select_ret != WS_SELECT_TIMEOUT)) {
                            break;
                        }
                    }
                    rxSz = wolfSSH_stream_read(threadCtx->ssh, buf, sizeof(buf));
                    if (rxSz <= 0) {
                        rxSz = wolfSSH_get_error(threadCtx->ssh);
                    }
                } while ((rxSz == WS_WANT_READ || rxSz == WS_WANT_WRITE) && sSettings.sSSH.SSHEnable);
                if (!sSettings.sSSH.SSHEnable) {
                    stop = true;
                    break;
                }

                if (rxSz > 0) {
                    cli_getchar(cli_state, buf[0], true);    // TODO handle rxSz > 1
                    if (buf[0] == 4 || cli_state->state == STATE_CLOSE) {       // EOF
                        stop = 1;
                    }
                } else {
                    stop = 1;
                }
            } while (!stop);
            free_state(cli_state);
        } else {
            cli_send(threadCtx->ssh, pcWarningMessage, pcWarningMessageLen);
        }
    }
    wolfSSH_stream_exit(threadCtx->ssh, 0);
    WCLOSESOCKET(threadCtx->fd);
    wolfSSH_free(threadCtx->ssh);
    free(threadCtx);

    return 0;
}

#ifndef NO_FILESYSTEM
static int load_file(const char* fileName, byte* buf, word32 bufSz)
{
    FILE* file;
    word32 fileSz;
    word32 readSz;

    if (fileName == NULL) return 0;

    if (WFOPEN(&file, fileName, "rb") != 0)
        return 0;
    fseek(file, 0, SEEK_END);
    fileSz = (word32)ftell(file);
    rewind(file);

    if (fileSz > bufSz) {
        fclose(file);
        return 0;
    }

    readSz = (word32)fread(buf, 1, fileSz, file);
    if (readSz < fileSz) {
        fclose(file);
        return 0;
    }

    fclose(file);

    return fileSz;
}
#endif /* !NO_FILESYSTEM */

/* returns buffer size on success */
static int load_key(byte isEcc, byte* buf, word32 bufSz)
{
    word32 sz = 0;

#ifndef NO_FILESYSTEM
    const char* bufName;
    bufName = isEcc ? "./keys/server-key-ecc.der" :
                       "./keys/server-key-rsa.der" ;
    sz = load_file(bufName, buf, bufSz);
#else
    /* using buffers instead */
    if (isEcc) {
        if (sizeof_ecc_key_der_256 > bufSz) {
            return 0;
        }
        WMEMCPY(buf, ecc_key_der_256, sizeof_ecc_key_der_256);
        sz = sizeof_ecc_key_der_256;
    }
    else {
        if (sizeof_rsa_key_der_2048 > bufSz) {
            return 0;
        }
        WMEMCPY(buf, rsa_key_der_2048, sizeof_rsa_key_der_2048);
        sz = sizeof_rsa_key_der_2048;
    }
#endif

    return sz;
}


static INLINE void c32toa(word32 u32, byte* c)
{
    c[0] = (u32 >> 24) & 0xff;
    c[1] = (u32 >> 16) & 0xff;
    c[2] = (u32 >>  8) & 0xff;
    c[3] =  u32 & 0xff;
}

static int wsUserAuth(byte authType, WS_UserAuthData* authData, void* ctx)
{
    user_level_t *user_id = ctx;

    if (ctx == NULL) {
        printf("wsUserAuth: ctx not set");
        return WOLFSSH_USERAUTH_FAILURE;
    }

    if (authType != WOLFSSH_USERAUTH_PASSWORD) {

        return WOLFSSH_USERAUTH_FAILURE;
    }

    // the incoming password is not zero-terminated
    char password[MAX_WEB_PASSWD_LEN];
    strncpy(password, (char *)authData->sf.password.password, sizeof(password));
    password[min(MAX_WEB_PASSWD_LEN - 1, authData->sf.password.passwordSz)] = 0;
    *user_id = cli_auth_user((const char *)authData->username, password, LOG_LOGIN_SSH);
    if (*user_id != MAX_USER_LEVELS) {
        return WOLFSSH_USERAUTH_SUCCESS;
    }

    return WOLFSSH_USERAUTH_INVALID_USER;
}


static void ssh_server(void *arg)
{
    (void)arg;
    WSTARTTCP();
    #ifdef DEBUG_WOLFSSH
        wolfSSH_Debugging_ON();
    #endif
    wolfSSH_Init();
    WOLFSSH_CTX* ctx = NULL;
    static SOCKET_T ssh_listen_fd = 0;
    word32 defaultHighwater = EXAMPLE_HIGHWATER_MARK;
    word32 threadCount = 0;
    const char multipleConnections = 0;
    char useEcc = 1;

    if (wolfSSH_Init() != WS_SUCCESS) {
        printf("Couldn't initialize wolfSSH.\n");
        exit(EXIT_FAILURE);
    }

    ctx = wolfSSH_CTX_new(WOLFSSH_ENDPOINT_SERVER, NULL);
    if (ctx == NULL) {
        printf("Couldn't allocate SSH CTX data.\n");
        exit(EXIT_FAILURE);
    }

    wolfSSH_SetUserAuth(ctx, wsUserAuth);
    wolfSSH_CTX_SetBanner(ctx, TELNET_CLI_WELCOME_MESSAGE);

    {
        byte buf[SCRATCH_BUFFER_SZ];
        word32 bufSz;

        bufSz = load_key(useEcc, buf, SCRATCH_BUFFER_SZ);
        if (bufSz == 0) {
            printf("Couldn't load key.\n");
            exit(EXIT_FAILURE);
        }
        if (wolfSSH_CTX_UsePrivateKey_buffer(ctx, buf, bufSz,
                                             WOLFSSH_FORMAT_ASN1) < 0) {
            printf("Couldn't use key buffer.\n");
            exit(EXIT_FAILURE);
        }
    }

    while (sSettings.sSSH.SSHEnable) {
        fl_reinit_ssh = false;
        word16 port = sSettings.sSSH.port;
        tcp_listen(&ssh_listen_fd, &port, 1, false, false);

        do {
            SOCKET_T      clientFd = 0;
            SOCKADDR_IN_T clientAddr;
            socklen_t     clientAddrSz = sizeof(clientAddr);
#ifndef SINGLE_THREADED
            THREAD_TYPE   thread;
#endif
            WOLFSSH*      ssh;
            thread_ctx_t* threadCtx;

            threadCtx = (thread_ctx_t*)malloc(sizeof(thread_ctx_t));
            if (threadCtx == NULL) {
                printf("Couldn't allocate thread context data.\n");
                exit(EXIT_FAILURE);
            }

            ssh = wolfSSH_new(ctx);
            if (ssh == NULL) {
                printf("Couldn't allocate SSH data.\n");
                exit(EXIT_FAILURE);
            }
            /* Use the session object for its own highwater callback ctx */
            if (defaultHighwater > 0) {
                wolfSSH_SetHighwaterCtx(ssh, (void*)ssh);
                wolfSSH_SetHighwater(ssh, defaultHighwater);
            }

            clientFd = accept(ssh_listen_fd, (struct sockaddr*)&clientAddr,
                                                                     &clientAddrSz);
            if (clientFd == -1)
                err_sys("tcp accept failed");

            if (nonBlock)
                tcp_set_nonblocking(&clientFd);

            wolfSSH_set_fd(ssh, (int)clientFd);

            threadCtx->ssh = ssh;
            threadCtx->fd = clientFd;
            threadCtx->id = threadCount++;

#ifndef SINGLE_THREADED
            ThreadStart(server_worker, threadCtx, &thread);

            if (multipleConnections)
                ThreadDetach(thread);
            else
                ThreadJoin(thread);
#else
            server_worker(threadCtx);
#endif /* SINGLE_THREADED */
        } while (multipleConnections);
        if (ssh_listen_fd) {
            WCLOSESOCKET(ssh_listen_fd);
            ssh_listen_fd = 0;
        }
    }

    wolfSSH_CTX_free(ctx);
    if (wolfSSH_Cleanup() != WS_SUCCESS) {
        printf("Couldn't clean up wolfSSH.\n");
        exit(EXIT_FAILURE);
    }
#if defined(HAVE_ECC) && defined(FP_ECC) && defined(HAVE_THREAD_LS)
    wc_ecc_fp_free();  /* free per thread cache */
#endif
    vTaskDelete(NULL);
}

void ssh_server_init(void)
{
    xTaskCreate(ssh_server, ( char * ) "ssh_server", 24*configMINIMAL_STACK_SIZE + EXAMPLE_BUFFER_SZ, NULL, tskIDLE_PRIORITY + 1, NULL);
}

void ssh_server_restart(void)
{
}