/*
 * pkd_daemon.c -- a sample public-key testing daemon using libssh
 *
 * Uses public key authentication to establish an exec channel and
 * echo back payloads to the user.
 *
 * (c) 2014 Jon Simons
 */

#include "config.h"

#include <errno.h>
#include <netinet/in.h>
#include <pthread.h>
#include <signal.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/socket.h>

#include <libssh/callbacks.h>
#include <libssh/libssh.h>
#include <libssh/server.h>
#include <libssh/kex.h>

#include "torture.h" // for ssh_fips_mode()
#include "pkd_daemon.h"

#include <setjmp.h> // for cmocka
#include <cmocka.h>

static int pkdout_enabled;
static int pkderr_enabled;

static void pkdout(const char *fmt, ...) PRINTF_ATTRIBUTE(1, 2);
static void pkderr(const char *fmt, ...) PRINTF_ATTRIBUTE(1, 2);

static void pkdout(const char *fmt, ...) {
    va_list vargs;
    if (pkdout_enabled) {
        va_start(vargs, fmt);
        vfprintf(stdout, fmt, vargs);
        va_end(vargs);
    }
}

static void pkderr(const char *fmt, ...) {
    va_list vargs;
    if (pkderr_enabled) {
        va_start(vargs, fmt);
        vfprintf(stderr, fmt, vargs);
        va_end(vargs);
    }
}

/*
 * pkd state: only one thread can run pkd at a time ---------------------
 */

static struct {
    int rc;
    pthread_t tid;
    int keep_going;
    volatile int pkd_ready;
} ctx;

static struct {
    int server_fd;
    int req_exec_received;
    int close_received;
    int eof_received;
} pkd_state;

static void pkd_sighandler(int signum) {
    (void) signum;
}

static int pkd_init_libssh(void)
{
    int rc = ssh_threads_set_callbacks(ssh_threads_get_pthread());
    return (rc == SSH_OK) ? 0 : 1;
}

static int pkd_init_server_fd(short port) {
    int rc = 0;
    int yes = 1;
    struct sockaddr_in addr;

    int server_fd = socket(PF_INET, SOCK_STREAM, 0);
    if (server_fd < 0) {
        rc = -1;
        goto out;
    }

    rc = setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(int));
    if (rc != 0) {
        goto outclose;
    }

    memset(&addr, 0x0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_port = htons(port);
    addr.sin_addr.s_addr = INADDR_ANY;
    rc = bind(server_fd, (struct sockaddr *)&addr, sizeof(addr));
    if (rc != 0) {
        goto outclose;
    }

    rc = listen(server_fd, 128);
    if (rc == 0) {
        goto out;
    }

outclose:
    close(server_fd);
    server_fd = -1;
out:
    pkd_state.server_fd = server_fd;
    return rc;
}

static int pkd_accept_fd(void)
{
    int fd = -1;
    struct sockaddr_in addr;
    socklen_t len = sizeof(addr);

    do {
        fd = accept(pkd_state.server_fd, (struct sockaddr *) &addr, &len);
    } while ((ctx.keep_going != 0) && (fd < 0) && (errno == EINTR));

    return fd;
}

static void pkd_eof(ssh_session session,
                    ssh_channel channel,
                    void *userdata) {
    (void) session;
    (void) channel;
    (void) userdata;
    pkdout("pkd_eof\n");
    pkd_state.eof_received = 1;
}

static void pkd_chan_close(ssh_session session,
                           ssh_channel channel,
                           void *userdata) {
    (void) session;
    (void) channel;
    (void) userdata;
    pkdout("pkd_chan_close\n");
    pkd_state.close_received = 1;
}

static int pkd_req_exec(ssh_session s,
                        ssh_channel c,
                        const char *cmd,
                        void *userdata) {
    (void) s;
    (void) c;
    (void) cmd;
    (void) userdata;
    /* assumes pubkey authentication has already succeeded */
    pkdout("pkd_req_exec\n");
    pkd_state.req_exec_received = 1;
    return 0;
}

/* assumes there is only ever a single channel */
static struct ssh_channel_callbacks_struct pkd_channel_cb = {
    .channel_eof_function = pkd_eof,
    .channel_close_function = pkd_chan_close,
    .channel_exec_request_function = pkd_req_exec,
};

