#!/bin/bash -

# Name: massh
# Modified: 20100822
# Description:
# Mass SSH'er that can also push and run scripts.  
# Copyright: 2010 Michael Marschall

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

###############################################################################
# Be Sure to Eat All Your Variables                                           #
###############################################################################
# We Debuggin?
echo "$@" | grep '\-D' &>/dev/null && set -x

# Wait, what airline is this?
OS=$(uname)

# Get basename and use it like it's hot. 
IDOWOT=`basename $0`

# Time Stamp for first line of hosts.results.
TIMESTAMP=`date +'%Y-%m-%d %H:%M:%S'`

# So Random!
RAND="$RANDOM"

# Trap Wisely
trap "kill 0" EXIT

# System Wide Config File and Massh Dir 
[ -f /etc/$IDOWOT ] && source /etc/$IDOWOT
[ -z $ALLFILES ] && ALLFILES="/var/massh"

# User Specific Config File and Massh Dir 
[ -f $HOME/.$IDOWOT ] && source $HOME/.$IDOWOT
[ -z $MYFILES ] && MYFILES="$HOME/massh"

# Set PATH
[ -d "$HOME/bin" ] \
&& PATH=/usr/bin:/bin:/usr/sbin:/sbin:/usr/bin:/usr/local/sbin:$HOME/bin

###############################################################################
# Usage                                                                       #
###############################################################################
function usage {  
    echo "Usage: "
    echo "      $IDOWOT [-f file | -r range] [-c cmd | -s script | -p file]"
    echo "            [-o -D [-S -R] [-F -R]]"
    echo "      -f        File containing hostnames or ranges "
    echo "                the following directories:"
    echo "                    * $ALLFILES"
    echo "                    * $MYFILES"
    echo "                    * \$CWD, full path, relative path"
    echo "      -r        Arbitrary or file defined host groupings"
    echo "      -c        Command to run. Quote commands to insure execution."
    echo "      -s        Script to push to all hosts and run."
    echo "      -p        File to push to all hosts."
    echo "      -o        Full, hostname tagged output"
    echo "      -P        Number of parallel session. Default is $PARA"
    echo "      -F        Output only hostnames where remote command failed."
    echo "      -S        Output only hostnames where remote command succeeded."
    echo "      -R        Regurgitate F,S output and feed back to $IDOWOT -r"
    echo "      -D        Debug"
    echo "      -M        Massh Menu"
    echo "      -O        SSH options in additinon to the following:"
    echo "                 ${SSHOPTS[0]} ${SSHOPTS[1]}"
    echo "                 ${SSHOPTS[2]} ${SSHOPTS[3]}"
    echo "Output: "
    echo "Succeeded  - The command, push, push/run ran with no errors."
    echo "   Failed  - The connection, command, push, push/run failed."
    echo "  Skipped  - Host was listed in one of the following files:"
    echo "                    * $ALLFILES/hosts.down"
    echo "                    * $MYFILES/hosts.down"
    echo "Examples: "
    echo "    $IDOWOT -f /tmp/hosts.txt -c 'cat /etc/passwd | grep ^root:' -o"
    echo "    $IDOWOT -r dbservers -p ~/.hushlogin"
    echo "    $IDOWOT -r web.[1,5,[10..15]].google.com -c 'last | head -10' -o"
    echo "    $IDOWOT -r webservers:dbservers:appservers -c 'df -h' -o"
    echo
    exit 69
}

###############################################################################
# Command Line Args                                                           #
###############################################################################
while getopts ":MDhf:P:c:t:op:O:s:r:FSR" Option
do
    case $Option in
        f) FILE=$OPTARG;; 
        r) RANG=$OPTARG;; 
        c) COMM=$OPTARG;;
        s) SCRI=$OPTARG;;
        p) PUSH=$OPTARG;;
        P) PARA=$OPTARG;;
        t) TIME=$OPTARG;;
        O) OPTS=$OPTARG;;
        D) DEBU=$OPTARG;;
        M) MENU=yes    ;;
        o) OUTP=yes    ;;
        F) FAIL=yes    ;;
        S) SUCC=yes    ;;
        R) REGU=yes    ;;
        h) usage       ;;
    esac
done

###############################################################################
# Options                                                                     #
###############################################################################
# Change these or supply something different on the command line.

# Parallel SSH's to run, -P on the command line.
[ -z $PARA ] && PARA=30

# SSH Timeout in seconds, -t on the command line.
[ -z $TIME ] && TIME=5

# SSH Defaults
CONTIME="-o ConnectTimeout=$TIME"
[ "$OS" = "SunOS" ] && CONTIME=""

# Array for SSH options.
SSHOPTS[0]="-o StrictHostKeyChecking=no"  
SSHOPTS[1]="-o LogLevel=QUIET"
SSHOPTS[2]="-o BatchMode=yes"
SSHOPTS[3]="$CONTIME $OPTS"

