/* 
 * Track new process on Linux
 * 2012 Laurent Ghigonis <laurent@p1sec.com>
 *
 * inspired from http://bewareofgeek.livejournal.com/2945.html
 * and pstree sources
 */

#include <sys/socket.h>
#include <linux/netlink.h>
#include <linux/connector.h>
#include <linux/cn_proc.h>
#include <signal.h>
#include <err.h>
#include <errno.h>
#include <stdbool.h>
#include <unistd.h>
#include <string.h>
#include <stdlib.h>
#include <stdio.h>
#include <time.h>

#define PROC_BASE    "/proc"

static char*
get_current_date() {
    time_t ts;
    struct tm *tm;
    static char date_buf[64];

    time(&ts);
    tm = localtime(&ts);
    strftime(date_buf, 64, "%b %d %X", tm);
    return date_buf;
}

static int
get_proc_infos(int pid, char **cmd) {
    char path[128];
    char readbuf[1024];
    int size;
    FILE *file;
    char *comm, *tmpptr;

    //if ((pid = (pid_t) atoi(de->d_name)) == 0) {
    snprintf(path, sizeof(path), "%s/%d/stat", PROC_BASE, pid);
    if ((file = fopen(path, "r")) == NULL) {
        warn("Error: PID %d: cannot open %s\n", pid, path);
        return -1;
    }
    size = fread(readbuf, 1, sizeof(readbuf), file);
    if (ferror(file) != 0) {
        warn("Error: PID %d: error reading %s", pid, path);
        return -1;
    }
    readbuf[size] = 0;
    /* commands may have spaces or ) in them.
     * so don't trust anything from the ( to the last ) */
    if (!(comm = strchr(readbuf, '('))
            || !(tmpptr = strrchr(comm, ')'))) {
        warn("Error: PID %d: could not parse %s", pid, path);
        return -1;
    }
    ++comm;
    *tmpptr = 0;
    /* We now have readbuf with pid and cmd, and tmpptr+2
     * with the rest */
    fclose(file);

    *cmd = comm;
    return 0;
}

/*
 * connect to netlink
 * returns netlink socket, or -1 on error
 */
static int nl_connect()
{
    int rc;
    int nl_sock;
    struct sockaddr_nl sa_nl;

    nl_sock = socket(PF_NETLINK, SOCK_DGRAM, NETLINK_CONNECTOR);
    if (nl_sock == -1) {
        perror("socket");
        return -1;
    }

    sa_nl.nl_family = AF_NETLINK;
    sa_nl.nl_groups = CN_IDX_PROC;
    sa_nl.nl_pid = getpid();

    rc = bind(nl_sock, (struct sockaddr *)&sa_nl, sizeof(sa_nl));
    if (rc == -1) {
        perror("bind");
        close(nl_sock);
        return -1;
    }

    return nl_sock;
}

/*
 * subscribe on proc events (process notifications)
 */
static int set_proc_ev_listen(int nl_sock, bool enable)
{
    int rc;
    struct __attribute__ ((aligned(NLMSG_ALIGNTO))) {
        struct nlmsghdr nl_hdr;
        struct __attribute__ ((__packed__)) {
            struct cn_msg cn_msg;
            enum proc_cn_mcast_op cn_mcast;
        };
    } nlcn_msg;

    memset(&nlcn_msg, 0, sizeof(nlcn_msg));
    nlcn_msg.nl_hdr.nlmsg_len = sizeof(nlcn_msg);
    nlcn_msg.nl_hdr.nlmsg_pid = getpid();
    nlcn_msg.nl_hdr.nlmsg_type = NLMSG_DONE;

    nlcn_msg.cn_msg.id.idx = CN_IDX_PROC;
    nlcn_msg.cn_msg.id.val = CN_VAL_PROC;
    nlcn_msg.cn_msg.len = sizeof(enum proc_cn_mcast_op);

    nlcn_msg.cn_mcast = enable ? PROC_CN_MCAST_LISTEN : PROC_CN_MCAST_IGNORE;

    rc = send(nl_sock, &nlcn_msg, sizeof(nlcn_msg), 0);
    if (rc == -1) {
        perror("netlink send");
        return -1;
    }

    return 0;
}

/*
 * handle a single process event
 */
