/*
 * Run the login shell or command as the given user in a new pty to prevent
 * terminal injection attacks.
 *
 * Copyright (C) 2016  Simon Ruderich
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#define _GNU_SOURCE

#include <assert.h>
#include <errno.h>
#include <fcntl.h>
#include <grp.h>
#include <limits.h>
#include <poll.h>
#include <pwd.h>
#include <signal.h>
#include <stdarg.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <termios.h>
#include <unistd.h>

/* Default PATH for new process.*/
#ifndef PTYAS_DEFAULT_PATH
/* Default user PATH from Debian's /etc/profile, change as needed. */
# define PTYAS_DEFAULT_PATH "/usr/local/bin:/usr/bin:/bin:/usr/local/games:/usr/games"
#endif


static void die(const char *s) {
    perror(s);
    exit(EXIT_FAILURE);
}
static void die_fmt(const char *fmt, ...) {
    va_list ap;

    va_start(ap, fmt);
    vfprintf(stderr, fmt, ap);
    va_end(ap);

    exit(EXIT_FAILURE);
}

static void open_pty_or_die(int *pty_master, int *pty_slave, uid_t uid) {
    char *slave_path;

    *pty_master = posix_openpt(O_RDWR | O_NOCTTY);
    if (*pty_master == -1) {
        die("posix_openpt");
    }
    slave_path = ptsname(*pty_master);
    if (!slave_path) {
        die("ptsname");
    }
    if (grantpt(*pty_master) != 0) {
        die("grantpt");
    }
    if (unlockpt(*pty_master) != 0) {
        die("unlockpt");
    }

    *pty_slave = open(slave_path, O_RDWR | O_NOCTTY);
    if (*pty_slave == -1) {
        die("open slave tty");
    }
    /* The user must be able to write to the new TTY. Normally grantpt() would
     * do this for us, but we don't trust the user and thus don't want to pass
     * the pty_master to a process running under that uid. */
    if (chown(slave_path, uid, (gid_t)-1) != 0) {
        die("chown slave tty");
    }
}

static void close_or_die(int fd) {
    if (close(fd) != 0) {
        die("close");
    }
}
static void dup2_or_die(int oldfd, int newfd) {
    if (dup2(oldfd, newfd) != newfd) {
        die("dup2");
    }
}
static int snprintf_or_assert(char *str, size_t size, const char *format, ...) {
    int ret;
    va_list ap;

    va_start(ap, format);
    ret = vsnprintf(str, size, format, ap);
    assert(size <= (size_t)INT_MAX);
    assert(ret < (int)size); /* assert output fit into buffer */
    va_end(ap);

    return ret;
}

static void drop_privileges_or_die(uid_t uid, gid_t gid) {
    /* Drop all supplementary group IDs. */
    if (setgroups(0, NULL) != 0) {
        die("setgroups");
    }
    if (getgroups(0, NULL) != 0) {
        die_fmt("failed to drop all groups");
    }

    /* Dropping groups may require privileges, do that first. */
    if (setresgid(gid, gid, gid) != 0) {
        die("setresgid");
    }
    if (setresuid(uid, uid, uid) != 0) {
        die("setresuid");
    }

    /* Ensure we dropped all privileges. */
    {
        uid_t ruid, euid, suid;
        gid_t rgid, egid, sgid;

        if (getresuid(&ruid, &euid, &suid) != 0) {
            die("getresuid");
        }
        if (getresgid(&rgid, &egid, &sgid) != 0) {
            die("getresgid");
        }
        if (       uid != ruid || uid != euid || uid != suid
                || gid != rgid || gid != egid || gid != sgid) {
            die_fmt("failed to drop privileges");
        }
    }
    /* Just to be safe. */
    if (setuid(0) != -1) {
        die_fmt("failed to drop privileges (setuid)");
    }
}

static void quit_with_matching_code(int status) {
    if (WIFEXITED(status)) {
        exit(WEXITSTATUS(status));
    } else if (WIFSIGNALED(status)) {
        kill(getpid(), WTERMSIG(status));
        /* Fall-through, should not happen. */
    }
    abort(); /* Should never happen, die painfully. */
}

static bool read_from_write_to(int from, int to) {
    char buf[4096];

    ssize_t r = read(from, buf, sizeof(buf));
    if (r < 0) {
        return false;
    }

    size_t left = (size_t)r;
    char *data = buf;

    while (left > 0) {
        ssize_t w = write(to, data, left);
        if (w < 0) {
            if (errno == EINTR) {
                continue;
            }
            return false;
        }
        left -= (size_t)w;
        data += (size_t)w;
    }

    return true;
}