static int pkd_auth_pubkey_cb(ssh_session s,
                              const char *user,
                              ssh_key key,
                              char state,
                              void *userdata) {
    (void) s;
    (void) user;
    (void) key;
    (void) state;
    (void) userdata;
    pkdout("pkd_auth_pubkey_cb keytype %s, state: %d\n",
           ssh_key_type_to_char(ssh_key_type(key)), state);
    if ((state == SSH_PUBLICKEY_STATE_NONE) ||
        (state == SSH_PUBLICKEY_STATE_VALID)) {
        return SSH_AUTH_SUCCESS;
    }
    return SSH_AUTH_DENIED;
}

static int pkd_service_request_cb(ssh_session session,
                                  const char *service,
                                  void *userdata) {
    (void) session;
    (void) userdata;
    pkdout("pkd_service_request_cb: %s\n", service);
    return (0 == (strcmp(service, "ssh-userauth"))) ? 0 : -1;
}

static ssh_channel pkd_channel_openreq_cb(ssh_session s,
                                          void *userdata) {
    ssh_channel c = NULL;
    ssh_channel *out = (ssh_channel *) userdata;

    /* assumes pubkey authentication has already succeeded */
    pkdout("pkd_channel_openreq_cb\n");

    c = ssh_channel_new(s);
    if (c == NULL) {
        pkderr("ssh_channel_new: %s\n", ssh_get_error(s));
        return NULL;
    }

    ssh_callbacks_init(&pkd_channel_cb);
    pkd_channel_cb.userdata = userdata;
    if (ssh_set_channel_callbacks(c, &pkd_channel_cb) != SSH_OK) {
        pkderr("ssh_set_channel_callbacks: %s\n", ssh_get_error(s));
        ssh_channel_free(c);
        c = NULL;
    }

    *out = c;

    return c;
}

static struct ssh_server_callbacks_struct pkd_server_cb = {
    .auth_pubkey_function = pkd_auth_pubkey_cb,
    .service_request_function = pkd_service_request_cb,
    .channel_open_request_session_function = pkd_channel_openreq_cb,
};

