#!/bin/bash

################################################################################
# Defaults

: ${tmpdir=/tmp}

batch_stdin=false   # read stdin once and copy it as input to each command
batch_stdout=false  # collect the stdout from each job and display coherently
batch_stderr=false  # collect the stderr from each job and display coherently
batch_merge=false   # collect the stdout & stderr from each job in a single stream, and display those streams coherently
stdout_mon=false    # send the local output to stdout (instead of stderr)

colour=false        # use coloured output (default if using a terminal)
target_group=       # one name for the collect of target hosts
keep_going=false    # don't stop if one (non-parallel) task returns non-zero exit status
parallel=false      # run on all hosts in parallel
rm_temps=true       # don't clean up the temp files -- normally only for debugging
verbose=true        # print messages indicating the source of each block of output

exclude_list=       # a list of hosts which we don't want to include, presumably because they're offline
exclude_files=$HOME/.rx-exclude # a file which lists the above

stdin=/tmp/rx.$$.stdin   # name of a temp file to hold stdin
stdout=/tmp/rx.$$.stdout # prefix of name of temp files to hold stdout
stderr=/tmp/rx.$$.stderr # prefix of name of temp files to hold stderr

# Is output to a terminal?
test -t 2 && colour=true

################################################################################
#

function enum {
    local pref=$1 lo=$2 hi=$3 sufx=$4 i
    for ((i=lo;i<=hi;i++)) do echo $pref$i$sufx ; done
    let "lo<=hi"  # return "fail" if the expansion was empty
}

################################################################################
# Options

while let $#
do  case $1 in
    (-he|-hel|-help|--h|--he|--hel|--help) cat <<- EOM ; exit 0 ;;
rx - run a command on multiple hosts

