#!/usr/bin/perl -w
#
# Copyright (c) 2000, 2005, 2009  Peter Pentchev
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.
#
# $Ringlet: pslist 3732 2009-06-01 11:54:51Z roam $

use strict;

my($PS, $PSflags) = ("/bin/ps", "axco pid,ppid,command");
my(%proc) = ();
my(%parproc) = ();

# Function:
#	help			- display usage information
# Inputs:
#	$err			- error code; non-zero for program termination
# Returns:
#	nothing; calls exit() if $err is non-zero
# Modifies:
#	nothing; writes to STDOUT
#	If $err is non-zero, writes to STDERR and exits.

sub help {
  my($err) = @_;
  my($s) =
    "Usage:\n"
    ."  pslist           [pid/name...]    - list child pids\n"
    ."  rkill   [-sig]    pid/name...     - kill a process and its children\n"
    ."  rrenice [+/-]pri  pid/name...     - renice a process and its children\n"
    ."\n"
    ."The priority for rrenice may be specified as an absolute value, or as\n"
    ."a +/- delta to current process priority.\n"
    ."At least one pid/name argument is mandatory for rkill and rrenice.\n"
    ;

  if ($err) {
    print STDERR $s;
    exit($err);
  } else {
    print $s;
  }
}

# Function:
#	version			- output program version information
# Inputs:
#	none
# Returns:
#	nothing
# Modifies:
#	nothing; writes to STDOUT

sub version {
  print "pslist 1.3 - control processes and their descendants\n";
}

# Function:
#	proc_gather		- parse ps output to fill out process arrays
# Inputs:
#	none
# Returns:
#	nothing
# Modifies:
#	fills out %proc and %parproc
#	invokes a pipe to 'ps'

sub proc_gather {
  my($line);

  open(PS, "$PS $PSflags |") or die("failed to invoke '$PS $PSflags' - $!\n");
  while(defined($line = <PS>)) {
    chomp $line;
    if ($line =~ /^\s*(\d+)\s+(\d+)\s+(\S+)(.*)$/) {
      my($pid, $ppid, $cmd, $args) = ($1, $2, $3, $4);

      $proc{$pid} = {'ppid'=>$ppid, 'cmd'=>$cmd, 'args'=>$args};
      $parproc{$ppid} .= "$pid ";
    }
  }
  close PS;
}

# Function:
#	proc_display		- display a process and its children
# Inputs:
#	$pid			- process ID to display
# Returns:
#	nothing
# Modifies:
#	nothing by itself; writes to STDOUT and calls proc_get_children_r()

sub proc_display {
  my($pid) = shift;
  my(@arr);

  @arr = proc_get_children_r($pid);
  print "$pid $proc{$pid}{'cmd'} @arr\n";
}

# Function:
#	proc_display_all	- display all processes and children lists
# Inputs:
#	none
# Returns:
#	nothing
# Modifies:
#	nothing by itself; calls proc_display()

sub proc_display_all {
  my($pid);

  foreach $pid (keys(%parproc)) {
    proc_display($pid);
  }
}

# Function:
#	proc_kill		- recursively kill a process and its children
# Inputs:
#	$pid			- PID to kill
#	$sig			- signal to send
# Returns:
#	0 on success
#	negative number of unkilled children on failure, $! is set