static int pkd_exec_hello(int fd, struct pkd_daemon_args *args)
{
    int rc = -1;
    ssh_bind b = NULL;
    ssh_session s = NULL;
    ssh_event e = NULL;
    ssh_channel c = NULL;
    enum ssh_bind_options_e opts = -1;

    int level = args->opts.libssh_log_level;
    enum pkd_hostkey_type_e type = args->type;
    const char *hostkeypath = args->hostkeypath;
    const char *default_kex = NULL;
    char *all_kex = NULL;
    size_t kex_len = 0;
    const char *all_ciphers = NULL;
    const uint64_t rekey_data_limit = args->rekey_data_limit;
    bool process_config = false;

    pkd_state.eof_received = 0;
    pkd_state.close_received  = 0;
    pkd_state.req_exec_received = 0;

    b = ssh_bind_new();
    if (b == NULL) {
        pkderr("ssh_bind_new\n");
        goto outclose;
    }

    if (type == PKD_RSA) {
        opts = SSH_BIND_OPTIONS_RSAKEY;
    } else if (type == PKD_ED25519) {
        opts = SSH_BIND_OPTIONS_HOSTKEY;
#ifdef HAVE_DSA
    } else if (type == PKD_DSA) {
        opts = SSH_BIND_OPTIONS_DSAKEY;
#endif
    } else if (type == PKD_ECDSA) {
        opts = SSH_BIND_OPTIONS_ECDSAKEY;
    } else {
        pkderr("unknown hostkey type: %d\n", type);
        rc = -1;
        goto outclose;
    }

    rc = ssh_bind_options_set(b, opts, hostkeypath);
    if (rc != 0) {
        pkderr("ssh_bind_options_set: %s\n", ssh_get_error(b));
        goto outclose;
    }

    rc = ssh_bind_options_set(b, SSH_BIND_OPTIONS_LOG_VERBOSITY, &level);
    if (rc != 0) {
        pkderr("ssh_bind_options_set log verbosity: %s\n", ssh_get_error(b));
        goto outclose;
    }

    rc = ssh_bind_options_set(b, SSH_BIND_OPTIONS_PROCESS_CONFIG,
                              &process_config);
    if (rc != 0) {
        pkderr("ssh_bind_options_set process config: %s\n", ssh_get_error(b));
        goto outclose;
    }

    if (!ssh_fips_mode()) {
        /* Add methods not enabled by default */
#define GEX_SHA1 "diffie-hellman-group-exchange-sha1"
        default_kex = ssh_kex_get_default_methods(SSH_KEX);
        kex_len = strlen(default_kex) + strlen(GEX_SHA1) + 2;
        all_kex = malloc(kex_len);
        if (all_kex == NULL) {
            pkderr("Failed to alloc more memory.\n");
            goto outclose;
        }
        snprintf(all_kex, kex_len, "%s," GEX_SHA1, default_kex);
        rc = ssh_bind_options_set(b, SSH_BIND_OPTIONS_KEY_EXCHANGE, all_kex);
        free(all_kex);
        if (rc != 0) {
            pkderr("ssh_bind_options_set kex methods: %s\n", ssh_get_error(b));
            goto outclose;
        }

        /* Enable all supported ciphers */
        all_ciphers = ssh_kex_get_supported_method(SSH_CRYPT_C_S);
        rc = ssh_bind_options_set(b, SSH_BIND_OPTIONS_CIPHERS_C_S, all_ciphers);
        if (rc != 0) {
            pkderr("ssh_bind_options_set Ciphers C-S: %s\n", ssh_get_error(b));
            goto outclose;
        }

        all_ciphers = ssh_kex_get_supported_method(SSH_CRYPT_S_C);
        rc = ssh_bind_options_set(b, SSH_BIND_OPTIONS_CIPHERS_S_C, all_ciphers);
        if (rc != 0) {
            pkderr("ssh_bind_options_set Ciphers S-C: %s\n", ssh_get_error(b));
            goto outclose;
        }
    }

    s = ssh_new();
    if (s == NULL) {
        pkderr("ssh_new\n");
        goto outclose;
    }

    rc = ssh_options_set(s, SSH_OPTIONS_REKEY_DATA, &rekey_data_limit);
    if (rc != 0) {
        pkderr("ssh_options_set rekey data: %s\n", ssh_get_error(s));
        goto outclose;
    }

    /*
     * ssh_bind_accept loads host key as side-effect.  If this
     * succeeds, the given 'fd' will be closed upon 'ssh_free(s)'.
     */
    rc = ssh_bind_accept_fd(b, s, fd);
    if (rc != SSH_OK) {
        pkderr("ssh_bind_accept_fd: %s\n", ssh_get_error(b));
        goto outclose;
    }

    /* accept only publickey-based auth */
    ssh_set_auth_methods(s, SSH_AUTH_METHOD_PUBLICKEY);

    /* initialize callbacks */
    ssh_callbacks_init(&pkd_server_cb);
    pkd_server_cb.userdata = &c;
    rc = ssh_set_server_callbacks(s, &pkd_server_cb);
    if (rc != SSH_OK) {
        pkderr("ssh_set_server_callbacks: %s\n", ssh_get_error(s));
        goto out;
    }

    /* first do key exchange */
    rc = ssh_handle_key_exchange(s);
    if (rc != SSH_OK) {
        pkderr("ssh_handle_key_exchange: %s\n", ssh_get_error(s));
        goto out;
    }

    /* setup and pump event to carry out exec channel */
    e = ssh_event_new();
    if (e == NULL) {
        pkderr("ssh_event_new\n");
        goto out;
    }

    rc = ssh_event_add_session(e, s);
    if (rc != SSH_OK) {
        pkderr("ssh_event_add_session\n");
        goto out;
    }

    /* poll until exec channel established */
    while ((ctx.keep_going != 0) &&
           (rc != SSH_ERROR) && (pkd_state.req_exec_received == 0)) {
        rc = ssh_event_dopoll(e, -1 /* infinite timeout */);
    }

    if (rc == SSH_ERROR) {
        pkderr("ssh_event_dopoll\n");
        goto out;
    } else if (c == NULL) {
        pkderr("poll loop exited but exec channel not ready\n");
        rc = -1;
        goto out;
    }

    rc = ssh_channel_write(c, args->payload.buf, args->payload.len);
    if (rc != (int)args->payload.len) {
        pkderr("ssh_channel_write partial (%d != %zd)\n", rc, args->payload.len);
    }

    rc = ssh_channel_request_send_exit_status(c, 0);
    if (rc != SSH_OK) {
        pkderr("ssh_channel_request_send_exit_status: %s\n",
                        ssh_get_error(s));
        goto out;
    }

    rc = ssh_channel_send_eof(c);
    if (rc != SSH_OK) {
        pkderr("ssh_channel_send_eof: %s\n", ssh_get_error(s));
        goto out;
    }

    rc = ssh_channel_close(c);
    if (rc != SSH_OK) {
        pkderr("ssh_channel_close: %s\n", ssh_get_error(s));
        goto out;
    }

    while ((ctx.keep_going != 0) &&
           (pkd_state.eof_received == 0) &&
           (pkd_state.close_received == 0)) {
        rc = ssh_event_dopoll(e, 1000 /* milliseconds */);
        if (rc == SSH_ERROR) {
            /* log, but don't consider this fatal */
            pkdout("ssh_event_dopoll for eof + close: %s\n", ssh_get_error(s));
            rc = 0;
            break;
        } else {
            rc = 0;
        }
    }

    while ((ctx.keep_going != 0) &&
           (ssh_is_connected(s))) {
        rc = ssh_event_dopoll(e, 1000 /* milliseconds */);
        if (rc == SSH_ERROR) {
            /* log, but don't consider this fatal */
            pkdout("ssh_event_dopoll for session connection: %s\n", ssh_get_error(s));
            rc = 0;
            break;
        } else {
            rc = 0;
        }
    }
    goto out;

outclose:
    close(fd);
out:
    if (c != NULL) {
        ssh_channel_free(c);
    }
    if (e != NULL) {
        ssh_event_remove_session(e, s);
        ssh_event_free(e);
    }
    if (s != NULL) {
        ssh_disconnect(s);
        ssh_free(s);
    }
    if (b != NULL) {
        ssh_bind_free(b);
    }
    return rc;
}

