#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <libssh/callbacks.h>
#include <libssh/server.h>

#include "server_ssh.h"
#include "log.h"

#if LIBSSH_VERSION_MAJOR != 0 || LIBSSH_VERSION_MINOR < 7 || \
    LIBSSH_VERSION_MICRO < 3
#pragma message "Unsupported libssh version < 0.7.3"
#endif

typedef struct ssh_data {
    ssh_bind sshbind;
    ssh_session session;
} ssh_data;

struct fwd_callbacks potd_ssh_callbacks = {
    .on_listen = ssh_on_listen,
    .on_shutdown = ssh_on_shutdown
};

static int set_default_keys(ssh_bind sshbind, int rsa_already_set,
                            int dsa_already_set, int ecdsa_already_set);
static int gen_default_keys(void);
static int gen_export_sshkey(enum ssh_keytypes_e type, int length, const char *path);
static void ssh_log_cb(int priority, const char *function, const char *buffer, void *userdata);


int ssh_init_cb(struct forward_ctx *ctx)
{
    N("libssh version: %s", ssh_version(0));
    if (ssh_version(SSH_VERSION_INT(0,7,3)) == NULL)
    {
        W("%s", "Unsupported libssh version < 0.7.3");
    }
    if (ssh_version(SSH_VERSION_INT(0,7,4)) != NULL ||
        ssh_version(SSH_VERSION_INT(0,7,90)) != NULL)
    {
        W("%s",
          "libssh versions > 0.7.3 may suffer "
          "from problems with the pki key generation/export");
    }
    ctx->fwd_cbs = potd_ssh_callbacks;

    if (ssh_init())
        return 1;

    ssh_data *d = (ssh_data *) calloc(1, sizeof(*d));
    assert(d);
    d->sshbind = ssh_bind_new();
    d->session = ssh_new();
    ctx->data = d;

    ssh_set_log_callback(ssh_log_cb);
    ssh_set_log_level(SSH_LOG_FUNCTIONS);

    if (!d->sshbind || !d->session)
        return 1;
    if (gen_default_keys())
        return 1;
    if (set_default_keys(d->sshbind, 0, 0, 0))
        return 1;

    return 0;
}

int ssh_on_listen(struct forward_ctx *ctx, const char *host,
                  const char *port)
{
    ssh_data *d = (ssh_data *) ctx->data;

    if (ssh_bind_options_set(d->sshbind, SSH_BIND_OPTIONS_BINDADDR,
                             host))
        return 1;
    if (ssh_bind_options_set(d->sshbind, SSH_BIND_OPTIONS_BINDPORT_STR,
                             port))
        return 1;

    if (ssh_bind_listen(d->sshbind) < 0) {
        E("Error listening to SSH socket: %s", ssh_get_error(d->sshbind));
    }

    return 0;
}

int ssh_on_shutdown(struct forward_ctx *ctx)
{
    return 0;
}

static int set_default_keys(ssh_bind sshbind, int rsa_already_set,
                             int dsa_already_set, int ecdsa_already_set)
{
    if (!rsa_already_set) {
        if (ssh_bind_options_set(sshbind, SSH_BIND_OPTIONS_RSAKEY,
                                 "./ssh_host_rsa_key")) {
            E2("Faled to set RSA key: %s", ssh_get_error(sshbind));
            return 1;
        }
    }
    if (!dsa_already_set) {
        if (ssh_bind_options_set(sshbind, SSH_BIND_OPTIONS_DSAKEY,
                                 "./ssh_host_dsa_key")) {
            E2("Failed to set DSA key: %s", ssh_get_error(sshbind));
            return 1;
        }
    }
    if (!ecdsa_already_set) {
        if (ssh_bind_options_set(sshbind, SSH_BIND_OPTIONS_ECDSAKEY,
                                 "./ssh_host_ecdsa_key")) {
            E2("Failed to set ECDSA key: %s", ssh_get_error(sshbind));
            return 1;
        }
    }

    return 0;
}

static int gen_default_keys(void)
{
    int s = 0;

    if (gen_export_sshkey(SSH_KEYTYPE_RSA, 1024, "./ssh_host_rsa_key")) {
        W("libssh %s key generation failed, using fallback ssh-keygen", "RSA");
        remove("./ssh_host_rsa_key");
        s |= system("ssh-keygen -t rsa -b 1024 -f ./ssh_host_rsa_key -N '' >/dev/null 2>/dev/null");
    }
    if (gen_export_sshkey(SSH_KEYTYPE_DSS, 1024, "./ssh_host_dsa_key")) {
        W("libssh %s key generation failed, using fallback ssh-keygen", "DSA");
        remove("./ssh_host_dsa_key");
        s |= system("ssh-keygen -t dsa -b 1024 -f ./ssh_host_dsa_key -N '' >/dev/null 2>/dev/null");
    }
    if (gen_export_sshkey(SSH_KEYTYPE_ECDSA, 1024, "./ssh_host_ecdsa_key")) {
        W("libssh %s key generation failed, using fallback ssh-keygen", "ECDSA");
        remove("./ssh_host_ecdsa_key");
        s |= system("ssh-keygen -t ecdsa -b 256 -f ./ssh_host_ecdsa_key -N '' >/dev/null 2>/dev/null");
    }

    return s != 0;
}

static int gen_export_sshkey(enum ssh_keytypes_e type, int length, const char *path)
{
    ssh_key priv_key;
    const char *type_str = NULL;
    int s;

    assert(path);
    assert(length == 1024 || length == 2048 ||
           length == 4096);

    switch (type) {
        case SSH_KEYTYPE_DSS:
            type_str = "DSS";
            break;
        case SSH_KEYTYPE_RSA:
            type_str = "RSA";
            break;
        case SSH_KEYTYPE_ECDSA:
            type_str = "ECDSA";
            break;
        default:
            W2("Unknown SSH key type: %d", type);
            return 1;
    }
    N2("Generating %s key with length %d bits and save it on disk: %s",
        type_str, length, path);
    s = ssh_pki_generate(type, length, &priv_key);
    if (s != SSH_OK) {
        E2("Generating %s key failed: %d", type_str, s);
        return 1;
    }
    s = ssh_pki_export_privkey_file(priv_key, "", NULL,
                                    NULL, path);
    ssh_key_free(priv_key);

    if (s != SSH_OK) {
        E2("SSH private key export to file failed: %d", s);
        return 1;
    }

    return 0;
}

static void ssh_log_cb(int priority, const char *function,
                       const char *buffer, void *userdata)
{
    switch (priority) {
        case 0:
            W("libssh: %s", buffer);
            break;
        case 1:
            N("libssh: %s", buffer);
            break;
        default:
            D("libssh: %s", buffer);
            break;
    }
}