/*
 * Patrick MARIE <mycroft@virgaria.org>
 * Licensed under BSD
 * http://www.opensource.org/licenses/bsd-license.php
 */
#include <sys/types.h>
#include <sys/stat.h>
#include <dirent.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <unistd.h>

#define PROC_DIR "/proc"
#define PROC_PID_MAX "/proc/sys/kernel/pid_max"

#define DELAY 5 

/* 1000 signifie: si un process bouffe 10 (1000 / 100) secondes CPU 
durant un laps de temps ou le dit process est Running sans arret,
on le killera. 
 */
#define MAX_CPUTIME DELAY * 4 * 100 

/* signal envoyé */
#define USE_SIGNAL 9 

/* Il faudra 20 cycles de DELAY secondes avant que le process
   soit retiré du tableau */
#define MAX_EXPIRE 4  

/*
 * Si on kill un process de nom "blabla" d'un user "lambda",
 * killer aussi tout les process "blabla" de cet user.
 *
 * (Commenter pour désactiver)
 */
#define FEATURE_KILLALL

/*
 * On monitore tout les process, et pas que les process 'R'unning
 * Si cette option est activée, penser à augmenter le nombre d'entrées
 * dans le tableau (TABLEAU_SIZE)
 * (Commenter pour désactiver)
 */
#define FEATURE_ALL_PROC

#ifdef FEATURE_KILLALL
#define TABLEAU_SIZE 2048 
#else
#define TABLEAU_SIZE 64
#endif

#define DEBUG

int monitored = 0;
int next_free = 0;
pid_t pid_max = 0;

typedef struct proc_load {
  pid_t pid;
  uid_t uid;
  char *proc_name;
  unsigned long cputime;
  time_t lasttime;
  int matched;
} plopstruct;

typedef struct {
  int proc_pid;
  char proc_name[129];
  char proc_state;
  unsigned long proc_utime, proc_stime;
  uid_t proc_uid;
} proc_read;

/* There should be enought ... */
plopstruct tableau[TABLEAU_SIZE];

/* Prototypes */
int die(const char *str);
int killall(char *name, uid_t uid);
int addproc(pid_t pid, uid_t uid, char *name, unsigned long temps_cpu);
int watchproc(void);
proc_read* get_process(int pid);
int hand(int pid);
int getinfo(void);

int
die(const char *str)
{
  fprintf(stderr, str);
  exit(-1);
}

int
killall(char *name, uid_t uid)
{
  DIR *proc;
  struct dirent *rep;
  proc_read *tmp;

  proc = opendir(PROC_DIR);
  if(proc == NULL) {
    perror("opendir");
    die("No " PROC_DIR " available.\n");
  }

  while((rep = readdir(proc)) != NULL) {
    if(*rep->d_name < '0' || *rep->d_name > '9')
      continue;
    tmp = get_process(atoi(rep->d_name));
    if(tmp == NULL)
      continue; /* May be I'll have more luck next time. */
    if(strcmp(tmp->proc_name, name) == 0 && tmp->proc_uid == uid) {
      kill(atoi(rep->d_name), USE_SIGNAL);
    }
  }

  closedir(proc);

  return 0;
}

int
addproc(pid_t pid, uid_t uid, char *name, unsigned long temps_cpu)
{
  div_t divi;
  int i, j;
  time_t nowtime = time(NULL);

  if(monitored == 0) {
    tableau[next_free].pid = pid;
    tableau[next_free].uid = uid;
    tableau[next_free].proc_name = strdup(name);
    tableau[next_free].lasttime = nowtime;
    tableau[next_free].cputime = temps_cpu;
    tableau[next_free].matched = 1;
    next_free ++;
    monitored ++;
  } else {
    j = monitored; i = 0;
    while(i < j && j < TABLEAU_SIZE) {
      if(tableau[i].pid == 0) {
        j++; i++; continue;
      }
      if(tableau[i].pid == pid) {
        /* Update ... */
        divi = div((temps_cpu - tableau[i].cputime),
                   (nowtime - tableau[i].lasttime));
        tableau[i].matched ++;

#ifdef DEBUG
        fprintf(stderr,
                "proc: %d - cpu: %lu deb: %lu - %lu (max: %lu) - %lu to %lu\n",
                tableau[i].pid,
                temps_cpu,
                tableau[i].cputime,
                temps_cpu - tableau[i].cputime,
                MAX_CPUTIME,
                tableau[i].lasttime,
                nowtime);

        fprintf(stderr, 
                "proc: %s - use: %d.%d (on %d sec.)\n", 
                tableau[i].proc_name, 
                divi.quot, 
                divi.rem,
                nowtime - tableau[i].lasttime);
#endif

        /*
         * printf("update: %d %d\n", divi.quot, divi.rem);
         */
        if((temps_cpu - tableau[i].cputime) >= MAX_CPUTIME) {
#ifdef FEATURE_KILLALL
          fprintf(stderr,
                  "Use signal %d against process %s owned by %d\n",
                  USE_SIGNAL,
                  tableau[i].proc_name,
                  tableau[i].uid);
          killall(tableau[i].proc_name, tableau[i].uid);
#else
          fprintf(stderr, 
                  "Use signal %d against process %d (%s) owned by %d\n", 
                  USE_SIGNAL, 
                  pid,
                  tableau[i].proc_name,
                  tableau[i].uid);
          kill(pid, USE_SIGNAL);
#endif
        }
        break;
      }
      i++;
    }

    /* Ajout d'une nouvelle structure si i == monitored */
    if(i == monitored && monitored < TABLEAU_SIZE) {
      tableau[next_free].pid = pid;
      tableau[next_free].uid = uid;
      tableau[next_free].proc_name = strdup(name);
      tableau[next_free].lasttime = nowtime;
      tableau[next_free].cputime = temps_cpu;
      tableau[next_free].matched = 1;
      monitored ++;
      /* Recherche d'un nouveau next_free ... */
      for(i = next_free + 1; i < TABLEAU_SIZE; i++) {
        if(tableau[i].pid == 0) {
          next_free = i;
          break;
        }
      }
    } 
  }

  return 0;
}