# Option Combo Checking
if [ -n "$FILE" ]
then
    [ -z "$COMM" ] && [ -z "$PUSH" ] && [ -z "$SCRI" ] && \
    echo "ERROR: You must use the f option with -c -p or -s" && usage
fi

if [ -n "$RANG" ]
then
    [ -z "$COMM" ] && [ -z "$PUSH" ] && [ -z "$SCRI" ] && \
    echo "ERROR: You must use the r option with -c -p or -s" && usage
fi

# [Try] to log all massh runs by default. This can be overridden via ~/.massh
! [ "$WELOGIN" = "no" ] && \
    $(type -p logger) -i -t massh -- "USER=$USER COMMAND=$IDOWOT $@"

# [Re]Create hosts.results for S,R options.
[ "$SUCC" = "yes" ] || [ "$FAIL" = "yes" ] && rm -f $MYFILES/hosts.results
[ "$SUCC" = "yes" ] || [ "$FAIL" = "yes" ] && \
    echo "# $TIMESTAMP $IDOWOT $@" > $MYFILES/hosts.results

# Temp File Location
[ -d $HOME/tmp ] && TEMPDIR=$HOME/tmp || TEMPDIR=/tmp

# No Options 
[ -z "$1" ] && usage

###############################################################################
# Massh Menu                                                                  #
###############################################################################
function mmenu_hosts {
    echo "Enter a filename, full file path, abitrary range or pre-defined"
    echo "range containing the hosts that you want to Massh to."
}

function mmenu {
    clear
    echo "Massh Menu" 
    echo "1) Massh a Command"
    echo "2) Massh a File"
    echo "3) Massh a Script"
    echo "4) Help Text"
    echo -n "Choice? : "
    read firstchoice
    echo
    [ "$firstchoice" = "4" ] && usage
    if [ "$firstchoice" = "1" ]
    then
        mmenu_hosts
        echo
        echo -n "HOSTS? : "
        read RANG
        echo
        echo "Enter the command that will be Massh'ed to your remote hosts."
        echo
        echo -n "COMMAND? : "
        read COMM
        echo 
        echo "Would you like Full Output or Simple Verification"
        echo "1) Full Output"
        echo "2) Simple Verification"
        echo
        echo -n "OUTPUT TYPE? : "
        read MENUOUT
        [ "$MENUOUT" = "1" ] && OUTP="yes"
        echo "Running..."
    elif [ "$firstchoice" = "2" ]
    then
        mmenu_hosts
        echo
        echo -n "HOSTS? : "
        read RANG
        echo
        echo "Enter the file that will be Massh'ed to your remote hosts."
        echo
        echo -n "FILE? : "
        read PUSH
        echo "Running..."
    elif [ "$firstchoice" = "3" ]
    then
        mmenu_hosts
        echo
        echo -n "HOSTS? : "
        read RANG
        echo
        echo "Enter the script that will be Massh'ed to your remote hosts."
        echo
        echo -n "SCRIPT? : "
        read SCRI
        echo 
        echo "Would you like Full Output or Simple Verification"
        echo "1) Full Output"
        echo "2) Simple Verification"
        echo
        echo -n "OUTPUT TYPE? : "
        read MENUOUT
        [ "$MENUOUT" = "1" ] && OUTP="yes"
        echo "Running..."
    fi
}

[ -n "$MENU" ] && mmenu

###############################################################################
# Get Hosts from File or Range                                                #
###############################################################################
[ -n "$FILE" ] && HOSTLIST=(`ambit $FILE`) 
[ -n "$RANG" ] && HOSTLIST=(`ambit $RANG`)

