1 files changed, 126 insertions, 72 deletions

@@ -1,126 +1,175 @@
 #!/bin/sh
 # mrgrctrnl: configurable ssh tunneler
 # Author: Case Duckworth <acdw@acdw.net>
-# Version: 0.2
+# Version: 0.3
 # License: MIT
 
 PRGN="${0##*/}"
 
-usage() {
-        cat <<END
-$PRGN: make magic ssh tunnels
-usage:  $PRGN [-h] [-k]
-        $PRGN [-n] [-c CONF | -s SSH]
-
-options:
-        -h      show this help
-        -k      kill all processes and exit
-        -r      restart $PRGN
-        -n      do a dry run: just print what would happen
-        -c CONFIG       use a different CONFIG file instead of
-                        \$XDG_CONFIG_HOME/$PRGN/config
-        -s CMDLINE      directly input an ssh CMDLINE --
-                        don't load the config file.
-END
-}
-
 # entry point
 main() {
+        # flags
+        _RUN=true
+        _LOG=true
+        _USE_CONFIG=true
+        # options
         __CONFIG="${XDG_CONFIG_HOME:-$HOME/.config}/$PRGN/config"
-        __SSHCMD="-N \"\$user@\$machine\" -L \"\$local:\$remote\""
         __PIDF="/tmp/$PRGN.pid"
-        _run=true
-        _quiet=false
-        _use_config=true
+        __CMDF="/tmp/$PRGN.cmd"
+        __SSH="$(command -v autossh || command -v ssh)"
+        __SSH_ARGS="-nT -N \"\$user@\$machine\" -L \"\$local:\$remote\""
 
-        while getopts hqkrnc:s: OPT; do
+        # save command invocation (for -r)
+        printf '%s %s\n' "$PRGN" "$*" >"$__CMDF"
+
+        # parse options
+        while getopts hkrnqSAc:s: OPT; do
                 case "$OPT" in
                 h)
                         usage
                         exit 0
                         ;;
-                q) _quiet=true ;;
                 k)
-                        tear_down
+                        untunnel
                         pkill -x "$PRGN"
-                        exit
+                        exit "$?"
                         ;;
                 r)
-                        tear_down
-                        exec "$PRGN"
+                        untunnel
+                        exec "$(cat "$__CMDF")"
                         ;;
-                n) _run=false ;;
+                n) _RUN=false ;;
+                q) _LOG=false ;;
+                S) __SSH='ssh' ;;
+                A) __SSH='autossh' ;;
                 c) __CONFIG="$OPTARG" ;;
                 s)
-                        __SSHCMD="$OPTARG"
-                        _use_config=false
-                        ;;
-                *)
-                        usage
-                        exit 1
+                        __SSH_ARGS="$OPTARG"
+                        _USE_CONFIG=false
                         ;;
+                *) exit 4 ;;
                 esac
         done
+        shift $((OPTIND - 1))
 
-        if $_use_config && [ ! -r "$__CONFIG" ]; then
-                echo "!! $PRGN: Cannot find config file '$__CONFIG'."
-                echo "Aborting."
+        # make sure ssh is installed
+        if ! command -v "$__SSH" >/dev/null; then
+                log "Not installed: $__SSH"
                 exit 2
         fi
 
-        if $_use_config; then
+        # make sure the config is readable
+        if $_USE_CONFIG && [ ! -r "$__CONFIG" ]; then
+                log "Cannot find config file $__CONFIG; aborting."
+                exit 3
+        fi
+
+        # load the config
+        if $_USE_CONFIG; then
                 awk '{sub(/#.*$/,"");print}' "$__CONFIG" |
                         while read -r machine user local remote key rest; do
-                                if [ -z "$machine" ] || [ -z "$user" ] ||
-                                        [ -z "$local" ] || [ -z "$remote" ]; then
+                                if [ -z "$machine" ] ||
+                                        [ -z "$user" ] ||
+                                        [ -z "$local" ] ||
+                                        [ -z "$remote" ]; then
                                         continue
                                 fi
 
+                                case "$local" in
+                                :*) local="localhost$local" ;;
+                                esac
+                                case "$remote" in
+                                :*) remote="localhost$remote" ;;
+                                esac
+
                                 # shellcheck disable=2030
                                 if [ -n "$key" ]; then
