#!/bin/bash
#   parssh: Parallel SSH orchestration in a Bash session.
#   Copyright (C) 2015  James Pannacciulli <jpnc@jpnc.info>

#   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/>.

# help function
_parssh.usage () {
    printf '  %s\n'\
        "parssh: Parallel SSH orchestration in a Bash session."\
        ""\
        "SYNOPSIS"\
        "parssh [-NUM] [-r|--rinput FILE] [-s|--servers FILE] [-b|--bare] [-c|--config]"\
        "  [COMMANDS] [< SERVERS]"\
        ""\
        "DESCRIPTION"\
        "  -NUM  (default: 4)"\
        "     Number of concurrent ssh connections to maintain.  E.g. '-40'."\
        ""\
        "  -r FILE, --rinput FILE  (default: unset/inactive)"\
        "     File to send as STDIN redirection on remote servers."\
        "     (Can be used as replacement for or in conjunction with COMMANDS)."\
        ""\
        "  -s FILE, --servers FILE  (default: unset/inactive)"\
        "     File to use as list of servers to run COMMANDS on."\
        "     (Cannot be used in conjunction with a server list on STDIN)."\
        ""\
        "  -b, --bare  (default: unset/inactive)"\
        "     Disable prepending of hostname to each output line returned by COMMANDS on SERVERS."\
        ""\
	"  -C, --savecmd  (default: unset/inactive)"\
	"     Save the remote command line and local meta data as a header on the output."\
	""\
        "  -c SSH_OPTIONS, --config SSH_OPTIONS  (default: unset/inactive)"\
        "     Comma separated list of configuration parameters to be passed to SSH via '-o' flag."\
        ""\
        "  -o PREFIX, --out PREFIX  (default: unset/inactive)"\
        "     Redirect STDOUT / STDERR to PREFIX.out / PREFIX.err, respectively."\
        ""\
        "  -t PREFIX, --tee PREFIX  (default: unset/inactive)"\
        "     Copy STDOUT / STDERR to PREFIX.out / PREFIX.err, respectively."\
        ""\
        "  COMMANDS"\
        "     The list of commands to be executed remotely by SSH on each SERVER."\
        ""\
        "  < SERVERS"\
        "     Unless '-s' flag is used, STDIN will be used as list of remote servers."\
        "     (Deliniated by newlines)."\
        ""\
        "NOTES"\
        "  Given the nature of entering passwords manually, this function will likely not be"\
        "     very useful without properly authorized SSH keys on all SERVERS."
    return 1
}

# text processing functions
_parssh.host_prepend ()
    while read -r; do
        printf "${_parssh_prepend_host+$host: }%s\n" "$REPLY"
    done

_parssh.out () {
    local _parssh_out_stream="$1"

    while read -r; do
        [[ "${_parssh_out_stdout}" == "true" ]] && \
            printf "%s\n" "$REPLY"
        [[ "${_parssh_out_file}" == "true" ]] && \
            printf "%s\n" "$REPLY" >> "${_parssh_out_prefix}.${_parssh_out_stream}"
    done
}

_parssh.savecmd () {
    CMDLINE=$(fc -ln -0)
    CMDLINE="${CMDLINE#[[:space:]][[:space:]]}"
    printf "### # ##\n"
    printf "### %s: %s\n"\
        date "$(date +%F@%R%z)"\
        user "$USER($UID)"\
	pwd "$PWD"\
	cmdline "$CMDLINE"
    printf "### BEGIN parssh remote commands ###\n"
    printf "%s\n### END parssh remote commands ###\n### ## #\n"\
        "$@"
}

# fd handlers
_parssh.serverlist_fd () {
    [[ -z "$_parssh_servers" ]] && {
        [[ -t 0 ]] && {
            echo "No list of servers provided."
            return 79
        }
        exec 9<&0
    } || {
        exec 9< <(printf "%s\n" "$_parssh_servers")
    }
}

_parssh.serverlist_fd_close () {
    exec 9>&-
}

# ssh workers
_parssh.ssh () {
    ssh -n ${_parssh_ssh_config+-o} ${_parssh_ssh_config//,/ -o } $host -- "$@"
}

_parssh.ssh_rinput () {
    ssh -T ${_parssh_ssh_config+-o} ${_parssh_ssh_config//,/ -o } $host -- "$@"\
        < "$_parssh_rinput"
}

# main
parssh () {
    (( $# )) || {
        _parssh.usage
        return $?
    }

    # preserve current bash options and ensure monitor (job control) is set
    local _parssh_origopts=$-
    set -m

    local _parssh_out_stdout=true
    local _parssh_out_file=false
    local _parssh_prepend_host=true
    local _parssh_ssh=_parssh.ssh

    while [[ "$1" == -* ]]; do
        case ${1#-} in
            -)
                shift
                break
            ;;
            [0-9]*)
                [[ "${1#-}" != *[^0-9]* ]] && {
                    local _parssh_concurrency=${1#-}
                }
            ;;
            b|-bare)
                unset _parssh_prepend_host
            ;;
            c|-config)
                local _parssh_ssh_config="$2"
                shift
            ;;
	    C|-savecmd)
	        local _parssh_savecmd="true"
	    ;;
            h|-help)
                _parssh.usage
                return $?
            ;;
            o|-out)
                touch "${2}."{out,err} || {
                    echo "'${2}.{out,err}': unable to write / modify file"
                    return 97
                }
                local _parssh_out_prefix="$2"
                local _parssh_out_stdout=false
                local _parssh_out_file=true
                shift
            ;;
            r|-rinput)
                [[ -r "$2" ]] || {
                    echo "'$2': invalid file name or permissions issue"
                    return 98
                }
                local _parssh_rinput="$2"
                local _parssh_ssh=_parssh.ssh_rinput
                shift
            ;;
            s|-servers)
                [[ -r "$2" ]] || {
                    echo "'$2': invalid file name or permissions issue"
                    return 99
                }
                local _parssh_servers="$2"
                shift
            ;;
            t|-tee)
                touch "${2}."{out,err} || {
                    echo "'${2}.{out,err}': unable to write / modify file"
                    return 96
                }
                local _parssh_out_prefix="$2"
                local _parssh_out_stdout=true
                local _parssh_out_file=true
                shift
            ;;
        esac
        shift
    done

    # load list of servers on fd 9
    _parssh.serverlist_fd

    # run main command loop / capture outputs
    {
        [[ "$_parssh_savecmd" == "true" ]] && _parssh.savecmd
        while read host; do
            while (( $(jobs -pr | wc -l) >= ${_parssh_concurrency:-4} )); do
                sleep 1
            done
            $_parssh_ssh "$@" 2> >(_parssh.host_prepend >&2) |\
                _parssh.host_prepend &
        done <&9 2> >(_parssh.out err >&2)
    } | _parssh.out out

    wait

    # close serverlist fd
    _parssh.serverlist_fd close

    # restore bash options if modified
    [[ "${_parssh_origopts//[^m]/}" == "m" ]] || set +m
}