int
watchproc(void)
{
  int i, j;
  time_t nowtime = time(NULL);

  if(monitored == 0)
    return 0;

  j = monitored;
  for(i = 0; i < j && j < TABLEAU_SIZE; i ++) {
    if(tableau[i].pid == 0) {
      j++; continue;
    }

    /* Check if the process is not anymore in high system use */
    if((nowtime - tableau[i].lasttime) > DELAY * MAX_EXPIRE) {
      /* Expiration */
      tableau[i].pid = 0;
      free(tableau[i].proc_name);
      monitored --;
      if(i < next_free)
        next_free = i;
/*
      if(i == monitored)
        break;
*/
    }
  }

  return 0;
}

proc_read* 
get_process(int pid)
{
  char buf[1024];
  int ret, fd;
  struct stat stbuf;

  proc_read *tmp = malloc(sizeof(proc_read));
  if(tmp == NULL) /* arg ... */
    return(NULL);

  char proc_name[129]; 
  int proc_pid;
  char proc_state;
  unsigned long proc_utime, proc_stime;

  snprintf(buf, 32, "/proc/%d/stat", pid);
  fd = open(buf, O_RDONLY, 0);
  if(fd == -1) {
    perror("open");
    free(tmp);
    return(NULL);
  }

  ret = stat(buf, &stbuf); 
  if(ret == -1) {
    perror("stat");
    free(tmp);
    return(NULL);
  }

  ret = read(fd, buf, sizeof buf - 1);
  close(fd);
  if(ret < 42) {
    free(tmp);
    return(NULL);
  }

  ret = sscanf(buf,
     "%d %128[^ ] "
     "%c "
     "%*d %*d %*d %*d %*d "
     "%*lu %*lu %*lu %*lu %*lu %lu %lu "
     "%*ld %*ld %*ld %*ld %*ld %*ld "
     "%*lu %*lu "
     "%*ld "
     "%*lu %*lu %*lu %*lu %*lu %*lu "
     "%*u %*u %*u %*u "
     "%*lu %*lu %*lu",
     &proc_pid, proc_name,
     &proc_state,
     &proc_utime, &proc_stime
  );

  strcpy(tmp->proc_name, proc_name);
  tmp->proc_pid = proc_pid;
  tmp->proc_state = proc_state;
  tmp->proc_utime = proc_utime;
  tmp->proc_stime = proc_stime;
  tmp->proc_uid = stbuf.st_uid;

  return(tmp);
}

int
hand(int pid)
{
  proc_read *tmp = get_process(pid);

#ifdef FEATURE_ALL_PROC
  if(tmp != NULL && tmp->proc_uid != 0) {
#else
  if(tmp != NULL && tmp->proc_state == 'R' && tmp->proc_uid != 0) {
#endif
/*
    printf("%i (%c)", ret, proc_state);
    printf("temps cpu du %i: %lu\n", proc_pid, proc_utime);
 */
    addproc(tmp->proc_pid, tmp->proc_uid, tmp->proc_name, tmp->proc_utime);
  }

  if(tmp != NULL)
    free(tmp);

  return 0;
}

int
getinfo(void)
{
  DIR *proc;
  struct dirent *rep;

  proc = opendir(PROC_DIR);
  if(proc == NULL) {
    perror("opendir");
    die("No " PROC_DIR " available.\n");
  }

  while((rep = readdir(proc)) != NULL) {
    if(*rep->d_name < '0' || *rep->d_name > '9')
      continue;
    hand(atoi(rep->d_name));
  }

  closedir(proc);

  return 0;
}

int
main(void)
{
  memset(&tableau, 0, sizeof(tableau));

  while(sleep(DELAY) == 0) {
#ifdef DEBUG
    fprintf(stderr, "-- watch --\n");
#endif
    getinfo();
    watchproc();
  }

  return 0;
}