#include <stdio.h>
#include <assert.h>
#include <sys/types.h>
#include <sys/wait.h>

#include "pseccomp.h"
#include "capabilities.h"
#include "log.h"
#include "log_colored.h"
#include "utils.h"
#include "redirector.h"
#include "protocol_ssh.h"
#include "forward.h"
#include "jail.h"
#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

static void jail_preinit(const char *jail_ports[], jail_ctx *ctx[],
                         const size_t siz);
static pid_t jail_init(jail_ctx *ctx[], const size_t siz);
static void ssh_protocol_preinit(const char *ssh_ports[], protocol_ctx *ctx[],
                                 const char *jail_ports[], const size_t siz);
static void ssh_protocol_init(protocol_ctx *ctx[], const size_t siz);
static void rdr_preinit(const char *rdr_ports[], redirector_ctx *ctx[],
                        const size_t siz);
static pid_t rdr_init(redirector_ctx *ctx[], const size_t siz);


static void jail_preinit(const char *jail_ports[], jail_ctx *ctx[],
                         const size_t siz)
{
    for (size_t i = 0; i < siz; ++i) {
        D("Initialising jail service on port %s", jail_ports[i]);

        jail_init_ctx(&ctx[i], MAX_STACKSIZE);
        //jail[i]->newroot = strdup("/home/lns/git/busybox/sysroot");
        ctx[i]->newroot = strdup("/home/toni/git/busybox/_install");
        ABORT_ON_FATAL( jail_setup(ctx[i], "127.0.0.1", jail_ports[i]),
            "Jail daemon setup" );
        ABORT_ON_FATAL( jail_validate_ctx(ctx[i]),
            "Jail validation" );
    }
}

static pid_t jail_init(jail_ctx *ctx[], const size_t siz)
{
    pid_t jail_pid;
    event_ctx *event = NULL;

    ABORT_ON_FATAL( jail_setup_event( ctx, siz, &event ),
        "Jail daemon epoll setup" );
    jail_pid = jail_daemonize(&event, ctx, siz);
    ABORT_ON_FATAL( jail_pid < 1, "Jail daemon startup" );

    return jail_pid;
}

static void ssh_protocol_preinit(const char *ssh_ports[], protocol_ctx *ctx[],
                                 const char *jail_ports[], const size_t siz)
{
    for (size_t i = 0; i < siz; ++i) {
        ABORT_ON_FATAL( proto_init_ctx(&ctx[i], ssh_init_cb),
            "SSH Protocol init" );
        ABORT_ON_FATAL( proto_setup(ctx[i], "127.0.0.1", ssh_ports[i],
            "127.0.0.1", jail_ports[i]), "SSH Protocol setup" );
        ABORT_ON_FATAL( proto_validate_ctx(ctx[i]),
            "SSH validation" );
    }
}

static void ssh_protocol_init(protocol_ctx *ctx[], const size_t siz)
{
    for (size_t i = 0; i < siz; ++i) {
        ABORT_ON_FATAL( proto_listen(ctx[i]),
            "SSH Protocol listen" );
    }
}

static void rdr_preinit(const char *rdr_ports[], redirector_ctx *ctx[],
                        const size_t siz)
{
    for (size_t i = 0; i < siz; ++i) {
        D("Initialising redirector service on port %s", rdr_ports[i]);

        ABORT_ON_FATAL( redirector_init_ctx(&ctx[i]),
            "Redirector init" );
        ABORT_ON_FATAL( redirector_setup(ctx[i], NULL, rdr_ports[i],
            "127.0.0.1", "22222"), "Redirector setup" );
        ABORT_ON_FATAL( redirector_validate_ctx(ctx[i]),
            "Redirector validation" );
    }
}

static pid_t rdr_init(redirector_ctx *ctx[], const size_t siz)
{
    pid_t rdr_pid;
    event_ctx *event = NULL;

    D2("%s", "Redirector event setup");
    ABORT_ON_FATAL( redirector_setup_event( ctx, siz, &event ),
        "Redirector event setup" );

    N("%s", "Redirector epoll mainloop");
    rdr_pid = redirector_daemonize( event, ctx, siz );
    ABORT_ON_FATAL( rdr_pid < 1, "Server epoll mainloop" );

    return rdr_pid;
}

int main(int argc, char *argv[])
{
    const size_t rdr_siz = 3;
    const size_t proto_siz = 2;
    const size_t jail_siz = 2;
    const char *rdr_ports[rdr_siz];
    const char *proto_ports[proto_siz];
    const char *jail_ports[jail_siz];
    redirector_ctx *rdr[rdr_siz];
    protocol_ctx *ssh_proto[proto_siz];
    jail_ctx *jail[jail_siz];
    int proc_status;
    pid_t daemon_pid, rdr_pid, jail_pid, child_pid;
    pseccomp_ctx *psc = NULL;

    (void) argc;
    (void) argv;
    arg0 = argv[0];

    LOG_SET_FUNCS_VA(LOG_COLORED_FUNCS);
#ifdef HAVE_CONFIG_H
    N("%s (C) 2018 Toni Uhlig (%s)", PACKAGE_STRING, PACKAGE_BUGREPORT);
#endif

    if (geteuid() != 0) {
        E("%s", "I was made for root!");
        exit(EXIT_FAILURE);
    }

    caps_default_filter();
    pseccomp_init(&psc, 0);
    if (pseccomp_default_rules(psc))
        FATAL("%s", "SECCOMP: adding default rules");
    pseccomp_free(&psc);

    D("%s", "Forking into background/foreground");
    daemon_pid = daemonize(1);
    ABORT_ON_FATAL( daemon_pid > 0, "Forking" );
    if (daemon_pid == 0) {
        set_procname("[potd] main");
    } else {
        FATAL("Forking (fork returned %d)", daemon_pid);
    }
    D2("Master pid: %d", getpid());
    ABORT_ON_FATAL( set_master_sighandler(),
        "Master sighandler" );

    memset(jail, 0, sizeof(jail));
    jail_ports[0] = "33333";
    jail_ports[1] = "33334";
    jail_preinit(jail_ports, jail, SIZEOF(jail_ports));
    jail_pid = jail_init(jail, SIZEOF(jail_ports));

    memset(ssh_proto, 0, sizeof(proto_ports));
    proto_ports[0] = "22222";
    proto_ports[1] = "22223";
    assert(SIZEOF(proto_ports) == SIZEOF(jail_ports));
    ssh_protocol_preinit(proto_ports, ssh_proto, jail_ports, proto_siz);
    ssh_protocol_init(ssh_proto, proto_siz);

    memset(rdr, 0, sizeof(rdr));
    rdr_ports[0] = "2222";
    rdr_ports[1] = "2223";
    rdr_ports[2] = "22050";
    rdr_preinit(rdr_ports, rdr, SIZEOF(rdr_ports));
    rdr_pid = rdr_init(rdr, SIZEOF(rdr_ports));

    while (1) {
        child_pid = wait(&proc_status);
        if (child_pid == jail_pid ||
            child_pid == rdr_pid) {
            E2("%s daemon with pid %d terminated, exiting",
                (child_pid == jail_pid ? "Jail" : "Redirector"),
                (child_pid == jail_pid ? jail_pid : rdr_pid));
            kill(getpid(), SIGTERM);
            break;
        } else W2("Process with pid %d terminated", child_pid);
    }

    return 0;
}