#!/usr/bin/perl
# bgcheck 0.5 by blue [blue@altavista.net]

# i wrote this as neat as i could; please be very careful running this
# script in making sure that you added the user that runs your webserver
# (usually "nobody", which is already defaultly on the exceptions list)
# and anything else that is running more applications than defined in
# this script and needs to be excused.  the best way to find these is to
# edit the configuration (below) and set $test to 1; what this will do is
# perform the script and output what would have happened if it were
# actually running, but not take action.  have fun, and play it safe.

# NOTE: the method that bgcheck uses to select the processes to kill is
# that it selects the ones that are using the most memory.

# props to my mentor/guru anut [nads@bleh.org] for moral/tech support

# ,-----------------------------------------.
# |       C O N F I G U R A T I O N         |
# |-----------------------------------------|
# | the arguments to give ps, do not change |
$ps_args = "auxw";
# |-----------------------------------------|
# | maximum processes to allow              |
$maximum = 2;
# |-----------------------------------------|
# | default permissions for bgcheck         |
$perms = 0700;
# |-----------------------------------------|
# | the username that handles ftpd procs    |
$ftp_user = "ftp";
# |-----------------------------------------|
# | what ftpd appears to be according to ps |
$ftp_server = "ftpd";
# |-----------------------------------------|
# | exception list, play the favorites      |
# | do not remove root; you were warned     |
@except = ("root","nobody","bin","$ftp_user");
# |-----------------------------------------|
# | program exception list (case sensitive) |
@except_program = ();
# |-----------------------------------------|
# | enable this if you have long usernames  |
# | (over eight characters)                 |
$long_usernames = 0;
# |-----------------------------------------|
# | enable this if you want bgcheck to ask  |
# | you before killing processes            |
$ask = 0; # 0=disabled, 1=enabled
# |-----------------------------------------|
# | enabling this will mail the user each   |
# | time they exceed the maximum processes  |
# | NOTE: you must have sendmail for this!  |
$mail = 1; # 0=disabled, 1=enabled
# |-----------------------------------------|
# | path to the sendmail binary             |
# | ($mail/$violate = 1)                    |
$sendmail = "/usr/sbin/sendmail";
# |-----------------------------------------|
# | shown in the From: field in e-mail      |
# | ($mail/$violate = 1)                    |
$from = "bgcheck <root>";
# |-----------------------------------------|
# | closing of the e-mail                   |
# | ($mail/$violate = 1)                    |
$closing = "Yours truly,";
# |-----------------------------------------|
# | number of violations before notifying   |
# | the admin (via email), 0=for never      |
$violate = 3; # 0 = never
# |-----------------------------------------|
# | the file to keep track of each user's   |
# | violations ($violate = 1+)              |
$violate_db = "/root/.bgcheck.db";
# |-----------------------------------------|
# | default permissions for $violate_db     |
$db_perms = 0600;
# |-----------------------------------------|
# | the admin's e-mail ($violate > 0)       |
$admin = "root";
# |-----------------------------------------|
# | enable this if you want to simply test, |
# | i.e. it'll output but not take action   |
$test = 0; # 0=disabled, 1=enabled
# |-----------------------------------------|
# | -max #, -ask, -no-ask, -test, -no-test  |
# | -mail, -no-mail                         |
# | all also work as parameters to bgcheck  |
# `-----------------------------------------'

$version = "0.5";

chmod ($perms, $0);

open (MEMINFO, "/proc/meminfo");
while (<MEMINFO>) {
  if (/^MemTotal\:\s*(\d+)\skB/) {
    $kb_of_ram = $1;
  }
}
close (MEMINFO);

foreach ($i = 0; $_ = $ARGV[$i]; $i++) {
  if (/^-(max|maximum)$/) {
    $maximum = $ARGV[$i+1];
    splice (@ARGV, $i, 2);
    $i -= 1;
  }
  if (/^-violate$/) {
    $violate = $ARGV[$i+1];
    splice (@ARGV, $i, 2);
    $i -= 1;
  }
  $test = 1 if (/^-test$/);
  $ask = 1 if (/^-ask$/);
  $mail = 1 if (/^-mail$/);
  $test = 0 if (/^-no-test$/);
  $ask = 0 if (/^-no-ask$/);
  $mail = 0 if (/^-no-mail$/);
  $violate = 0 if (/^-no-violate$/);
  die "bgcheck $version [http://blue.dhs.org/bgcheck]\nauthor: blue (blue\@altavista.net)\n" if (/^\-(v|V)$/);
}

print "bgcheck ${version}; checking processes...\n";
$except_list = join(",", @except);
$except_prog_list = join(",", @except_program);
print "excepting users: $except_list\n" if ($except_list);
print "excepting progs: $except_prog_list\n" if ($except_prog_list);

$ps_args .= "n" if ($long_usernames);

open (PS, "ps $ps_args|");
while (<PS>) {
  chomp();
  if (/^\s*(\S+)\s*(\d+)\s*\S+\s*(\S+)\s*\S+\s*\S+\s*(\S+)\s*(\S+\s\S+|\S+)\s*(\S+\s\S+|\S+)\s*\S+\s*(.*)/ && $4 eq "?" && check_except($1) && check_except_program($7)) {
    my ($cmdline) = remove_spaces($7); my ($user) = $1; my ($pid) = $2; my ($mem) = $3;
    if ($long_usernames) {
      ($user) = getpwuid($user);
      next unless (check_except($user));
    }
    next if ($cmdline =~ /^$ftp_server\: .*\:/);
    my ($truncmd) = $cmdline; $truncmd =~ s/^(.{25}).*/\1/;
    $pid{$pid} = $truncmd;
    if ($data{$user} ne "") {
      $data{$user} .= ",${pid}:${mem}";
    } else {
      $data{$user} = "${pid}:${mem}";
    }
  }
}
close (PS);

