#!/usr/local/bin/perl
#
# masshosts - John Mechalas, Intel Corporation
#

#
# See the license.txt file for licensing and copying information.
#

# The one hardcoded piece.  Set this to the full path/location of your
# config file.

$MASSHOSTS_CONFIG= '/usr/local/lib/masshosts.pl';

#----------------------------------------------------------------------------
# You shouldn't have to change anything below here.
#----------------------------------------------------------------------------

use Getopt::Std;
use IO::Socket;
use POSIX;

#----------------------------------------------------------------------------
# The main program body.  Looks deceptively short, doesn't it?
#----------------------------------------------------------------------------

require $MASSHOSTS_CONFIG;

$GETHOSTS_PL= '/usr/local/lib/gethosts.pl' unless $GETHOSTS_PL;
$CONNECT_TIMEOUT= 10 unless $CONNECT_TIMEOUT;
$EXCLUDE_FILE= '/usr/local/etc/masshosts.exclude' unless $EXCLUDE_FILE;
$FILTER_FILE= '/usr/local/etc/masshosts.filters' unless $FILTER_FILE;
$RSH_CMD= 'rsh' unless $RSH_CMD;

require $GETHOSTS_FUNC;

($prog= $0)=~ s#^.*/##;

&ParseOptions or &Usage;

$shell_port= (getservbyname('shell', 'tcp'))[2] if $opt{c};

if ( &HostList(\@Hosts) ) {
	&Action(\@Hosts);
}

exit 0;

#----------------------------------------------------------------------------
# Either print our host list, execute commands on the local machine or rsh
# to the remote host, depending upon our command-line options.
#----------------------------------------------------------------------------

sub Action {
	my ($aref)= shift;

	if ( $opt{e} ) {
		if ( $opt{o} ) {
			open(SAVEOUT, ">&STDOUT");
			open(SAVEERR, ">&STDERR");
		}

		if ( $opt{p} ) {
			&parallelExecute($aref);
		} else {
			&serialExecute($aref);
		}

		if ( $opt{o} ) {
			close(SAVEOUT);
			close(SAVERR);
		}
	} else {
		print join("\n", @$aref) . "\n";
	}

	1;
}

#----------------------------------------------------------------------------
# Execute our commands in a serial fashion, waiting for one command to 
# complete before starting the next one.
#----------------------------------------------------------------------------

sub serialExecute {
	my ($aref)= shift;
	my ($host, $cmd, $prefix, $status);

	foreach $host (@$aref) {
		($cmd= $opt{e})=~ s/%HOST%/$host/g;

		print STDERR "Executing command for $host\n" if $opt{V};

		&setOutErr($host) if $opt{o};

		if ( $opt{r} ) {
			$status= system $RSH_CMD, $host, $cmd;
		} else {
			$status= system $cmd;
		}
		exit(0) if ( $status &= WIFSIGNALLED );

		&restoreOutErr if $opt{o};
	}
}

#----------------------------------------------------------------------------
# Execute our commands in a parallel fashion, keeping N processes active
# at any one time, spawning new processes as the old ones exit.
#----------------------------------------------------------------------------

sub parallelExecute {
	my ($aref)= shift;
	my ($limit, $current, $host, $newpid);
	my (%runpids);

	$limit= $opt{p};
	$current= 0;

	while(@$aref) {
		$current= scalar keys %runpids;
		if ( $current < $limit ) {
			$host= shift(@$aref);
			if ( $newpid= &spawn($host) ) {
				$runpids{$newpid}= $host;
			}
			next;
		}

		&reapChildren(\%runpids);
	}

	print STDERR "Spawns complete\n" if ( $opt{v} );


	$current= scalar keys %runpids;
	while ( $current > 0 ) {
		print STDERR 'Waiting for: ' . join(' ', values %runpids) .
			"\n" if ( $opt{v} && ( $before != $after ) );

		$before= $current;

		&reapChildren(\%runpids);
		$current= scalar keys %runpids;

		$after= $current;
	}

	1;
}

sub spawn {
	my ($host)= shift;
	my ($cmd, $pid, $status);

	($cmd= $opt{e})=~ s/%HOST%/$host/g;

	if ( $pid= fork ) {					# parent
		print STDERR "Spawned $pid ($host)\n" if $opt{V};
	} elsif ( defined $pid ) {				# child
		&setOutErr($host) if $opt{o};
		if ( $opt{r} ) {
			if ( $opt{c} ) {
				my ($s)= IO::Socket::INET->new(
					PeerAddr	=> $host,
					PeerPort	=> $shell_port,
					Proto		=> 'tcp',
					Type		=> SOCK_STREAM,
					Timeout		=> $CONNECT_TIMEOUT
				);
				if ( defined($s) ) {
					close($s);
				} else {
					die "$host: $@\n";
				}
			} 

			alarm($opt{t}) if ( $opt{t} );

			if ( $opt{i} ) {
				open(RSH, "-|") || exec $RSH_CMD, $host, $cmd;
				while(<RSH>) {
					print join(': ', $host, $_);
				}
				close(RSH);
				exit $?;
			} else {
				exec $RSH_CMD, $host, $cmd;
				die "exec: $!\n";
			}
		} else {
			alarm($opt{t}) if ( $opt{t} );
			exec $cmd;
		}
		die "exec: $!\n";
	} else {						# fork error
		warn "fork: $!\n";
		return 0;
	}

	$pid;
}

