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


static void sig_handler(int signal) {
    (void)signal;
}
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 open_tty(void) {
    int ptm;
    const char *name;

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

    /* Prevent write access for other users so they can't use wall to send
     * messages to this program. */
    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;
}

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

    pid_t pid;

    /* Skip argv[0]. */
    argv++;

    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 */

        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 ((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);
    }
}


int main(int argc, char **argv) {
    int ptm, pts;
    char *name;

    if (argc < 2) {
        fprintf(stderr, "usage: %s <cmd args..>\n", argv[0]);
        exit(EXIT_FAILURE);
    }

    ptm = open_tty();
    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);
    if (pts < 0) {
        perror(name);
        exit(EXIT_FAILURE);
    }

    /* 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;
}