static volatile bool need_exit = false;
static int handle_proc_ev(int nl_sock)
{
    int rc;
    struct __attribute__ ((aligned(NLMSG_ALIGNTO))) {
        struct nlmsghdr nl_hdr;
        struct __attribute__ ((__packed__)) {
            struct cn_msg cn_msg;
            struct proc_event proc_ev;
        };
    } nlcn_msg;

    while (!need_exit) {
        rc = recv(nl_sock, &nlcn_msg, sizeof(nlcn_msg), 0);
        if (rc == 0) {
            /* shutdown? */
            printf("%s: stop listening to netlink events\n", 
                   get_current_date());
            return 0;
        } else if (rc == -1) {
            if (errno == EINTR) continue;
            perror("netlink recv");
            return -1;
        }
        switch (nlcn_msg.proc_ev.what) {
            static char *cmd;

            /* see /usr/include/linux/cn_proc.h for struct proc_event */
            case PROC_EVENT_NONE:
                printf("%s: start listening to netlink events...\n",
                       get_current_date());
                break;

            case PROC_EVENT_FORK:
                get_proc_infos(nlcn_msg.proc_ev.event_data.fork.parent_pid,
                               &cmd);
                if (nlcn_msg.proc_ev.event_data.fork.child_tgid
                    != nlcn_msg.proc_ev.event_data.fork.child_pid)
                  printf("%s: fork %s %d -> %d tid %d\n",
                      get_current_date(),
                      cmd,
                      nlcn_msg.proc_ev.event_data.fork.parent_pid,
                      nlcn_msg.proc_ev.event_data.fork.child_pid,
                      nlcn_msg.proc_ev.event_data.fork.child_tgid);
                else
                  printf("%s: fork %s %d -> %d\n",
                      get_current_date(),
                      cmd,
                      nlcn_msg.proc_ev.event_data.fork.parent_pid,
                      nlcn_msg.proc_ev.event_data.fork.child_pid);
                break;

            case PROC_EVENT_EXEC:
                get_proc_infos(nlcn_msg.proc_ev.event_data.exec.process_pid,
                               &cmd);
                printf("%s: exec %d -> %s\n",
                       get_current_date(),
                       nlcn_msg.proc_ev.event_data.exec.process_pid,
                       cmd);
                break;
#if 0
            case PROC_EVENT_UID:
                printf("uid change: tid=%d pid=%d from %d to %d\n",
                        nlcn_msg.proc_ev.event_data.id.process_pid,
                        nlcn_msg.proc_ev.event_data.id.process_tgid,
                        nlcn_msg.proc_ev.event_data.id.r.ruid,
                        nlcn_msg.proc_ev.event_data.id.e.euid);
                break;

            case PROC_EVENT_GID:
                printf("gid change: tid=%d pid=%d from %d to %d\n",
                        nlcn_msg.proc_ev.event_data.id.process_pid,
                        nlcn_msg.proc_ev.event_data.id.process_tgid,
                        nlcn_msg.proc_ev.event_data.id.r.rgid,
                        nlcn_msg.proc_ev.event_data.id.e.egid);
                break;

            case PROC_EVENT_EXIT:
                printf("exit: tid=%d pid=%d exit_code=%d\n",
                        nlcn_msg.proc_ev.event_data.exit.process_pid,
                        nlcn_msg.proc_ev.event_data.exit.process_tgid,
                        nlcn_msg.proc_ev.event_data.exit.exit_code);
                break;

            default:
                printf("unhandled proc event\n");
                break;
#endif
        }
    }

    printf("%s: stop listening to netlink events\n", 
           get_current_date());
    return 0;
}

static void on_sigint(int unused)
{
    need_exit = true;
}

int main(int argc, const char *argv[])
{
    int nl_sock;
    int rc = EXIT_SUCCESS;

    signal(SIGINT, &on_sigint);
    siginterrupt(SIGINT, true);

    nl_sock = nl_connect();
    if (nl_sock == -1)
        exit(EXIT_FAILURE);

    rc = set_proc_ev_listen(nl_sock, true);
    if (rc == -1) {
        rc = EXIT_FAILURE;
        goto out;
    }

    rc = handle_proc_ev(nl_sock);
    if (rc == -1) {
        rc = EXIT_FAILURE;
        goto out;
    }

    set_proc_ev_listen(nl_sock, false);

out:
    close(nl_sock);
    exit(rc);
}