-                                        __SSHCMD="$__SSHCMD -i \"$key\""
+                                        __SSH_ARGS="$__SSH_ARGS -i \"$key\""
                                 fi
                                 if [ -n "$rest" ]; then
-                                        __SSHCMD="$__SSHCMD $rest"
+                                        __SSH_ARGS="$__SSH_ARGS $rest"
                                 fi
 
-                                eval "tunnel $__SSHCMD" &
+                                eval "tunnel_${__SSH##*/} $__SSH_ARGS" &
                         done
         else
-                # shellcheck disable=2031
-                eval "tunnel $__SSHCMD" &
+                #shellcheck disable=2031
+                eval "tunnel_${__SSH##*/} $__SSH_ARGS" &
         fi
 
         wait
 }
 
-tunnel() {
-        # shellcheck disable=2086
-        $_quiet || echo ssh "$@"
-        $_run && {
-                touch "$__PIDF"
-                while [ -e "$__PIDF" ]; do
-                        grep -q -e "$*" "$__PIDF" && {
-                                sleep 3
-                                continue
-                        }
-                        # shellcheck disable=2015
-                        if ssh "$@" & then
-                                printf '%s\t%s\n' "$!" "$*" >>"$__PIDF"
-                        else
-                                badpid="$!"
-                                kill "$badpid"
-                                sed "/^$badpid/d" "$__PIDF" >"$__PIDF~" &&
-                                        mv "$__PIDF~" "$__PIDF"
-                                continue
-                        fi
+usage() {
+        cat <<END
+$PRGN: make magic ssh tunnels
+usage:  $PRGN [-h] [-k] [-r]
+        $PRGN [-n] [-q] [-S|-A] [-c FILE | -s ARGS]
+
+options:
+        -h      show this help
+        -k      demolish all tunnels and exit
+        -r      demolish all tunnels and restart
+        -n      dry run: just print what would happen
+        -q      quiet: don't print anything
+
+        -S      force use of 'ssh'
+        -A      force use of 'autossh'
+                default: use 'autossh' if installed, otherwise 'ssh'
+
+        -c FILE load FILE as config file
+                default:        \$XDG_CONFIG_HOME/$PRGN/config
+
+        -s ARGS run CMD with ARGS, do not read config file
+                default:        '-N USER@MACHINE -L LOCAL:REMOTE'
+                                (populated from config file)
+
+config example:
+        # a hash mark comments to the end of the line
+        # if either 'local' or 'remote' begin with a colon,
+        # they'll be prefixed with 'localhost'
+        # machine       user    local   remote  key     other
+        example.com     bill    :1299   :6667   # key and other are optional
+END
+}
+
+tunnel_ssh() {
+        set -- ssh "$@"
+        log "$@"
+        $_RUN || return
+
+        touch "$__PIDF"
+        while [ -e "$__PIDF" ]; do # quit if the pid file is removed
+                # make sure we haven't tunneled this connection already
+                if grep -q -e "$*" "$__PIDF"; then
                         sleep 3
-                done
-        }
+                        continue
+                fi
+
+                # try to start ssh tunnel
+                if "$@" & then
+                        printf '%s\t%s\n' "$!" "$*" >>"$__PIDF"
+                else
+                        badpid="$!"
+                        kill "$badpid"
+                        sed "/^$badpid/d" "$__PIDF" >"$__PIDF~" &&
+                                mv "$__PIDF~" "$__PIDF"
+                        continue
+                fi
+                sleep 3
+        done
 }
 
-tear_down() {
-        printf '%s...' "Killing extant tunnels"
+tunnel_autossh() {
+        set -- autossh -f "$@"
+        log "$@"
+        $_RUN || return
+
+        "$@" && printf '%s\t%s\n' "$!" "$*" >"$__PIDF"
+}
+
+untunnel() {
+        log "Killing tunnels"
         if [ -e "$__PIDF" ]; then
                 while read -r pid _; do
                         kill "$pid" 2>/dev/null
@@ -128,8 +177,13 @@ tear_down() {
                 rm "$__PIDF"
         else
                 pkill -x ssh
+                pkill -x autossh
         fi
-        echo Done.
+        log "Done"
+}
+
+log() {
+        $_LOG && printf '%s: %s\n' "$PRGN" "$*" >&2
 }
 
 main "$@"