/*
 * Created by yuuta on 12/23/22.
 */

#include "common.h"
#include "log.h"

#define IOCTL(a, b, ...) \
if (ioctl(a, b, __VA_ARGS__) < 0) { LOGFV("ioctl(%d, 0x%x): %m", a, b, errno); return errno; }

#include <stdio.h>
#include <fcntl.h>
#include <unistd.h>
#include <errno.h>
#include <stdlib.h>
#include <limits.h>
#include <bsd/unistd.h>

#include <sys/ioctl.h>
#include <sys/wait.h>
#include <linux/kd.h>
#include <linux/vt.h>

static const char *conspath[] = {
        "/proc/self/fd/0",
        "/proc/self/fd/1",
        "/proc/self/fd/2",
        "/dev/tty",
        "/dev/tty0",
        "/dev/vc/0",
        "/dev/systty",
        "/dev/console"
};

static int fd_con = -1;
static struct vt_stat vts;
static char vt_dev[PATH_MAX + 1];
static int vt_fd = -1;
static pid_t chld = -1;

static int vt_switch(int to) {
    if (to) {
        IOCTL(fd_con, VT_ACTIVATE, p_env.vt)
        IOCTL(fd_con, VT_WAITACTIVE, p_env.vt)
    } else {
        IOCTL(fd_con, VT_ACTIVATE, vts.v_active)
        IOCTL(fd_con, VT_WAITACTIVE, vts.v_active)
    }
    return 0;
}

static int find_current_tty(void) {
    for (unsigned int i = 0; i < sizeof(conspath) / sizeof(char *); i++) {
        if ((fd_con = open(conspath[i], O_RDONLY)) == -1) {
            LOGDV("Cannot open '%s': %m.",
                  conspath[i],
                  errno);
            continue;
        }
        char arg;
        errno = 0;
        if (!isatty(fd_con) || ioctl(fd_con, KDGKBTYPE, &arg) ||
            ((arg != KB_101) && (arg != KB_84))) {
            if (errno) {
                LOGDV("Cannot open '%s': %m.",
                      conspath[i]);
            } else {
                LOGDV("Cannot open '%s': Not a TTY.", conspath[i]);
            }
            close(fd_con);
            fd_con = -1;
            continue;
        }
        LOGDV("Opened '%s': %d.", conspath[i], fd_con);
        break;
    }
    if (fd_con == -1) {
        LOGE("Cannot get the current console. Try running as root to get tty0.");
        return 1;
    }
    return 0;
}

static int main_sak(void) {
    setproctitle("sak");
    /* We have to take the ownership here - otherwise SAK will kill the parent. */
    if (setsid() == -1) {
        LOGFV("Cannot set session ID: %m.", errno);
        return errno;
    }
    int fd;
    if ((fd = open(vt_dev, O_RDWR)) == -1) {
        LOGFV("Cannot open '%s': %m.", vt_dev, errno);
        return errno;
    }
    LOGDV("Opened '%s': %d.", vt_dev, fd);
    IOCTL(fd, TIOCSCTTY, 0)
    vt_switch(1);

    dprintf(fd, "\033[r\033[H\033[J"
                   "Press SysRq + K to continue.\n");
    /* Wait for signal. */
    pause();
    return 0;
}

int sd_setup(void) {
    if (find_current_tty()) {
        return 1;
    }
    IOCTL(fd_con, VT_GETSTATE, &vts)
    IOCTL(fd_con, VT_OPENQRY, &p_env.vt)
    snprintf(vt_dev, PATH_MAX, "/dev/tty%d", p_env.vt);

    if (!(chld = fork())) {
        exit(main_sak());
    }
    if (chld == -1) {
        LOGFV("Cannot fork: %m.", errno);
        return errno;
    }
    int chld_status = -1;
    if (waitpid(chld, &chld_status, 0) != chld) {
        LOGFV("Cannot wait for child: %m.", errno);
        return errno;
    }

    if (!WIFSIGNALED(chld_status)) {
        LOGF("Not killed by signal. Considered insecure.");
        return 1;
    }
    switch (WTERMSIG(chld_status)) {
        case SIGKILL: {
            break;
        }
        case SIGINT:
        case SIGTERM: {
            return -1;
        }
        default: {
            /* Consider insecure.
             * Switch back and use the insecure view. */
            LOGF("Killed by an incorrect signal. Considered insecure.");
            return 1;
        }
    }

    /* Take ownership */
    if ((vt_fd = open(vt_dev, O_RDWR)) == -1) {
        LOGFV("Cannot open '%s': %m.", vt_dev, errno);
        return errno;
    }
    LOGDV("Opened '%s': %d.", vt_dev, vt_fd);
    if (setsid() == -1) {
        LOGFV("Cannot set session ID: %m.", errno);
        return errno;
    }
    IOCTL(vt_fd, TIOCSCTTY, 1)

    int in_bak, out_bak, err_bak;
    if ((in_bak = dup(p_env.in)) == -1 ||
        (out_bak = dup(p_env.out)) == -1 ||
        (err_bak = dup(p_env.err)) == -1) {
        LOGFV("dup: %m", errno);
        return errno;
    }

    p_env.in = in_bak;
    p_env.out = out_bak;
    p_env.err = err_bak;

    if (dup2(vt_fd, STDIN_FILENO) == -1 ||
        dup2(vt_fd, STDOUT_FILENO) == -1 ||
        dup2(vt_fd, STDERR_FILENO) == -1) {
        LOGFV("dup2: %m", errno);
    }

    return 0;
}

void sd_cleanup(void) {
    if (vt_fd >= 0) {
        if (ioctl(vt_fd, TIOCNOTTY) < 0) {
            LOGEV("Cannot deallocate console: %m", errno);
        }
        if (p_env.in != STDIN_FILENO) p_env.in = dup2(p_env.in, STDIN_FILENO);
        if (p_env.out != STDOUT_FILENO) p_env.out = dup2(p_env.out, STDOUT_FILENO);
        if (p_env.err != STDERR_FILENO) p_env.err = dup2(p_env.err, STDERR_FILENO);
        /* It should already be closed? */
        close(vt_fd);
        vt_fd = -1;
    }
    if (fd_con >= 0) {
        if (vt_switch(0)) {
            LOGEV("Cannot return to the previous console: %m", errno);
        }
        if (ioctl(fd_con, VT_DISALLOCATE, p_env.vt) < 0) {
            LOGEV("Cannot deallocate console: %m", errno);
        }
        close(fd_con);
        fd_con = -1;
    }
}