/*
 * Receive wall messages and pass them to a notification program via stdin.
 *
 * Copyright (C) 2014-2015  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/>.
 */


#include <config.h>

#include <assert.h>
#include <fcntl.h>
#include <limits.h>
#include <pwd.h>
#include <signal.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/select.h>
#include <sys/stat.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>

#ifdef USE_UTEMPTER
# include <utempter.h>
#endif
#ifdef USE_UTMPX
# include <utmpx.h>
#endif
#ifndef DONT_USE_X11
# include <pthread.h>
# include <X11/Xlib.h>
#endif


static sig_atomic_t signaled = 0;


static void sig_handler(int signal) {
    (void)signal;

    signaled = 1;
}
static void setup_signals(void) {
    struct sigaction action;

    memset(&action, 0, sizeof(action));
    sigemptyset(&action.sa_mask);
    action.sa_handler = sig_handler;

    /* Handle all important signals which might be sent to us so we break out
     * of the read()-loop below and can perform our cleanup. */
    sigaction(SIGHUP,  &action, NULL);
    sigaction(SIGINT,  &action, NULL);
    sigaction(SIGQUIT, &action, NULL);
    sigaction(SIGUSR1, &action, NULL);
    sigaction(SIGUSR2, &action, NULL);
}

static int set_cloexec(int fd) {
    int flags = fcntl(fd, F_GETFD);
    if (flags == -1) {
        return 0;
    }
    if (fcntl(fd, F_SETFD, flags | FD_CLOEXEC) == -1) {
        return 0;
    }
    return 1;
}

static int open_tty(int public) {
    int ptm;
    const char *name;

    ptm = posix_openpt(O_RDWR);
    if (ptm < 0) {
        return -1;
    }
    if (grantpt(ptm) != 0) {
        return -1;
    }
    if (!set_cloexec(ptm)) {
        return -1;
    }

    /* Prevent write access for other users so they can't use wall to send
     * messages to this program. */
    if (!public) {
        name = ptsname(ptm);
        if (!name) {
            return -1;
        }
        if (chmod(name, S_IRUSR | S_IWUSR) != 0) {
            return -1;
        }
    }

    if (unlockpt(ptm) != 0) {
        return -1;
    }

    return ptm;
}

#ifdef USE_UTMPX
static const char *skip_prefix(const char *string, const char *prefix) {
    size_t length = strlen(prefix);

    if (!strncmp(string, prefix, length)) {
        return string + length;
    } else {
        return string;
    }
}
static int set_utmpx(short type, int ptm) {
    struct utmpx entry;

    const char *tty, *user, *id, *line;
    struct timeval now;

    user = getpwuid(getuid())->pw_name;
    gettimeofday(&now, NULL);

    tty = ptsname(ptm);
    if (!tty) {
        return 0;
    }

    id   = skip_prefix(tty, "/dev/pts/");
    line = skip_prefix(tty, "/dev/");

    /* Create utmp entry for the given terminal. */
    memset(&entry, 0, sizeof(entry));

    snprintf(entry.ut_user, sizeof(entry.ut_user), "%s", user);
    snprintf(entry.ut_id,   sizeof(entry.ut_id),   "%s", id);
    snprintf(entry.ut_line, sizeof(entry.ut_line), "%s", line);

    entry.ut_pid  = getpid();
    entry.ut_type = type;
    entry.ut_tv.tv_sec  = now.tv_sec;
    entry.ut_tv.tv_usec = now.tv_usec;

    /* Write the entry to the utmp file. */
    setutxent();
    if (!pututxline(&entry)) {
        return 0;
    }
    endutxent();

    return 1;
}
#endif
static int login(int ptm) {
#if defined(USE_UTEMPTER)
    int result = utempter_add_record(ptm, NULL);
    /* Exit value of utempter_*() is not correct on all systems, e.g. FreeBSD
     * always returns 0. Checking the utmpx database manually is difficult
     * because we don't know the exact values for ut_id and ut_line, therefore
     * we only check the return value on systems known to return a useful
     * value. */
# ifndef __linux
    result = 1;
# endif
    return result;
#elif defined(USE_UTMPX)
    return set_utmpx(USER_PROCESS, ptm);
#else
# error "neither USE_UTEMPTER nor USE_UTMPX defined"
#endif
}
static int logout(int ptm) {
#if defined(USE_UTEMPTER)
    int result = utempter_remove_record(ptm);
    /* See above. */
# ifndef __linux
    result = 1;
# endif
    return result;
#elif defined(USE_UTMPX)
    return set_utmpx(DEAD_PROCESS, ptm);
#else
# error "neither USE_UTEMPTER nor USE_UTMPX defined"
#endif
}

