#!/usr/bin/perl
#
# Copyright (c) 2000, 2005, 2009, 2016, 2018  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.

use v5.010;
use strict;
use warnings;

use constant VERSION => '1.4.0';

my ($defaultPS, $defaultPSflags) = ("/bin/ps", "axco pid,ppid,command");
my ($PS, $PSflags);
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) = @_;
	$err //= 1;

	my $s = <<'EOUSAGE' ;
Usage:
  pslist           [pid/name...]    - list child pids
  rkill   [-sig]    pid/name...     - kill a process and its children
  rrenice [+/-]pri  pid/name...     - renice a process and its children

The priority for rrenice may be specified as an absolute value, or as
a +/- delta to current process priority.
At least one pid/name argument is mandatory for rkill and rrenice.
EOUSAGE

	if ($err) {
		die $s;
	} else {
		print $s;
	}
}

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

sub version() {
	say 'pslist '.VERSION;
}

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

sub features() {
	say "Features: pslist=".VERSION;
}

# 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() {
	open my $ps, "$PS $PSflags |" or die "failed to invoke '$PS $PSflags' - $!\n";
	while (defined(my $line = <$ps>)) {
		chomp $line;
		if ($line =~ /^
		    \s* (?<pid>\d+) \s+
		    (?<ppid>\d+) \s+
		    (?<cmd>\S+)
		    (?<args>.*)
		    $/x) {
			$proc{$+{pid}} = {
				ppid => $+{ppid},
				cmd => $+{cmd},
				args => $+{args},
			};
			$parproc{$+{ppid}} .= "$+{pid} ";
		}
	}
	close $ps;
}

# 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) = @_;
	my @arr;
	my $s;

	return () unless defined $parproc{$pid};

	$s = $parproc{$pid};
	while ($s =~ /^(?<pid>\d+) \s+ (?<rest>.*)/x) {
		push @arr, $+{pid};
		$s = $+{rest};
	}

	return @arr;
}

# 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($);

sub proc_get_children_r($) {
	my ($pid) = @_;

	my @chi = proc_get_children $pid;
	my @res;
	for my $child (@chi) {
		push @res, $child, proc_get_children_r $child;
	}
	return @res;
}

# 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) = @_;

	my @arr = proc_get_children_r $pid;
	say "$pid $proc{$pid}{cmd} @arr";
}

# 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() {
	proc_display $_ for keys %parproc;
}

# 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) = @_;
	die "bad pid ($pid)\n" unless $pid =~ /^\d+$/;
	die "non-existent pid ($pid)\n" unless defined $proc{$pid};

	my @arr = ($pid, proc_get_children_r $pid);
	print "@arr ";
	if (scalar @arr != kill $sig, @arr) {
		die "Could not kill all the requested processes (@arr): $!\n"
	}
}

# 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) = @_;
	die "bad pid ($pid)\n" unless $pid =~ /^\d+$/;
	die "non-existent pid ($pid)\n" unless defined $proc{$pid};

	my @arr = ($pid, proc_get_children_r $pid);
	print "@arr ";

	my $exact = $delta !~ /^[+-]/;
	foreach my $kpid (@arr) {
		my $val = $exact? $delta: $delta + getpriority 0, $kpid;
		setpriority 0, $kpid, $val or
		    die "Could not renice $kpid to $val: $!\n";
	}
}

# Main block

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

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

	# Let the user override the ps program location and flags
	$PS = $ENV{'PS'} // $defaultPS;
	$PSflags = $ENV{'PSflags'} // $defaultPSflags;
	proc_gather;

	# no arguments, no kill requested - display all and exit
	if (!@ARGV && !($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;
			} elsif ($niceem) {
				proc_nice $killpid, $sig;
			} else {
				proc_display $killpid;
			}
		} else {
			# process name..
			foreach my $pid (keys %proc) {
				if ($proc{$pid}{cmd} eq $killpid) {
					if ($killem) {
						proc_kill $pid, $sig;
					} elsif ($niceem) {
						proc_nice $pid, $sig;
					} else {
						proc_display $pid;
					}
				}
			}
		}
	}
	print "\n" if $killem || $niceem;
}