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

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

    local _parssh_origopts=$-
    set -m

    local _parssh_outfun=_parssh.out
    local _parssh_prepend_host=true
    local _parssh_ssh=_parssh.ssh

    while [[ "$1" == -* ]]; do
        case ${1#-} in
            [0-9]*)
                [[ "${1#-}" != *[^0-9]* ]] && {
                    local _parssh_concurrency=${1#-}
                }
            ;;
            s|-servers)
                [[ -r "$2" ]] || {
                    echo "'$2': invalid file name or permissions issue"
                    return 99
                }
                local _parssh_servers="$2"
                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
            ;;
            b|-bare)
                unset _parssh_prepend_host
            ;;
            o|-out)
                touch "${2}."{out,err} || {
                    echo "'${2}.{out,err}': unable to write / modify file"
                    return 97
                }
                local _parssh_outprefix="$2"
                local _parssh_outfun=_parssh.out_redirect
                shift
            ;;
            t|-tee)
                touch "${2}."{out,err} || {
                    echo "'${2}.{out,err}': unable to write / modify file"
                    return 96
                }
                local _parssh_outprefix="$2"
                local _parssh_outfun=_parssh.out_copy
                shift
            ;;
            h|-help)
                _parssh.usage
                return $?
            ;;
            -)
                shift
                break
            ;;
        esac
        shift
    done

    [[ -z "$_parssh_servers" ]] && {
        [[ -t 0 ]] && {
            echo "No list of servers provided."
            return 79
        }
        exec 9<&0
    } || {
        exec 9<"$_parssh_servers"
    }

    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_outfun err >&2) |\
        $_parssh_outfun out

    wait

    exec 9>&-
    [[ "${_parssh_origopts//[^m]/}" == "m" ]] || set +m
}

_parssh.ssh () {
    ssh -no StrictHostKeyChecking=no $host -- "$@"
}

_parssh.ssh_rinput () {
    ssh -To StrictHostKeyChecking=no $host -- "$@" < "$_parssh_rinput"
}

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

_parssh.out () {
    cat -
}

_parssh.out_redirect () {
    local _parssh_out_stream="$1"
    cat - > "${_parssh_outprefix}.${_parssh_out_stream}"
}

_parssh.out_copy () {
    local _parssh_out_stream="$1"
    tee "${_parssh_outprefix}.${_parssh_out_stream}"
}

_parssh.usage () {
    printf '  %s\n'\
        "parssh: Parallel SSH orchestration in a Bash session."\
        ""\
        "SYNOPSIS"\
        "parssh [-NUM] [-r|--rinput FILE] [-s|--servers FILE] [-b|--bare] [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."\
        ""\
        "  -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)."
    return 1
}