static int wait_for_write(int fd, int timeout) {
    fd_set rfds;
    struct timeval tv;
    int result;

    FD_ZERO(&rfds);
    FD_SET(fd, &rfds);

    tv.tv_sec  = timeout;
    tv.tv_usec = 0;

    result = select(fd + 1, &rfds, NULL, NULL, &tv);
    if (result < 0) {
        perror("select");
        return 0;
    }
    if (result == 0) {
        /* Timeout. */
        return 0;
    }

    /* Got more data to read. */
    return 1;
}

#ifdef USE_UTMPX
static void drop_privileges(void) {
    uid_t uid, ruid, euid, suid;
    gid_t gid, rgid, egid, sgid;

    uid = getuid();
    gid = getgid();

    /* Drop all privileges. */
    if (setresgid(gid, gid, gid) != 0) {
        perror("setresgid");
        exit(EXIT_FAILURE);
    }
    if (setresuid(uid, uid, uid) != 0) {
        perror("setresuid");
        exit(EXIT_FAILURE);
    }

    /* Verify all privileges were dropped. */
    if (getresuid(&ruid, &euid, &suid) != 0) {
        perror("getresuid");
        exit(EXIT_FAILURE);
    }
    if (getresgid(&rgid, &egid, &sgid) != 0) {
        perror("getresgid");
        exit(EXIT_FAILURE);
    }
    if (uid == ruid && uid == euid && uid == suid
            && gid == rgid && gid == egid && gid == sgid) {
        /* Everything fine. */
        return;
    }

    fprintf(stderr, "failed to drop privileges, aborting\n");
    exit(EXIT_FAILURE);
}
#endif

static void pass_buffer_to_program(const char *buffer, size_t length, char **argv) {
    int fds[2];
    FILE *fh;

    pid_t pid;

    if (pipe(fds) != 0) {
        perror("pipe");
        return;
    }

    fh = fdopen(fds[1] /* write side */, "w");
    if (!fh) {
        perror("fdopen");
        close(fds[0]);
        close(fds[1]);
        return;
    }

    pid = fork();
    if (pid < 0) {
        perror("fork");
        goto out;
    } else if (pid == 0) {
        /* child */

#ifdef USE_UTMPX
        drop_privileges();
#endif

        close(fds[1]); /* write side */

        /* Pass read side as stdin to the program. */
        if (dup2(fds[0], STDIN_FILENO) < 0) {
            perror("dup2");
            exit(EXIT_FAILURE);
        }
        close(fds[0]);

        /* Double fork so `wall-notify` doesn't have to wait for the
         * notification process to terminate. We can't just use
         * signal(SIGCHLD, SIG_IGN); because utempter on at least FreeBSD
         * doesn't work if SIGCHLD is ignored. */
        pid = fork();
        if (pid < 0) {
            perror("fork");
            exit(EXIT_FAILURE);
        } else if (pid == 0) {
            execvp(argv[0], argv);
            perror("execvp");
            exit(EXIT_FAILURE);
        }

        exit(EXIT_SUCCESS);
    }
    /* father */

    if (fwrite(buffer, 1, length, fh) != length) {
        perror("fwrite");
        /* continue to perform cleanup */
    }

out:
    close(fds[0]); /* read side */
    fclose(fh);

    if (waitpid(pid, NULL, 0) < 0) {
        perror("waitpid");
    }
}
static void handle_wall(int fd, char **argv) {
    char buffer[4096];
    ssize_t r;

    assert(SSIZE_MAX <= SIZE_MAX);
    while (!signaled && (r = read(fd, buffer, sizeof(buffer))) > 0) {
        size_t space;
        ssize_t r2;

        /* To prevent partial messages (sometimes it takes multiple reads to
         * get the complete message) wait for a short time to get the rest of
         * the message. */
        space = sizeof(buffer) - (size_t)r;
        while (space > 0 && wait_for_write(fd, 1 /* second */)) {
            r2 = read(fd, buffer + r, space);
            if (r2 <= 0) {
                break;
            }
            r += r2;
            space -= (size_t)r2;
        }

        pass_buffer_to_program(buffer, (size_t)r, argv);
    }
}
#ifndef DONT_USE_X11
static void *x11_event_loop_thread(void *unused) {
    Display *display;
    XEvent event;

    (void)unused;

    pthread_detach(pthread_self());

    display = XOpenDisplay(NULL);
    if (!display) {
        fprintf(stderr, "failed to connect to X server\n");
        exit(EXIT_FAILURE);
    }

    /* Do nothing. We just want to die if the X11 session is closed. */
    while (1) {
        XNextEvent(display, &event);
    }
}
#endif

