#include <stdio.h>
#include <stdlib.h>
#include <sched.h>
#include <signal.h>
#include <pty.h>
#include <sys/wait.h>
#include <sys/epoll.h>
#include <sys/prctl.h>
#include <sys/stat.h>
#include <assert.h>

#include "jail.h"
#include "socket.h"
#include "server.h"
#include "utils.h"
#include "log.h"

typedef struct jail_prisoner_process {
    psocket client_psock;
    char host_buf[NI_MAXHOST], service_buf[NI_MAXSERV];
    char *newroot;
} jail_prisoner_process;

typedef enum connection_state {
    CON_OK, CON_IN_TERMINATED, CON_OUT_TERMINATED,
    CON_IN_ERROR, CON_OUT_ERROR
} connection_state;

static int jail_mainloop_epoll(int epoll_fd, jail_ctx *ctx[], size_t siz)
    __attribute__((noreturn));
static int jail_accept_client(jail_ctx *ctx[],
                              size_t siz, struct epoll_event *event);
static int jail_childfn(jail_prisoner_process *ctx)
    __attribute__((noreturn));
static int jail_socket_tty_epoll(jail_prisoner_process *ctx, int tty_fd);
static connection_state
jail_socket_tty_io_epoll(struct epoll_event *ev, int dest_fd);


void jail_init_ctx(jail_ctx **ctx, size_t stacksize)
{
    assert(ctx);
    if (stacksize > MAX_STACKSIZE ||
        stacksize < MIN_STACKSIZE)
    {
        stacksize = MAX_STACKSIZE;
    }
    if (!*ctx)
        *ctx = (jail_ctx *) malloc(sizeof(**ctx));
    assert(*ctx);

    memset(*ctx, 0, sizeof(**ctx));
    (*ctx)->stacksize = stacksize;
    (*ctx)->stack_ptr = calloc(1, (*ctx)->stacksize);
    (*ctx)->stack_beg =
        (unsigned char *) (*ctx)->stack_ptr
            + (*ctx)->stacksize;
}

int jail_setup(jail_ctx *ctx,
               const char *listen_addr, const char *listen_port)
{
    assert(ctx);
    assert(listen_addr || listen_port);

    D2("Try to listen on %s:%s",
       (listen_addr ? listen_addr : "*"), listen_port);
    if (fwd_setup_server(&ctx->fwd_ctx, listen_addr, listen_port))
        return 1;
    if (fwd_listen_sock(&ctx->fwd_ctx, NULL))
        return 1;

    return 0;
}

int jail_validate_ctx(const jail_ctx *ctx)
{
    assert(ctx);
    assert(ctx->fwd_ctx.sock.addr_len > 0);
    assert(ctx->stack_ptr);

    return 0;
}

int jail_setup_epoll(jail_ctx *ctx[], size_t siz)
{
    int s, fd = epoll_create1(0); /* flags == 0 -> obsolete size arg is dropped */
    struct epoll_event ev;

    assert(ctx);
    assert(siz > 0 && siz < POTD_MAXFD);
    if (fd < 0)
        return -1;

    for (size_t i = 0; i < siz; ++i) {
        memset(&ev, 0, sizeof(ev));
        ev.data.fd = ctx[i]->fwd_ctx.sock.fd;
        ev.events = EPOLLIN | EPOLLET;

        s = socket_addrtostr_in(&ctx[i]->fwd_ctx.sock,
                                ctx[i]->host_buf, ctx[i]->service_buf);
        if (s) {
            E_GAIERR(s, "Convert socket address to string");
            return -2;
        }
        N("Jail service listening on %s:%s: %d",
            ctx[i]->host_buf, ctx[i]->service_buf, ev.data.fd);

        s = epoll_ctl(fd, EPOLL_CTL_ADD, ctx[i]->fwd_ctx.sock.fd, &ev);
        if (s) {
            close(fd);
            return -3;
        }
    }

    return fd;
}

pid_t jail_daemonize(int epoll_fd, jail_ctx *ctx[], size_t siz)
{
    pid_t p;
	int s;
	size_t i;

    assert(ctx);
    assert(siz > 0);

	for (i = 0; i < siz; ++i) {
        assert(ctx[i]);
        s = socket_addrtostr_in(&ctx[i]->fwd_ctx.sock,
                                ctx[i]->host_buf, ctx[i]->service_buf);
        if (s) {
            E_GAIERR(s, "Could not initialise jail daemon socket");
            return 1;
        }
    }

	p = fork();
	switch (p) {
        case -1:
            W_STRERR("%s", "Jail daemonize");
	        return -1;
        case 0:
            N("Jail daemon mainloop on Epoll fd %d", epoll_fd);
            jail_mainloop_epoll(epoll_fd, ctx, siz);
    }
    D2("Jail daemon pid: %d", p);

    close(epoll_fd);
    for (i = 0; i < siz; ++i)
        socket_close(&ctx[i]->fwd_ctx.sock);

    return p;
}

