#!/usr/local/bin/perl
################################################################################
#
# File:         $RCSfile: $
# Revision:	$Revision: $
# Description:  Remotely run processes on other machines
# Author:       Andrew@DeFaria.com
# Created:      Tue Jan  8 15:57:27 MST 2008
# Modified:	$Date: $
# Language:     perl
#
# (c) Copyright 2008, ClearSCM, Inc., all rights reserved
#
################################################################################
use strict;
use warnings;

use FindBin;
use Getopt::Long;
use Term::ANSIColor qw(:constants);
use POSIX ":sys_wait_h";

my $libs;

BEGIN {
  $libs = $ENV{SITE_PERLLIB} ? $ENV{SITE_PERLLIB} : "$FindBin::Bin/../lib";

  die "Unable to find libraries\n" if !$libs and !-d $libs;
}

use lib "$FindBin::Bin/../lib";
use lib $libs;

use Display;
use Logger;
use Machines;
use Rexec;
use Utils;

our $_host;
our $_skip			= 0;
our $_currentHost;

my $_log			= 0;
my $_quiet			= 0;
my $_alternateFile;
my $_parallel			= 0;

my $_totalMachines		= 0;
my $_totalExecutions		= 0;
my $_totalFailures		= 0;
my $_totalConnectFailures	= 0;
my $_totalSkips			= 0;

my (%_workerStatuses, %_workerNames);

sub Usage {
  my $msg = shift;

  display "ERROR: $msg\n" if defined $msg;

  display "rexec\t[-v] [-d] [-u] <cmd>";
  display "\t-v\tTurn on verbose mode";
  display "\t-d\tTurn on debug mode";
  display "\t-u\tThis usage message";
  display "<cmd>\tCommand to execute remotely";

  exit 1;
} # Usage

sub printStats {
  display YELLOW  . "Machines: "		. RESET . "$_totalMachines " .
          MAGENTA . "Executions/Failures: "	. RESET . "($_totalExecutions/$_totalFailures) " .
          BLUE    . "Connect Failures/Skips: "	. RESET . "($_totalConnectFailures/$_totalSkips)";
} # printStats

sub Interrupted {
  use Term::ReadKey;

  display BLUE . "\nInterrupted execution on $_host" . RESET;

  printStats;

  display_nolf "Executing on " . YELLOW . $_host . RESET . " - "
    . GREEN	. BOLD . "S" . RESET . GREEN	. "kip"		. RESET . ", "
    . CYAN	. BOLD . "C" . RESET . CYAN	. "ontinue"	. RESET . " or "
    . MAGENTA	. BOLD . "A" . RESET . MAGENTA	. "bort run"	. RESET . " ("
    . GREEN	. BOLD . "s" . RESET . "/"
    . CYAN	. BOLD . "C" . RESET . "/"
    . MAGENTA	. BOLD . "a" . RESET . ")?";

  ReadMode ("cbreak");
  my $answer = ReadKey (0);
  ReadMode ("normal");

  if ($answer eq "\n") {
    display "c";
  } else {
    display $answer;
  } # if

  $answer = lc $answer;

  if ($answer eq "s") {
    *STDOUT->flush;
    display "Skipping $_host";
    $_skip = 1;
    $_totalSkips++;
  } elsif ($answer eq "a") {
    display RED . "Aborting run". RESET;
    printStats;
    exit;
  } else {
    display "Continuing...";
    $_skip = 0;
  } # if
} # Interrupted

sub workerDeath {
  while ((my $worker = waitpid (-1, WNOHANG)) > 0) {
    my $status	= $?;

    # Ignore all child deaths except for processes we started
    next if !exists $_workerStatuses{$worker};

    $_workerStatuses{$worker} = $status;
  } # while

  $SIG{CHLD} = \&workerDeath;
} # workerDeath