rx [-p|-b|-bo|-be|-bm] [-bi] [-t target-group | -t targetNN-NN | -T target,target,...] [-x host] [-X file] command...
    -t      Specify target group (doesn't have to be last option)

    -p      Run on all hosts in parallel (default is serial)
    -b      Run in parallel but serialise stdout & stderr
    -bo     Run in parallel but serialise stdout
    -be     Run in parallel but serialise stderr
    -bi     Batch stdin (repeat it for each host)

    -c      Print colour escapes (default if output is a terminal)
    -C      Don't print colour escapes
    -o      Send monitoring messages to stdout instead of strerr
    -q      Don't print progress info (just output from command)
    -v      Print progress info

    -k      Don't stop if a host returns non-zero exit status
            (only useful in serial mode)
    -K      Don't delete temporary files (useful for debugging)
    -x      Skip a host (useful if the host is down)
    -X      Read list of hosts to skip from a file

Target groups are currently: alias gen grunt ironport mail recurdns

You may find it useful to have shell aliases like:
    alias rs="rx -t smtp "
EOM
    (-b|--batch)         parallel=true batch_stdout=true batch_stderr=true         ;;
    (-bi|--batch-stdin)  batch_stdin=true ;;
    (-bm|--batch-merge)  parallel=true batch_stdout=true batch_stderr=true batch_merge=true ;;
    (-bo|--batch-stdout) parallel=true batch_stdout=true                 ;;
    (-be|--batch-stderr) parallel=true         batch_stderr=true         ;;
    (-o|--monitor-to-stdout) stdout_mon=true ;;
    (-C|--bw|--no-colour|--no-color) colour=false ;;
    (-c|--colour|--color) colour=true ;;
    (-k|--keep-going) keep_going=true ;;
    (-p|--parallel) parallel=true ;;
    (-q|--quiet)    verbose=false ;;
    (-K|--keep|--keep-temps) rm_temps=false ;;
    (-v|--verbose)  verbose=true ;;
    (-t|--targetgroup) target_group=$2 ; shift ;;
    (-t*)           target_group=${1#-?} ;;
    (--targetgroup=*) target_group=${1#--*=} ;;
    (-T|--targets)  H=($(IFS=", ";echo $2)) ; shift ;;
    (-T*)           H=($(IFS=", ";echo ${1#-?})) ;;
    (-targets=*)    H=($(IFS=", ";echo ${1#--*=})) ;;
    (-x|--exclude)  exclude_list=$exclude_list:$2 ; shift ;;
    (-x*)           exclude_list=$exclude_list:${1#-?} ;;
    (--exclude=*)   exclude_list=$exclude_list:${1#--*=} ;;
    (-X|--exclude-from) exclude_files=$exclude_files:$2 ; shift ;;
    (-X*)           exclude_files=$exclude_files:${1#-?} ;;
    (--exclude-from=*) exclude_files=$exclude_files:${1#--*=} ;;
    (--)            shift ; break ;;
    (-*)            echo >&2 invalid option $1 ; exit 64 ;;
    (*) break ;;
    esac
    shift
done

################################################################################
# Read exclusion files, especially ~/.rx-exclude

for xf in $(IFS=:;echo $exclude_files)
do  test -f $xf || continue
    exclude_list=$(tr < $xf -s '
	 ' :)$exclude_list
done

################################################################################
# Clear temp files when we're done with them

$rm_temps && trap 'rm -f $stdin $stdout.* $stderr.*' EXIT HUP INT

################################################################################
# Verbose mode?

if $stdout_mon
then
    exec 4>&1
else
    exec 4>&2
fi

if ! $verbose
then
    # No debug? Do nothing then...
    function mon() { : ; }
    function mon_n() { : ; }
    function mon_b() { : ; }
else
    # Debug? Print messages...
    function mon()  { echo >&4 "$c1$@$c0" ; }
    function mon_n()  { echo >&4 -n "$c1$@$c0" ; }
    function mon_b() { echo >&4 "$c2$@$c0" ; }
fi

################################################################################
# Coloured output on a terminal?

c0= c1= c2= c3= c4= c5=

# Debug output to terminal? make it pretty colours
$colour && case $TERM in
(vt???*|?term*|screen*|linux*|ansi*)
    E="$(echo -e '\033[')"
    c0="${E}39;49;22m" # reset/default
    c1="${E}34;49m"    # dim; messages identifying which session is which
    c2="${E}33;49;1m"  # top-level messages
    c3="${E}31;49;1m"  # error messages
    c4="${E}33;49m"    # success messages
    c5=
    ;;
esac

#
# If we don't already have a target list, make one based on the "group"
#
if ! test "X${#H[*]}" = X0
then
    ################################################################################
    # If we don't otherwise have a target group, Set target group by matching $0

    test "X$target_group" = X &&
    case /$0 in
    ( */rx_* ) target_group=${0##*rx_} ;;
    (*) echo >&2 Target required ; exit 64 ;;
    esac

    ################################################################################
    #

    i=$HOME/.ssh/${target_group:-${H[0]%%[0-9]*}}_id
    u=root

    case $target_group in

    # a list of preset groups; for each group enumerate the valid hostnames
    (singer)        H=($(enum singer 1 16)) ;;
    (writer)        H=($(enum 192.168.132. 41 44)) ;;
    (poet)          H=(poet{2,3,5,7,11,13,17,19}) ;;
    (orator)        H=(orator-{a,b,c}) ;;
    (imposter)      H=(imposter{1,2,3,4,5,6} forger5{1,2,3} swindler7) ; u=admin ;;
    (smtp)          H=($(enum smtp 0 31)) ;;

    (*[^0-9]-* | *-*[^0-9]*) false ;;
    (*[0-9]-[0-9]*)
            h_end=${target_group##*-} ;
            h1=${target_group%-$h_end} ;
            h_base=${h1%%[0-9]*} ;
            h_start=${h1#$h_base} ;
            H=($( enum $h_base $h_start $h_end )) ;;
    (*)     false ;;
    esac ||
    { echo >&2 Undefined or invalid target $target_group ; exit 66 ; }
fi

for ((j=0;j<${#H[@]};j++))
do  hj="${H[j]}"
    case :$exclude_list: in
    (*:$hj:*) mon "${c4}Excluding $hj" ; H[j]='' ;;
    esac
done

H=(${H[*]})  # suppress blanks

################################################################################

let $# || $batch_stdin || { echo >&2 "Command required" ; exit 64 ; }

################################################################################

if $batch_stdin
then
    test -t 0 && mon_b "Spooling input; type ctrl-D to end"
    cat > $stdin
fi

mon_b "Running on: ${H[*]}"

    exec 5<&0 6>&1 7>&2

if $parallel
then

    for ((j=0;j<${#H[@]};j++))
    do  hj=${H[j]}
        mon_n Starting $hj

        $batch_stdin && exec 0< $stdin
        $batch_stdout && exec 1> $stdout.$j
        $batch_stderr && exec 2> $stderr.$j
        ssh -i $i -l $u -n $hj "$@" & pid=$!
        exec 1>&6 2>&7 0<&5

        mon " [$pid]"
        P[j]=$pid
    done

    $batch_merge && exec 7>&6

    trap : INT

    for ((j=0;j<${#H[@]};j++))
    do  hj=${H[j]}
        pid=${P[j]}
        wait $pid ; z=$?
        if $batch_stderr && test -s $stderr.$j ; then mon "Stderr from $hj [$pid]:" ; cat $stderr.$j 1>&7 ; fi
        if $batch_stdout && test -s $stdout.$j ; then mon "Stdout from $hj [$pid]:" ; cat $stdout.$j 1>&6 ; fi
        if let z
        then mon Completed $hj [$pid], ${c3}failed status $z
        else mon Completed $hj [$pid], ${c4}succeeded
        fi
    done
    rm -f $stdout.* $stderr.*

else

    for hj in ${H[@]}
    do  mon Starting $hj
        $batch_stdin && exec 0< $stdin
        ssh -i $i -l $u $hj "$@"
        z=$?
        exec 0<&5
        mon Completed $hj, status $z
        if let z
        then mon Completed $hj, ${c3}failed status $z
             $keep_going || exit $z
        else mon Completed $hj, ${c4}succeeded
        fi
    done

fi

mon_b "Completed on: ${H[*]}"