static int jail_mainloop_epoll(int epoll_fd, jail_ctx *ctx[], size_t siz)
{
    static struct epoll_event *events = NULL;
    sigset_t eset;

    if (!events)
        events = (struct epoll_event *) calloc(POTD_MAXEVENTS, sizeof(*events));

    assert(events);
    assert(ctx);
    assert(siz > 0 && siz < POTD_MAXFD);

    set_procname("[potd] jail");
    assert( set_child_sighandler() == 0 );
    sigemptyset(&eset);

    while (1) {
        int n, i;

        n = epoll_pwait(epoll_fd, events, POTD_MAXEVENTS, -1, &eset);
        if (n < 0)
            goto error;

        for (i = 0; i < n; ++i) {
            if ((events[i].events & EPOLLERR) ||
                (events[i].events & EPOLLHUP) ||
                (!(events[i].events & EPOLLIN)))
            {
                E_STRERR("Epoll for descriptor %d", events[i].data.fd);
                close(events[i].data.fd);
                continue;
            } else {
                if (jail_accept_client(ctx, siz, &events[i])) {
                    /* new client connection, accept succeeded */
                    continue;
                }
                W2("Jail daemon accept client failed:  [fd: %d , npoll: %d]", events[i].data.fd, n);
            }
        }
    }

    close(epoll_fd);
    exit(EXIT_SUCCESS);
error:
    close(epoll_fd);
    exit(EXIT_FAILURE);
}

static int jail_accept_client(jail_ctx *ctx[],
                              size_t siz, struct epoll_event *event)
{
    size_t i, rc = 0;
    int s;
    pid_t prisoner_pid;
    static jail_prisoner_process *args;

    for (i = 0; i < siz; ++i) {
        if (ctx[i]->fwd_ctx.sock.fd == event->data.fd) {
            args = (jail_prisoner_process *) calloc(1, sizeof(*args));
            assert(args);
            args->newroot = ctx[i]->newroot;

            if (socket_accept_in(&ctx[i]->fwd_ctx.sock, &args->client_psock)) {
                E_STRERR("Could not accept client connection for fd %d",
                    args->client_psock.fd);
                goto error;
            }

            s = socket_addrtostr_in(&args->client_psock,
                                    args->host_buf, args->service_buf);
            if (s) {
                E_GAIERR(s, "Convert socket address to string");
                goto error;
            }
            N2("New connection from %s:%s to %s:%s: %d",
                args->host_buf, args->service_buf,
                ctx[i]->host_buf, ctx[i]->service_buf,
                args->client_psock.fd);

            prisoner_pid = fork();
            switch (prisoner_pid) {
                case -1:
                    W_STRERR("%s", "Jail client fork");
                    goto error;
                case 0:
                    jail_childfn(args);
            }
            N2("Jail prisoner pid %d", prisoner_pid);

            rc = 1;
error:
            socket_close(&args->client_psock);
            free(args);
            args = NULL;
            return rc;
        }
    }

    return rc;
}

