sshx

#!/bin/bash
###############################################################################
# sshx: Run a Local Script over SSH
#
# Source: https://github.com/sourcesimian/bin
###############################################################################
set -e

function usage () {
   cat <<EOF
SSHX: Run a Local Script over SSH                     v1.0   by Source Simian

A SSH utility script that serialises a local command including the contents of
any script and argument files and runs it on a remote host. The total size of
the command and file arguments must be less than ARG_MAX since the upload is
done in the SSH command line buffer.

Usage:  sshx [options] -- [command]

Options:
  *                SSH options
  --workdir=<dir>  Change to directory before running command
  --env=<name>     Export local environment variable to the remote host
  --rcfile=<file>  Source this file before running the command
  --gzip           Compress the launcher script content with Gzip
  --debug          Run the launcher script with debug output
  --dump[=<file>]  Write the launcher script to the file. Default /dev/stdout

Command:
  The command and arguments you wish to run on the remote host. This may
  include a local script and/or local argument files.

Source:
  https://github.com/sourcesimian/bin

EOF
}

function require () {
    if ! type -P "$1" > /dev/null; then
        echo "Required binary missing: '$1'" >&2
        return 1
    fi
}

function resolve_path() {
    local RES
    if RES=$(which "$1"); then
        echo "$RES"
        return
    fi
    echo "$1"
}

function realpath () {
    [[ $1 = /* ]] && echo "$1" || echo "$PWD/${1#./}"
}

function stat_mode () {
    if [[ "$OSTYPE" == darwin* ]]; then
        stat -f %Lp "$1"
    else
        stat --format %a "$1"
    fi
}

function stat_mtime () {
    if [[ "$OSTYPE" == darwin* ]]; then
        stat -f '%m' "$1"
    else
        stat -c '%Y' "$1"
    fi
}

function cat_base64_file () {
    local FILE=${1:?[?FILE]}
    echo "base64 --decode <<<\"$(base64 "$FILE" | tr -d "\n")\"" "${@:2}"
}

function cat_gzip_file () {
    local FILE=${1:?[?FILE]}
    echo "gunzip < <(base64 --decode <<<\"$(base64 <(gzip <"$FILE") | tr -d "\n")\")"
}

function is_marshalable_file () {
    local FILE=${1:?[?FILE]}
    if [ -e "$FILE" ] && [ ! -d "$FILE" ] && [[ "$FILE" =~ ^/dev/fd/* || "$(realpath "$FILE")" =~ ^/tmp/*|^/mnt/*|^/media/*|^/home/*|^/Users/* ]]; then
        return 0
    else
        return 1
    fi
}

function match_count () {
    local NEEDLE=${1:?[?NEEDLE]}
    local HAYSTACK=${*:2}
    local X=${HAYSTACK//$NEEDLE}
    echo $(((${#HAYSTACK} - ${#X}) / ${#NEEDLE}))
}

function cat_marshalled_command () {
    local ARGS=()
    local FILE_LIST=()
    local B64_FILES=()
    local ARG_INDEX=0
    local ARG FILE_INDEX MODE FILE MTIME
    while [ -n "$1" ]; do
        ARG="$1"
        if [ "$ARG_INDEX" == "0" ]; then
            ARG=$(resolve_path "$ARG")
        fi
        if is_marshalable_file "$ARG"; then
            if [[ "$1" == /dev/fd/* ]]; then
                FILE="~dev~fd~$(basename "$ARG")"
                MODE="500"
            elif [ -p "$ARG" ]; then
                FILE="$(basename "$ARG")"
                MODE="500"
            else
                FILE="$(basename "$ARG")"
                MODE=$(stat_mode "$ARG")
            fi
            FILE="\$CTX/"$(printf '%q' "$FILE")
            MTIME=$(stat_mtime "$ARG")

            # Count the repeated file names and suffix index as necessary
            FILE_INDEX=$(match_count "$FILE" "${FILE_LIST[@]}")
            FILE_LIST+=("$FILE")
            if [ "$FILE_INDEX" != "0" ]; then
                FILE=$FILE"~$FILE_INDEX"
            fi

            B64_FILES+=("$(cat_base64_file "$ARG" ">$FILE;chmod $MODE $FILE;touch -d @$MTIME $FILE")")

            ARGS+=("$FILE")
        else
            ARGS+=($(printf '%q' "$ARG"))
        fi
        shift || break
        ARG_INDEX=$((ARG_INDEX+1))
    done

    if [ "$DEBUG" ]; then
        echo 'set -x'
    fi
    if [ "${#B64_FILES[@]}" != "0" ]; then
        echo 'CTX=$(TMPDIR="$HOME" mktemp -d -t .sshx.XXXXXXXX)'
        echo 'trap "{ rm -rf \"\$CTX\"; }" EXIT SIGINT SIGTERM SIGKILL'
        echo 'chmod 700 "$CTX"'
        printf "%s\n" "${B64_FILES[@]}"
    fi
    if [ "$WORKDIR" ]; then
        echo "cd $(printf '%q' "$WORKDIR")"
    fi
    if [ "$RCFILE" ]; then
        echo ". $(printf '%q' "$RCFILE")"
    fi

    for ENV in "${ENV_VARS[@]}"; do
        echo "export $ENV=$(printf '%q' "${!ENV}")"
    done

    echo "${ARGS[@]}"
}

function sshx () {
    local STAGE=ssh
    local SSH=()
    local CMD=()
    local ENV_VARS=()
    local DUMP DEBUG GZIP WORKDIR
    while [ -n "$1" ]; do
        if [ "$STAGE" == "ssh" ] && [ "$1" == "--" ]; then
            STAGE=cmd
            shift || break
            continue
        fi
        case "$STAGE" in
            ssh)
                case "$1" in
                    --help)
                        usage
                        return 1
                        ;;
                    --dump)
                        DUMP=/dev/stdout
                        ;;
                    --dump=*)
                        DUMP=${1/#*=/}
                        ;;
                    --debug)
                        DEBUG=true
                        ;;
                    --env)
                        ENV_VARS+=("$2")
                        shift
                        ;;
                    --env=*)
                        ENV_VARS+=("${1/#*=/}")
                        ;;
                    --gzip)
                        GZIP=true
                        ;;
                    --rcfile)
                        RCFILE=$2
                        shift
                        ;;
                    --rcfile=*)
                        RCFILE=${1/#*=/}
                        ;;
                    --workdir)
                        WORKDIR=$2
                        shift
                        ;;
                    --workdir=*)
                        WORKDIR=${1/#*=/}
                        ;;
                    *)
                        SSH+=("$1")
                        ;;
                esac
                ;;
            cmd)
                CMD+=("$1")
                ;;
        esac
        shift || break
    done

    if [ "$DUMP" ]; then
        cat_marshalled_command "${CMD[@]}" > "$DUMP"
        return
    fi

    if [ ! "${SSH[*]}" ] || [ ! "${CMD[*]}" ]; then
        usage
        exit 1
    fi

    SSHX=(ssh "${SSH[@]}")
    if [ "${SSH[*]}" ] && [ "${CMD[*]}" ]; then
        SSHX+=(/bin/bash)
        if [ "$GZIP" ]; then
            require gzip
            SSHX+=('<('$(cat_gzip_file <(cat_marshalled_command "${CMD[@]}"))')')
        else
            SSHX+=('<('$(cat_base64_file <(cat_marshalled_command "${CMD[@]}"))')')
        fi
    fi
    exec "${SSHX[@]}"
}

require base64
require tr
require printf
require basename
require stat

sshx "${@}"