static void proxy_input_between_ttys(int pty_master, int ctty, volatile pid_t *pid_to_wait_for) {
    struct pollfd fds[] = {
        { /* 0 */
            .fd = pty_master,
            .events = POLLIN,
        },
        { /* 1 */
            .fd = ctty,
            .events = POLLIN,
        },
    };

    sigset_t sigset, sigset_old;
    sigemptyset(&sigset);
    sigaddset(&sigset, SIGCHLD);
    if (sigprocmask(SIG_BLOCK, &sigset, &sigset_old) != 0) {
        die("sigprocmask block sigchld proxy");
    }

    /* Proxy data until our child has terminated. */
    while (*pid_to_wait_for != 0) {
        /*
         * If a signal happens here _and_ the child hasn't closed pty_slave,
         * we would hang in poll(); therefore ppoll() is necessary.
         */
        nfds_t nfds = sizeof(fds)/sizeof(*fds);
        if (ppoll(fds, nfds, NULL /* no timeout */, &sigset_old) == -1) {
            if (errno == EAGAIN || errno == EINTR) {
                continue;
            } else {
                perror("poll");
            }
            break;
        }

        /* Handle errors first. (Data available before the error occurred
         * might be skipped, but shouldn't matter here.) */
        if (fds[0].revents & (POLLERR | POLLNVAL)) {
            fprintf(stderr, "poll: error on master: %d\n", fds[0].revents);
            break;
        }
        if (fds[1].revents & (POLLERR | POLLNVAL)) {
            fprintf(stderr, "poll: error on ctty: %d\n", fds[1].revents);
            break;
        }

        /* Read data if available. */
        if (fds[0].revents & POLLIN) {
            if (!read_from_write_to(pty_master, ctty)) {
                perror("read from master write to ctty");
                break;
            }
        }
        if (fds[1].revents & POLLIN) {
            if (!read_from_write_to(ctty, pty_master)) {
                perror("read from ctty write to master");
                break;
            }
        }

        /* Finally we are done if either side of the pty has disconnected. */
        if ((fds[0].revents & POLLHUP) || (fds[1].revents & POLLHUP)) {
            break;
        }
    }

    if (sigprocmask(SIG_SETMASK, &sigset_old, NULL) != 0) {
        die("sigprocmask setmask proxy");
    }
}


/*
 * Not sig_atomic_t (as required by POSIX) but I don't know how to do that any
 * other way.
 */
static volatile pid_t pid_to_wait_for;
static int pid_to_wait_for_status;

static void sigchld_handler() {
    int status;
    pid_t pid;

    while ((pid = waitpid(-1, &status, WNOHANG)) > 0) {
        if (pid == pid_to_wait_for) {
            /* Mark that our child has died and we should exit as well. */
            pid_to_wait_for = 0;
            /* We must exit like our child, save status. */
            pid_to_wait_for_status = status;
        }
    }
}


