/* 
 * Track new process on Linux
 * 2012 Laurent Ghigonis <laurent@p1sec.com>
 *
 * new version of trackproc.c
 * that was 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>
#include <event.h>
#include <libglouglou.h>

#define GG_TRACKPROC_USER "_gg_trackproc"
#define CHROOT_PATH   "/proc"
#define PROC_BASE     "/"

struct gg_client *ggcli;
struct event_base *ev_base;

int loglevel = 0;

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

void cb_nl(evutil_socket_t fd, short what, void *arg)
{
  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;
  static char *cmd;
  struct gg_packet pkt;

  rc = recv(fd, &nlcn_msg, sizeof(nlcn_msg), 0);
  if (rc == 0) {
    /* shutdown? */
    if (loglevel)
      printf("%s: stop listening to netlink events\n", 
          get_current_date());
    event_base_loopexit(ev_base, NULL);
    return;
  } else if (rc == -1) {
    if (errno == EINTR)
      return;
    perror("netlink recv");
    event_base_loopexit(ev_base, NULL);
    return;
  }

  pkt.type = 0;
  pkt.ver = PACKET_VERSION;
  /* see /usr/include/linux/cn_proc.h for struct proc_event */
  switch (nlcn_msg.proc_ev.what) {
  case PROC_EVENT_NONE:
    /* XXX what is this event for ?
     * if (loglevel)
     *   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)) {
      if (loglevel)
        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 {
      if (loglevel)
        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);
    }
    pkt.type = PACKET_FORK;
    pkt.fork_pid = nlcn_msg.proc_ev.event_data.fork.parent_pid;
    pkt.fork_ppid = nlcn_msg.proc_ev.event_data.fork.parent_pid;
    pkt.fork_cpid = nlcn_msg.proc_ev.event_data.fork.child_pid;
    pkt.fork_tgid = nlcn_msg.proc_ev.event_data.fork.child_tgid;
    break;

  case PROC_EVENT_EXEC:
    get_proc_infos(nlcn_msg.proc_ev.event_data.exec.process_pid,
        &cmd);
    if (loglevel)
      printf("%s: exec %d -> %s\n",
          get_current_date(),
          nlcn_msg.proc_ev.event_data.exec.process_pid,
          cmd);
    pkt.type = PACKET_EXEC;
    pkt.exec_pid = nlcn_msg.proc_ev.event_data.fork.parent_pid;
    pkt.exec_cmdlen = strnlen(cmd, GG_PKTARG_MAX);
    strncpy((char *)pkt.exec_cmd, cmd, GG_PKTARG_MAX);
    break;
#if 0
  case PROC_EVENT_UID:
    if (loglevel)
      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:
    if (loglevel)
      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;
#endif

  case PROC_EVENT_EXIT:
    if (loglevel)
      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);
    pkt.type = PACKET_EXIT;
    pkt.exit_pid = nlcn_msg.proc_ev.event_data.exit.process_pid;
    pkt.exit_tgid = nlcn_msg.proc_ev.event_data.exit.process_tgid;
    pkt.exit_ecode = nlcn_msg.proc_ev.event_data.exit.exit_code;
    break;

  default:
    if (loglevel)
      printf("unhandled proc event\n");
    break;
  }

  if (pkt.type != 0)
    gg_client_send(ggcli, &pkt);
}

static void on_sigint(int unused)
{
  event_base_loopexit(ev_base, NULL);
}

#if defined(__OPENBSD__)
void __dead
#else
void
#endif
usage(void)
{
  extern char *__progname;

  fprintf(stderr, "usage: %s [-vi]", __progname);
  exit(1);
}

int main(int argc, char **argv)
{
  struct event *ev_nl;
  int nl_sock;
  int rc = EXIT_SUCCESS;
	int op;

	while ((op = getopt(argc, argv, "hv")) != -1) {
		switch (op) {
			case 'h':
				usage();
				/* NOTREACHED */
			case 'v':
				loglevel++;
				break;
			default:
				usage();
				/* NOTREACHED */
		}
	}

	if (geteuid() != 0)
		errx(1, "must be root");

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

  gg_verbosity_set(loglevel);

  ev_base = event_base_new();
  ggcli = gg_client_connect(ev_base, "127.0.0.1", GLOUGLOU_PROBE_DEFAULT_PORT,
                    NULL, NULL, NULL);

  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 = set_proc_ev_listen(nl_sock, true);
  if (rc == -1) {
    rc = EXIT_FAILURE;
    goto out;
  }
  ev_nl = event_new(ev_base, nl_sock, EV_READ|EV_PERSIST, cb_nl, NULL);
  event_add(ev_nl, NULL);

  droppriv(GG_TRACKPROC_USER, 1, CHROOT_PATH);

  event_base_dispatch(ev_base);

  set_proc_ev_listen(nl_sock, false);

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