sub execute ($$$) {
  my ($cmd, $host, $prompt) = @_;

  my @lines;

  verbose_nolf "Connecting to machine $host...";

  eval {
    $_currentHost = new Rexec (
      host	=> $host,
      prompt	=> $prompt,
    );
  };

  # Problem with creating Rexec object. Log error if logging and return.
  if ($@ or !$_currentHost) {
    if ($_log) {
      my $log = new Logger (name => $_host);

      $log->err ("Unable to connect to $host to execute command\n$cmd");
    } # if

    $_totalConnectFailures++;

    return (1, ());
  } # if

  verbose " connected";

  display YELLOW . "$host:" . RESET . UNDERLINE . "$cmd" . RESET unless $_quiet;

  @lines = $_currentHost->exec ($cmd);

  if ($_skip) {
    # Kick current connection
    kill INT => $_currentHost->{handle}->pid;
  } # if

  if ($_parallel != 0) {
    if ($_log) {
      my $log = new Logger (name => $_host);

      $log->err ("Unable to connect to $host to execute command\n$cmd");
    } # if

    $_totalConnectFailures++;
  } # if

  verbose "Disconnected from $host";

  my $status = $_currentHost->status;

  undef $_currentHost;

  return ($status, @lines);
} # execute

sub parallelize ($%) {
  my ($cmd, %machines) = @_;

  my $thread_count = 1;

  foreach $_host (sort keys %machines) {
    if ($thread_count <= $_parallel) {
      debug "Processing $_host ($thread_count)";
      $thread_count++;

      if (my $pid = fork) {
	# In parent process - record this host and its status
	$_workerNames{$pid} = $_host;
      } else {
	# In spawned child...
	$pid = $$;

	debug "Starting process for $_host [$pid]";

	$_workerNames{$pid} = $_host;
	
	my ($status, @lines) = execute $cmd, $_host, $machines{$_host};

	my $log = new Logger (name => $_host);

	$log->log ($_) foreach (@lines);

	exit $status;
      } # if
    } else {
      # Wait for somebody to finish;
      debug "Waiting for somebody to exit...";
      my $reaped = wait;

      debug "Reaped $_workerNames{$reaped} [$reaped] (Status: $?)";
      $_workerStatuses{$reaped} = $? >> 8 if $reaped != -1;

      $thread_count--;
    } # if
  } # foreach

  # Wait for all kids
  my %threads = %_workerNames;

  foreach (keys %threads) {
    if (waitpid ($_, 0) == -1) {
      delete $threads{$_};
    } else {
      $_workerStatuses{$_} = $? >> 8;
      debug "$threads{$_} [$_] exited with a status of $_workerStatuses{$_}";
    } # if
  } # foreach

  debug "All processed completed - Status:";

  if (get_debug) {
    foreach (sort keys %_workerStatuses) {
      debug "$_workerNames{$_}\t[$_]:\tStatus: $_workerStatuses{$_}";
    } # foreach
  } # if

  # Gather output...
  display "Output of all executions";
  foreach $_host (sort keys %machines) {
    if (-f "$_host.log") {
      display "$_host:$_" foreach (ReadFile ("$_host.log"));

      #unlink "$_host.log";
    } else {
      warning "Unable to find output for $_host ($_host.log missing)";
    } # if
  } # foreach
} # parallelize

# Print the totals if interrupted
$SIG{INT} = \&Interrupted;

# Get our options
GetOptions (
  "usage"	=> sub { Usage "" },
  "verbose"	=> sub { set_verbose },
  "debug"	=> sub { set_debug },
  "log"		=> \$_log,
  "quiet"	=> \$_quiet,
  "file=s"	=> \$_alternateFile,
  "parallel:i"	=> \$_parallel,
) || Usage "Unknown parameter";

my $cmd = join " ", @ARGV;

error "No command specified", 1 if !$cmd;

my $machines = Machines->new (file => $_alternateFile);
my %machines = $machines->all ();

if ($_parallel > 0) {
  parallelize ($cmd, %machines);
  printStats;
  exit;
} # if

display "NOTE: Logging output to <host>.log" if $_log;

foreach $_host (sort keys (%machines)) {
  $_totalMachines++;

  my ($status, @lines) = execute $cmd, $_host, $machines{$_host};

  if ($_skip) {
    $_skip = 0;
    next;
  } # if

  if (defined $status) {
    if ($status == 0) {
      $_totalExecutions++;
    } else {
      if ($_log) {
	my $log = new Logger (name => $_host);

	$log->log ("Host: $_host\nCommand: $cmd\nStatus: $status\nOutput:\n");
	$log->log ($_) foreach (@lines);
      } # if
	
      $_totalFailures++;

      next;
    } # if
  } # if

  if ($_log) {
    my $log = new Logger (name => $_host);

    $log->log ("Host: $_host\nCommand: $cmd\nStatus: $status\nOutput:\n");
    $log->log ($_) foreach (@lines);
  } else {
    display $_ foreach (@lines);
  } # if
} # foreach

printStats;