foreach $this_user (keys %data) {
  my (@pids) = split (/,/, $data{$this_user}); my ($kill) = 1;
  next if (@pids <= $maximum);
  my ($orig_pids) = $#pids+1;
  print "$this_user is running $orig_pids (max=${maximum}) background process(es)\n";
  for (1 ... $maximum) {
    my ($splice_me) = get_lowest_mem(@pids);
    splice (@pids, $splice_me, 1);
  }
  my ($excess) = $#pids+1;
  if ($ask) {
    print "kill the excessive ($excess) background processes [y/n]? ";
    chomp($yn = <STDIN>);
    if ("\L$yn\E" eq "y") {
      $kill = 1;
    } else {
      $kill = 0;
    }
  }
  my ($save_mem) = 0; my (@proc_list);
  foreach $pid (@pids) {
    my ($pid, $mem) = split(/:/, $pid);
    $save_mem += $mem;
    my ($this_mem) = int(($mem / 100) * $kb_of_ram);
    if ($kill) {
      print "++ killing process $pid ($pid{$pid}); mem=${mem}% (${this_mem}kB)\n";
      if (! $test) {
        kill(9, $pid);
        push(@proc_list, "$pid{$pid} (pid: ${pid})");
      }
    } else {
      print "++ not killing process $pid ($pid{$pid}); mem=${mem}% (${this_mem}kB)\n";
    }
  }
  $saved = int(($save_mem / 100) * $kb_of_ram);
  if ($kill) {
    print "++ saved: ${save_mem}% (${saved}kB) mem\n";
    if ($mail && ! $test) {
      open (SENDMAIL, "|$sendmail $this_user");
      print SENDMAIL <<EOF;
From: $from
To: $this_user
Subject: Excessive processes ($orig_pids)

You have exceeded minimum of $maximum background processes by $excess.
In result, we have kill(1)'ed the following:

EOF
      foreach $proc (@proc_list) {
        print SENDMAIL "  * $proc\n";
      }
      print SENDMAIL <<EOF;

Please refrain from this in the future.  These implications
are merely to sustain system resources, and to ensure that
we can maintain a stable machine for our users.

Thank you.

$closing
bgcheck [http://blue.dhs.org/bgcheck]
EOF
      close (SENDMAIL);
      print "++ sent e-mail to $this_user as a reminder\n";
    }
    if ($violate > 0 && $violate_db && ! $test) {
      my (@violators,$done,$this_v);
      open (DB, $violate_db);
      while (<DB>) {
        if (/^(\S+)\:(\d+)/) {
          if ($1 eq $this_user) {
            $this_v = $2 + 1; $done = 1;
            if ($this_v >= $violate) {
              open (SENDMAIL, "|$sendmail $admin");
              print SENDMAIL <<EOF;
From: $from
To: $admin
Subject: ${this_user}'s violations >= $violate

Hello.  I'm the perl script you entrusted to manage processes
on your machine.  How are you?  I'm here to report to you what
you requested of me, but don't get too grumpy.  It won't help.

'${this_user}' seems to be running too many processes (over $maximum)
too many times ($violate)!  This needs to be dealt with, because I
keep seeing his name being mentioned in my reports.  Anyway,
thank you for listening.  I'm here, as always.

$closing
bgcheck [http://blue.dhs.org/bgcheck]
EOF
              close (SENDMAIL);
            } else {
              push (@violators, "$1:$this_v");
            }
          } else {
            push (@violators, "$1:$2");
          }
        }
      }
      close (DB);
      if (! $done) {
        push (@violators, "$this_user:1");
        $this_v = 1;
      }
      open (DB, ">$violate_db");
      foreach $v (sort @violators) {
        print DB "$v\n";
      }
      close (DB);
      if ($this_v >= $violate) {
        print "++ $this_user: $this_v violations (reset! + e-mailed admin)\n";
      } else {
        print "++ $this_user: $this_v violations\n";
      }
    }
  } else {
    print "++ [would have] saved: ${save_mem}% (${saved}kB) mem\n";
  }
}

chmod ($db_perms, $violate_db) if (-f $violate_db);

sub check_except {
  my ($user) = @_; my ($except) = 1;
  foreach $e (@except) {
    if ($user eq $e) {
      $except = 0; last;
    }
  }
  return $except;
}

sub check_except_program {
  my ($program) = @_; my ($except) = 1;
  $program = remove_spaces($program); $program =~ s/^(\S+).*/\1/;
  foreach $e (@except_program) {
    if ($program eq $e) {
     $except = 0; last;
    }
  }
  return $except;
}

sub get_lowest_mem {
  my (@pids) = @_; my($l_mem) = 1000000;
  for ($i = 0; $_ = $pids[$i]; $i++) {
    my ($pid, $mem) = split(/:/, $_);
    if ($mem < $l_mem) {
      $l_mem = $mem;
      $l_pid = $i;
    }
  }
  return $l_pid;
}

sub remove_spaces {
  my ($str) = @_;
  while ($str =~ s/^\s//) {}
  while ($str =~ s/\s$//) {}
  while ($str =~ s/^\(//) {}
  while ($str =~ s/\)$//) {}
  return $str;
}