# Set Loop Var, Get Host Array Element Count
index=0
index_count=${#HOSTLIST[@]}

###############################################################################
# Pushing Files and Dirs                                                      #
###############################################################################
for sending in $PUSH $SCRI
do
    [ -n "$PUSH" ] && prefix="file"
    [ -n "$SCRI" ] && prefix="script"

    [ -f "$sending" ]                           \
        && SENDING="$sending"                   \
        && SENT=`basename $SENDING`

    [ -f $MYFILES/$prefix.$sending ]            \
        && SENDING="$MYFILES/$prefix.$sending"  \
        && SENT=`basename $SENDING`

    [ -f $ALLFILES/$prefix.$sending ]           \
        && SENDING="$ALLFILES/$prefix.$sending" \
        && SENT=`basename $SENDING`
done

###############################################################################
# Functions                                                                   #
###############################################################################
# Mass Command - Full Output
function showoutput {
    ssh ${SSHOPTS[*]} $HOST "$COMM"
    ! [ $? -eq 0 ] && echo "Failed"
}

# Mass Command - Simple Output - Succeeded, Failed, Regurgitate
function nooutput {
    ssh ${SSHOPTS[*]} $HOST "$COMM" &>/dev/null
    if [ $? -eq 0 ]
    then
        if [ -z "$FAIL" ]
        then
            if [ -n "$SUCC" ]
            then
                [ -n "$REGU" ] && echo -n "$HOST:" \
                    || echo "$HOST" | tee -a $MYFILES/hosts.results
            else
                echo "$HOST : Suceeded"
            fi
        fi
    else
        if [ -z "$SUCC" ]
        then
            if [ -n "$FAIL" ]
            then
                [ -n "$REGU" ] && echo -n "$HOST:" \
                    || echo "$HOST" | tee -a $MYFILES/hosts.results
            else
                echo "$HOST : Failed"
            fi
        fi
    fi
}

# Push File - Simple Output - Succeeded/Failed
function push {
    scp ${SSHOPTS[*]} $SENDING $HOST: &>/dev/null 
    [ $? -eq 0 ] && echo "$HOST : Succeeded" || echo "$HOST : Failed"
}

# Push Script - Simple Output - Push, Run, Cleanup
function pushandrun {
    scp ${SSHOPTS[*]} $SENDING $HOST: &> /dev/null \
        && ssh ${SSHOPTS[*]} $HOST "./$SENT; rm -rf ./$SENT" \
        &> /dev/null
    [ $? -eq 0 ] && echo "$HOST : Succeeded" || echo "$HOST : Failed"
}

# Push Script - Full Output - Push, Run, Display, Cleanup
function pushandrunout {
    scp ${SSHOPTS[*]} $SENDING $HOST: &> /dev/null \
        && ssh ${SSHOPTS[*]} $HOST "./$SENT; rm -rf ./$SENT"
    ! [ $? -eq 0 ] && echo "Failed"
}

###############################################################################
# Where Sausage is Made                                                       # 
###############################################################################
airbourne="0"
while [[ $index -le $index_count && "$airbourne" -lt 1 ]]
do
    # The guest of honor.
    HOST=${HOSTLIST[$index]}

    # Check to see if $HOST should be intentionally skipped. 
    if [ -f $ALLFILES/hosts.down ] && [ -n "$HOST" ]
    then
        DWNHOST="$ALLFILES/hosts.down"
        if grep $HOST $DWNHOST &> /dev/null
        then
            ! [ -n "$SUCC" ] || [ -n "$FAIL" ] \
                && echo "$HOST : Skipped"
            [ $index -lt $index_count ] && let "index++" && continue \
                || continue
        fi
    fi
    if [ -f $MYFILES/hosts.down ] && [ -n "$HOST" ]
    then
        DWNHOST="$MYFILES/hosts.down"
        if grep $HOST $DWNHOST &> /dev/null
        then
            ! [ -n "$SUCC" ] || [ -n "$FAIL" ] && echo "$HOST : Skipped"
            [ $index -lt $index_count ] && let "index++" && continue \
                || continue
        fi
    fi

    BGS=$(jobs -p | wc -l)

    if [ "$BGS" -lt "$PARA" ]
    then
            # Simple Output
        if  [ -z "$OUTP" ] && [ -z "$SCRI" ] && \
            [ -z "$PUSH" ] && [ -n "$HOST" ]
        then
            nooutput &
            let "index++"
            continue
        fi

        # Full Output
        if  [ -n "$OUTP" ] && [ -z "$SCRI" ] && \
            [ -z "$PUSH" ] && [ -n "$HOST" ]
        then
            showoutput &> $TEMPDIR/$HOST.$USER.$IDOWOT.$RAND &
            inflight[$!]=$HOST
        fi

            # Push
        if  [ -z "$OUTP" ] && [ -z "$SCRI" ] && \
            [ -n "$PUSH" ] && [ -n "$HOST" ]
        then
            push &
            let "index++"
            continue
        fi

        # Push and Run Simple Output
        if  [ -z "$OUTP" ] && [ -n "$SCRI" ] && \
            [ -z "$PUSH" ] && [ -n "$HOST" ]
        then
            pushandrun &
            let "index++"
            continue
        fi

        # Push and Run Full Output
        if  [ -n "$OUTP" ] && [ -n "$SCRI" ] && \
            [ -z "$PUSH" ] && [ -n "$HOST" ]
        then
            pushandrunout &> $TEMPDIR/$HOST.$USER.$IDOWOT.$RAND &
            inflight[$!]=$HOST
        fi

        # Inflight Entertainment
        for flying in ${!inflight[*]}
        do
            #[ "$TRAP" = "yes" ] && break

            if ! kill -0 $flying &> /dev/null
            then
                while read line
                do
                    echo "${inflight[$flying]} : $line"
                done <$TEMPDIR/${inflight[$flying]}.$USER.$IDOWOT.$RAND
                rm -f $TEMPDIR/${inflight[$flying]}.$USER.$IDOWOT.$RAND
                unset inflight[$flying]
            fi
        done
        [ ${#inflight[*]} -eq 0 ]   && let "airbourne++"
        [ $index -lt $index_count ] && let "index++"
    fi
done
rm -f $TEMPDIR/*$USER.$IDOWOT.$RAND &> /dev/null
wait