static int jail_childfn(jail_prisoner_process *ctx)
{
    const char *path_dev = "/dev";
    const char *path_devpts = "/dev/pts";
    const char *path_proc = "/proc";
    const char *path_shell = "/bin/sh";
    int s, master_fd;
    int unshare_flags = CLONE_NEWUTS|CLONE_NEWPID|CLONE_NEWIPC|
        CLONE_NEWNS|CLONE_NEWNET;
    pid_t self_pid, child_pid;

    assert(ctx);
    self_pid = getpid();

    if (prctl(PR_SET_PDEATHSIG, SIGKILL) != 0)
        FATAL("Jail child prctl for pid %d", self_pid);
    if (!ctx->newroot)
        FATAL("New root set for pid %d", self_pid);

    D2("Unshare prisoner %d", self_pid);
    if (unshare(unshare_flags))
        FATAL("Unshare prisoner %d", self_pid);

    D2("Safe change root to: '%s'", ctx->newroot);
    if (safe_chroot(ctx->newroot))
        FATAL("Safe jail chroot to '%s' failed", ctx->newroot);

    D2("Mounting rootfs to '%s'", ctx->newroot);
    mount_root();

    D2("Mounting devtmpfs to '%s%s'", ctx->newroot, path_dev);
    s = mkdir(path_dev, S_IRUSR|S_IWUSR|S_IXUSR|
                        S_IRGRP|S_IXGRP|
                        S_IROTH|S_IXOTH);
    if (s && errno != EEXIST)
        FATAL("Create directory '%s'", path_dev);
    if (!dir_is_mountpoint(path_dev) && mount_dev(path_dev))
        FATAL("Mount devtmpfs to '%s%s'", ctx->newroot, path_dev);

    D2("Mounting devpts to '%s%s'", ctx->newroot, path_devpts);
    s = mkdir(path_devpts, S_IRUSR|S_IWUSR|S_IXUSR|
                           S_IRGRP|S_IXGRP|
                           S_IROTH|S_IXOTH);
    if (s && errno != EEXIST)
        FATAL("Create directory '%s'", path_devpts);
    if (!dir_is_mountpoint(path_devpts) && mount_pts(path_devpts))
        FATAL("Mount devpts to '%s%s'", ctx->newroot, path_devpts);

    D2("Mounting proc to '%s%s'", ctx->newroot, path_proc);
    s = mkdir(path_proc, S_IRUSR|S_IWUSR|S_IXUSR|
                         S_IRGRP|S_IXGRP|
                         S_IROTH|S_IXOTH);
    if (s && errno != EEXIST)
        FATAL("Create directory '%s'", path_proc);
    if (!dir_is_mountpoint(path_proc) && mount_proc(path_proc))
        FATAL("Mount devpts to '%s%s'", ctx->newroot, path_proc)

    D2("Creating device files in '%s%s'", ctx->newroot, path_dev);
    if (create_device_files(path_dev)) {
        E2("Device file creation failed for rootfs '%s%s'",
            ctx->newroot, path_dev);
        exit(EXIT_FAILURE);
    }

    D2("Forking a new pty process for "
       "parent %d", self_pid);
    child_pid = forkpty(&master_fd, NULL, NULL, NULL);
    switch (child_pid) {
        case -1:
            FATAL("Forking a new process for the slave tty from "
                "parent pty with pid %d",
                self_pid);
            break;
        case 0:
            D2("Executing '%s'", path_shell);
            if (execl(path_shell, path_shell, (char *) NULL))
                FATAL("Execute a shell for pid %d", self_pid);
            break;
        default:
            if (set_fd_nonblock(master_fd))
                FATAL("Pty master fd nonblock for prisoner pid %d",
                    child_pid);
            N("Socket to tty I/O loop for prisoner pid %d",
                child_pid);
            if (jail_socket_tty_epoll(ctx, master_fd))
                FATAL("Socket to tty I/O loop for prisoner pid %d",
                    child_pid);
            waitpid(child_pid, &s, 0);
    }

    exit(EXIT_FAILURE);
}

static int jail_socket_tty_epoll(jail_prisoner_process *ctx, int tty_fd)
{
    int s, fd;
    struct epoll_event event = {0,{0}};
    struct epoll_event *events;
    sigset_t eset;

    assert(ctx);
    events = (struct epoll_event *) calloc(POTD_MAXEVENTS, sizeof(*events));
    assert(events);

    fd = epoll_create1(0);
    if (fd < 0)
        return -1;

    event.events = EPOLLIN | EPOLLET;
    event.data.fd = ctx->client_psock.fd;
    s = epoll_ctl(fd, EPOLL_CTL_ADD, ctx->client_psock.fd, &event);
    if (s)
        FATAL("Jail Socket Epoll for client %s:%s",
            ctx->host_buf, ctx->service_buf);
    event.data.fd = tty_fd;
    s = epoll_ctl(fd, EPOLL_CTL_ADD, tty_fd, &event);
    if (s)
        FATAL("Jail TTY Epoll for client %s:%s",
            ctx->host_buf, ctx->service_buf);

    sigemptyset(&eset);

    while (1) {
        int n, i;

        n = epoll_pwait(fd, events, POTD_MAXEVENTS, -1, &eset);
        if (n < 0)
            break;

        for (i = 0; i < n; ++i) {
        }
    }

    return 0;
}

static connection_state
jail_socket_tty_io_epoll(struct epoll_event *ev, int dest_fd)
{
    return CON_OK;
}