sub reapChildren {
	my ($href)= shift;
	my ($pid, $host, $collect);

	while( ($pid, $host)= each %$href ) {
		if ( waitpid($pid, &WNOHANG) > 0 ) {
			delete $href->{$pid};
			&removeEmpty($host) if $opt{z};
			print STDERR "Collected $pid ($host)\n" if $opt{V};
		}
	}
}

#----------------------------------------------------------------------------
# When the -o switch is specified, write out STDOUT and STDERR to a file
# named by the prfix (argument to -o) followed by .host.out or .host.err
# Restore the "real" STDOUT and STDERR when we are done.  If -z is
# specified, delete any zero-length (empty) files that we created.
#----------------------------------------------------------------------------

sub restoreOutErr {
	my ($host)= shift;

	open(STDOUT, ">&SAVEOUT");
	open(STDERR, ">&SAVEERR");

	&removeEmpty($host) if $opt{z};
}

sub removeEmpty {
	my ($host)= shift;

	my ($outfile)= join(".", $opt{o}, $host, 'out');
	my ($errfile)= join(".", $opt{o}, $host, 'err');

	unlink $outfile unless -s $outfile;
	unlink $errfile unless -s $errfile;
}

sub setOutErr {
	my ($host)= shift;
	my ($prefix);

	$prefix= join(".", $opt{o}, $host);
	open(STDOUT, ">$prefix.out");
	open(STDERR, ">$prefix.err");
}

#----------------------------------------------------------------------------
# Get a list of hosts according to the options specified on the command
# line.
#----------------------------------------------------------------------------

sub HostList {
	my ($aref)= shift;
	my (%Exclude, @filters);

	if ( $opt{l} ) {
		@$aref= @ARGV;
	} elsif ( $opt{L} ) {
		chomp(@$aref= <>);
	} elsif ( $opt{f} ) {
		&getHosts($aref, \@ARGV);
	} else {
		if ( $opt{F} ) {
			chomp(@filters= <>);
		} else {
			if ( $opt{K} ) {
				my (@keywords);
				chomp(@keywords= <>);
				&getFilters(\@keywords, \@filters);
			} else {
				&getFilters(\@ARGV, \@filters);
			}
		}
		&getHosts($aref, \@filters);
	}

	@ARGV= ();

	if ( $excludefile ) {
		open(EXC, $excludefile) or warn "open: $excludefile: $!\n";
		while(<EXC>) {
			chomp;
			$Exclude{$_}= 1;
		}
		close(EXC);
	}

	if ( @$aref ) {
		my (@tmparray, $elem, $prevelem);

		@tmparray= sort { $a cmp $b; } @$aref;
		@$aref= ();

		foreach $elem (@tmparray) {
			next if ( $elem eq $prevelem );
			next if ( $opt{x} && $Exclude{$elem} );
			next if ( $opt{n} && ! &network($elem, $opt{n}) );

			push(@$aref, $elem);
			$prevelem= $elem;
		}
	}

	scalar @$aref;
}

#----------------------------------------------------------------------------
# Take the keywords that were given on the command line, and turn them
# into filter expressions from the filters file.  These filters will be
# passed to the &GetHosts subroutine, which will look up our list of hosts 
# and find matches.
#----------------------------------------------------------------------------

sub getFilters {
	my ($arefKeys, $arefFilters)= @_;
	my ($filter, $regex);

	open(FILTERS, $FILTERS_FILE) or die "open: $FILTERS_FILE: $!\n";
	while(<FILTERS>) {
		chomp;
		s/#.*//;
		s/^\s+//;
		next if ( /^\s*$/ );

		($filter, $regex)= split(/\s+/);
		push (@$arefFilters, $filter) if ( grep(/\b($regex)\b/i,
			@$arefKeys) );

	}
	close(FILTERS);
}

#----------------------------------------------------------------------------
# Do the IP addresses for this host match the given network expression?
#----------------------------------------------------------------------------

sub network {
	my ($host, $regex)= @_;
	my (@addrs);

	@addrs= gethostbyname($host) or warn "$host: $!\n";
	foreach $addr ( map { inet_ntoa($_) } @addrs[4 .. $#addrs]) {
		return 1 if ( $addr=~ m/$regex/o );
	}

	0;
}

#----------------------------------------------------------------------------
# Parse our command-line options.
#----------------------------------------------------------------------------

sub ParseOptions {
	getopts('ce:fFhiKlLn:o:p:rt:vVxX:z', \%opt) or return 0;

	return 0 if ( $opt{h} || ( ( + $opt{f} + $opt{l} + $opt{F} +
		$opt{L} ) > 1 ) || ( $opt{x} && $opt{X} ) || !@ARGV );

	return 0 if ( !$opt{e} && ( $opt{c} || $opt{o} || $opt{p} ||
		$opt{r} ) );

	$opt{v}= ( $opt{v} || $opt{V} );

	$excludefile= $EXCLUDE_FILE if $opt{x};
	$excludefile= $opt{X} if $opt{X};

	1;
}

#----------------------------------------------------------------------------
# Print our usage
#----------------------------------------------------------------------------

sub Usage {
	print STDERR 
"usage: $prog [ -f | -F | -K | -l | -L ] [ -ivV ] [ -x | -X file ] [ -n net_expr ] [ [ -crz ] [ -p N ] [ -o prefix ] -e cmd ] arg1 arg2 ...
       $prog -h
";

	exit 1;
}