int main(int argc, char **argv) {
    char *exec_argv_shell[] = { NULL, NULL }; /* filled below */
    char **exec_argv = NULL;

    if (argc == 2) {
        /* exec_argv set below */
    } else if (argc > 2) {
        exec_argv = argv + 2;
    } else {
        die_fmt("%s <user> [<cmd>...]\n", argv[0]);
    }

    const char *user = argv[1];

    struct passwd *passwd = getpwnam(user);
    if (!passwd) {
        die_fmt("unknown user name '%s'\n", user);
    }

    uid_t uid = passwd->pw_uid;
    gid_t gid = passwd->pw_gid;

    if (!exec_argv) {
        assert(argc == 2);
        exec_argv_shell[0] = passwd->pw_shell;
        exec_argv = exec_argv_shell;
    }

    int pty_master, pty_slave;

    open_pty_or_die(&pty_master, &pty_slave, uid);

    int ctty = open("/dev/tty", O_RDWR | O_NOCTTY); /* controlling TTY */
    if (ctty == -1) {
        die("open /dev/tty");
    }

    sigset_t sigset, sigset_old;
    sigemptyset(&sigset);
    sigaddset(&sigset, SIGCHLD);
    if (sigprocmask(SIG_BLOCK, &sigset, &sigset_old) != 0) {
        die("sigprocmask block sigchld");
    }

    pid_t pid = fork();
    if (pid == -1) {
        die("fork parent");
    } else if (pid == 0) {
        /* child, will become a session leader */

        if (sigprocmask(SIG_SETMASK, &sigset_old, NULL) != 0) {
            die("sigprocmask setmask child");
        }

        struct winsize size;
        if (ioctl(ctty, TIOCGWINSZ, &size) == -1) {
            die("ioctl TIOCGWINSZ");
        }

        close_or_die(pty_master);
        close_or_die(ctty);

        /* Start a new session and attach controlling TTY. */
        if (setsid() == -1) {
            die("setsid");
        }
        if (ioctl(pty_slave, TIOCSCTTY, 0) == -1) {
            die("ioctl TIOCSCTTY");
        }

        if (ioctl(pty_slave, TIOCSWINSZ, &size) == -1) {
            die("ioctl TIOCSWINSZ");
        }

        pid_t pid = fork();
        if (pid == -1) {
            die("fork child");
        } else if (pid == 0) {
            drop_privileges_or_die(uid, gid);

            dup2_or_die(pty_slave, STDIN_FILENO);
            dup2_or_die(pty_slave, STDOUT_FILENO);
            dup2_or_die(pty_slave, STDERR_FILENO);
            close_or_die(pty_slave);

            const char *term_orig = getenv("TERM");
            const char *term = term_orig;
            if (!term) {
                term = ""; /* for strlen() below */
            }
            const char *home = passwd->pw_dir;

            char envp_user[strlen("USER=") + strlen(user) + 1];
            char envp_home[strlen("HOME=") + strlen(home) + 1];
            char envp_term[strlen("TERM=") + strlen(term) + 1];
            snprintf_or_assert(envp_user, sizeof(envp_user), "USER=%s", user);
            snprintf_or_assert(envp_home, sizeof(envp_home), "HOME=%s", home);
            snprintf_or_assert(envp_term, sizeof(envp_term), "TERM=%s", term);

            char *exec_envp[] = {
                "PATH=" PTYAS_DEFAULT_PATH,
                envp_user,
                envp_home,
                term_orig ? envp_term : NULL,
                NULL,
            };

            execve(exec_argv[0], exec_argv, exec_envp);
            die("execve");
        }
        close_or_die(pty_slave);
        close_or_die(STDIN_FILENO);
        close_or_die(STDOUT_FILENO);
        close_or_die(STDERR_FILENO);

        /* TODO: EINTR? */
        int status;
        if (waitpid(pid, &status, 0) <= 0) {
            die("waitpid child");
        }
        quit_with_matching_code(status);
    }
    close_or_die(pty_slave);

    pid_to_wait_for = pid;
    struct sigaction action = {
        .sa_handler = sigchld_handler,
    };
    if (sigaction(SIGCHLD, &action, NULL) != 0) {
        die("sigaction");
    }

    if (sigprocmask(SIG_SETMASK, &sigset_old, NULL) != 0) {
        die("sigprocmask setmask parent");
    }

    struct termios old_term, term;

    /* Change terminal to raw mode. */
    if (tcgetattr(ctty, &old_term) != 0) {
        die("tcgetattr");
    }
    term = old_term;
    /* From man 3 cfmakeraw; cfmakeraw is non-standard so set it manually. */
    term.c_iflag &= ~(tcflag_t)(IGNBRK | BRKINT | PARMRK | ISTRIP | INLCR | IGNCR | ICRNL | IXON);
    term.c_oflag &= ~(tcflag_t)(OPOST);
    term.c_lflag &= ~(tcflag_t)(ECHO | ECHONL | ICANON | ISIG | IEXTEN);
    term.c_cflag &= ~(tcflag_t)(CSIZE | PARENB);
    term.c_cflag |= CS8;
    if (tcsetattr(ctty, TCSADRAIN, &term) != 0) {
        die("tcsetattr");
    }

    proxy_input_between_ttys(pty_master, ctty, &pid_to_wait_for);

    /* Restore terminal mode. */
    if (tcsetattr(ctty, TCSADRAIN, &old_term) != 0) {
        die("tcsetattr restore");
    }

    /* Wait until we got the status code from our child. poll() might also
     * exit after POLLHUP while we haven't collected the child yet. */
    if (sigprocmask(SIG_BLOCK, &sigset, &sigset_old) != 0) {
        die("sigprocmask block sigchld loop");
    }
    while (pid_to_wait_for != 0) {
        sigsuspend(&sigset_old);
        if (errno != EINTR) {
            die("sigsuspend");
        }
    }
    if (sigprocmask(SIG_SETMASK, &sigset, &sigset_old) != 0) {
        die("sigprocmask setmask sigchld loop");
    }

    /* Try to exit the same way as the spawned process. */
    if (pid_to_wait_for == 0) {
        quit_with_matching_code(pid_to_wait_for_status);
    }
    return EXIT_FAILURE;
}