sub proc_kill {
  my($pid, $sig) = @_;
  my(@arr);

  die("bad pid ($pid)\n") if ($pid !~ /^\d+$/);
  die("non-existent pid ($pid)\n") unless defined($proc{$pid});
  
  $arr[0] = $pid;
  push(@arr, proc_get_children_r($pid));
  print "@arr ";
  return kill($sig, @arr) - ($#arr + 1);
}

# Function:
#	proc_nice		- recursively renice a process and its children
# Inputs:
#	$pid			- PID to kill
#	$sig			- signal to send
# Returns:
#	0 on success
#	negative number of unkilled children on failure, $! is set

sub proc_nice {
  my($pid, $delta) = @_;
  my(@arr);
  my($kpid, $exact) = (0, 1);

  die("bad pid ($pid)\n") if ($pid !~ /^\d+$/);
  die("non-existent pid ($pid)\n") unless defined($proc{$pid});
  
  $exact = 0 if ($delta =~ /^[+-]/);

  $arr[0] = $pid;
  push(@arr, proc_get_children_r($pid));
  print "@arr ";

  foreach $kpid (@arr) {
    if ($exact) {
      setpriority(0, $kpid, $delta) or return -1;
    } else {
      my($cpri) = getpriority(0, $kpid);
      setpriority(0, $kpid, $cpri + $delta) or return -1;
    }
  }

  return 0;
}

# Function:
#	proc_get_children_r	- get a list of PIDs of a process's child tree
# Inputs:
#	$pid			- PID to examine
# Returns:
#	array of children PIDs
# Modifies:
#	nothing; calls proc_get_children()

sub proc_get_children_r {
  my($pid, $i) = (shift || 0, 0);
  my(@chi) = ();
  my(@res) = ();

  @chi = proc_get_children($pid);
  return () if ($#chi == -1);

  for($i = 0; $i <= $#chi; $i++) {
    next if $chi[$i] == 0;
    $res[++$#res] = $chi[$i];
    push(@res, proc_get_children_r($chi[$i]));
  }
  return @res;
}

# Function:
#	proc_get_children	- get a list of a process's immediate children
# Inputs:
#	$pid			- process ID to examine
# Returns:
#	array of children PIDs
# Modifies:
#	nothing

sub proc_get_children {
  my($pid) = (shift || 0);
  my(@arr) = ();
  my($s);
  
  return () unless defined($parproc{$pid});
  
  $s = $parproc{$pid};
  while($s =~ /^(\d+) (.*)/) {
    $arr[++$#arr] = $1;
    $s = $2;
  }

  return @arr;
}

# Main block

MAIN:{
  my($killem, $niceem, $killpid) = (0, 0, 0);
  my($sig) = 15;

  # check for help, version
  if ($#ARGV > -1) {
    if (($ARGV[0] eq '-h') || ($ARGV[0] eq '--help')) {
      version(); help(0); exit(0);
    }
    if (($ARGV[0] eq '-v') || ($ARGV[0] eq '--version')) {
      version(); exit(0);
    }
  }
  
  # are we invoked as pslist, rkill, or rrenice?
  if ($0 =~ /kill$/) {
    help(1) if ($#ARGV == -1);
    if ($ARGV[0] =~ /^-(.*)/) {
      help(1) if ($#ARGV == 0);
      $sig = $1;
      shift @ARGV;
    }
    $killem = 1;
  } elsif ($0 =~ /nice$/) {
    help(1) if (($#ARGV < 1) || (($ARGV[0] !~ /^[+-]?\d+$/)));
    $sig = $ARGV[0];
    shift @ARGV;
    $niceem = 1;
  }

  # Let the user override the ps program location and flags
  $PS = $ENV{'PS'} if (defined($ENV{'PS'}));
  $PSflags = $ENV{'PSflags'} if (defined($ENV{'PSflags'}));
  proc_gather();

  # no arguments, no kill requested - display all and exit
  if (($#ARGV == -1) && !($killem || $niceem)) {
    exit(proc_display_all());
  }
 
  # either a kill request, or specific processes display
  foreach $killpid (@ARGV) {
    # pid or process name?
    if ($killpid =~ /^\d+$/) {
      # pid..
      die("nonexistent pid $killpid\n") unless defined($proc{$killpid});
      if ($killem) {
    	(proc_kill($killpid, $sig) == 0) || die("rkill($killpid) failed - $!\n");
      } elsif ($niceem) {
	(proc_nice($killpid, $sig) == 0) || die("rrenice($killpid) failed - $!\n");
      } else {
	proc_display($killpid);
      }
    } else {
      # process name..
      my($pid);
      foreach $pid (keys(%proc)) {
	if ($proc{$pid}{'cmd'} eq $killpid) {
	  if ($killem) {
	    (proc_kill($pid, $sig) == 0) || die("rkill($pid) failed - $!\n");
	  } elsif ($niceem) {
	    (proc_nice($pid, $sig) == 0) || die("rrenice($pid) failed - $!\n");
	  } else {
	    proc_display($pid);
	  }
	}
      }
    }
  }
  print "\n" if ($killem);
}