/*
 * main loop ------------------------------------------------------------
 */

static void *pkd_main(void *args) {
    int rc = -1;
    struct pkd_daemon_args *a = (struct pkd_daemon_args *) args;

    struct sigaction act = { .sa_handler = pkd_sighandler, };

    pkd_state.server_fd = -1;
    pkd_state.req_exec_received = 0;
    pkd_state.close_received = 0;
    pkd_state.eof_received = 0;

    /* SIGUSR1 is used to interrupt 'pkd_accept_fd'. */
    rc = sigaction(SIGUSR1, &act, NULL);
    if (rc != 0) {
        pkderr("sigaction: %d\n", rc);
        goto out;
    }

    /* Ignore SIGPIPE */
    signal(SIGPIPE, SIG_IGN);

    rc = pkd_init_libssh();
    if (rc != 0) {
        pkderr("pkd_init_libssh: %d\n", rc);
        goto out;
    }

    rc = pkd_init_server_fd(1234);
    if (rc != 0) {
        pkderr("pkd_init_server_fd: %d\n", rc);
        goto out;
    }

    ctx.pkd_ready = 1;

    while (ctx.keep_going != 0) {
        int fd = pkd_accept_fd();
        if (fd < 0) {
            if (ctx.keep_going != 0) {
                pkderr("pkd_accept_fd");
                rc = -1;
            } else {
                rc = 0;
            }
            break;
        }

        rc = pkd_exec_hello(fd, a);
        if (rc != 0) {
            pkderr("pkd_exec_hello: %d\n", rc);
            break;
        }
    }

    if (pkd_state.server_fd != -1) {
        close(pkd_state.server_fd);
    }
    pkd_state.server_fd = -1;
out:
    ctx.rc = rc;

    return NULL;
}

/*
 * pkd start and stop used by setup/teardown test scaffolding -----------
 */

int pkd_start(struct pkd_daemon_args *args) {
    int rc = 0;

    pkdout_enabled = args->opts.log_stdout;
    pkderr_enabled = args->opts.log_stderr;

    /* Initialize the pkd context. */
    ctx.rc = -1;
    ctx.keep_going = 1;
    ctx.pkd_ready = 0;
    rc = pthread_create(&ctx.tid, NULL, &pkd_main, args);
    assert_int_equal(rc, 0);

    /* Busy-spin until pkd thread is ready. */
    while (ctx.pkd_ready == 0);

    return rc;
}

void pkd_stop(struct pkd_result *out) {
    int rc = 0;

    ctx.keep_going = 0;
    close(pkd_state.server_fd);

    rc = pthread_kill(ctx.tid, SIGUSR1);
    assert_int_equal(rc, 0);

    rc = pthread_join(ctx.tid, NULL);
    assert_int_equal(rc, 0);

    assert_non_null(out);
    out->ok = (ctx.rc == 0);

    return;
}