static void usage(const char *argv0) {
    fprintf(stderr, "usage: %s [-X] [-m] [-v] <cmd args..>\n", argv0);
    fprintf(stderr, "Pass wall messages to <cmd args..>.\n");
    fprintf(stderr, "\n");
    fprintf(stderr, "-X quit when the current X session terminates\n");
    fprintf(stderr, "-m allow messages from normal users (like `mesg y`)\n");
    fprintf(stderr, "-v display version\n");
#ifdef DONT_USE_X11
    fprintf(stderr, "\n");
    fprintf(stderr, "compiled without X11 support, -X disabled\n");
#endif
    exit(EXIT_FAILURE);
}


int main(int argc, char **argv) {
    int option, enable_x11, mesg_yes;
    const char *argv0;
    int ptm, pts;
    char *name;

    /* Don't display error messages for unknown options. */
    opterr = 0;

    enable_x11 = 0;
    mesg_yes = 0;

    argv0 = argv[0];

    /*
     * Glibc violates POSIX by default and skips over non-option arguments and
     * parses later arguments which look like options as well. But we want to
     * pass everything except the options unmodified to execvp(). Prefixing
     * the optstring with "+" fixes this behaver. This is not POSIX
     * compatible, but the option should be ignored on other systems.
     */
    while ((option = getopt(argc, argv, "+Xmhv")) != -1) {
        switch (option) {
            case 'X':
                enable_x11 = 1;
                break;
            case 'm':
                mesg_yes = 1;
                break;
            case 'h':
                usage(argv0);
                break;
            case 'v':
                printf("%s%s\n",
                       PACKAGE_STRING,
                       strlen(GIT_VERSION) ? " (Git " GIT_VERSION ")" : "");
                exit(EXIT_SUCCESS);
                break;
            default:
                fprintf(stderr, "%s: unknown option '%s'!\n\n",
                                argv0, argv[optind - 1]);
                usage(argv0);
                break;
        }
    }

    /* No arguments remaining, abort. */
    if (!argv[optind]) {
        usage(argv0);
    }
    /* Arguments for notification program. */
    argv += optind;

    ptm = open_tty(mesg_yes);
    if (ptm < 0) {
        perror("open_tty");
        exit(EXIT_FAILURE);
    }
    name = ptsname(ptm);
    if (!name) {
        perror("ptsname");
        exit(EXIT_FAILURE);
    }

#ifdef DEBUG
    printf("%s\n", name);
#endif

    /* We need to open the slave or reading from the master yields EOF after
     * the first wall write to it. */
    pts = open(name, O_RDWR | O_CLOEXEC);
    if (pts < 0) {
        perror(name);
        exit(EXIT_FAILURE);
    }

    /* Start a thread which connects to X11. This way we die if the user logs
     * out of its X11 session. */
    if (enable_x11) {
#ifdef DONT_USE_X11
        fprintf(stderr, "%s: option -X is disabled!\n\n", argv0);
        usage(argv0);
#else
        pthread_t tid;

        if (pthread_create(&tid, NULL, x11_event_loop_thread, NULL) != 0) {
            perror("pthread_create");
            exit(EXIT_FAILURE);
        }
#endif
    }

    /* Cleanup on signals. Necessary before login(). */
    setup_signals();

    if (!login(ptm)) {
        perror("login");
        exit(EXIT_FAILURE);
    }

    /* Main loop. Handle all wall messages sent to our PTY. */
    handle_wall(ptm, argv);

    if (!logout(ptm)) {
        perror("logout");
        exit(EXIT_FAILURE);
    }

    close(ptm);
    close(pts);

    return EXIT_SUCCESS;
}