tunnels.sh

#!/usr/bin/env bash
# Usage: sub tunnels [start|stop|restart|status|edit]
# Summary: Manage shared and user-specified SSH tunnels into your infrastructure
# Help: Manage shared and user-specified SSH tunnels into your infrastructure.
#
# Expects ssh to be configured for proxying. Run `sub check` to check your
# environment.
#
# This will attempt to use autossh if it's available.
#
# On regular ssh control sockets will be stored in ~/.sub/tunnels.
# On autossh pids will be stored in ~/.sub/autosshpids.
#
# If you need to add or remove shared tunnels, just modify the tunnels() function.
#
# You can also add personal-use tunnels to ~/.sub/user_tunnels.sh.
# Run `sub tunnels edit` to open or create an example tunnel configuration
# file.

set -e

tunnels() {
  [[ -e "${user_tunnels_file}" ]] && source "${user_tunnels_file}"

  tunnel SOMEHOST \
      -L 8000:localhost:80
}

# Starts the tunnels.
start() {
  if [[ "${autossh_available}" = true ]]; then
    echo "Using autossh." >&2
    mkdir -p "${pid_dir}"
    tunnels
  else
    echo "Using regular ssh, no autossh available." >&2
    mkdir -p "${tunnel_dir}"
    tunnels | xargs -L 1 -P 50 ssh
  fi
}

# Sends exit signals to all sockets in the tunnel dir.
stop() {
  if [[ "${autossh_available}" = true ]]; then
    for pidfile in $(find ${pid_dir} -depth 1); do
      echo "Stopping `basename ${pidfile}`..."
      kill `cat "${pidfile}"`
    done
  else
    for socket in $(find ${tunnel_dir} -depth 1); do
      echo "Stopping `basename ${socket}`..."
      ssh -S "${socket}" -O exit localhost 2>/dev/null
    done
  fi
}

autossh_status() {
  for pidfile in $(find ${pid_dir} -not -type d); do
    local pid=`cat "${pidfile}"`

    echo "Checking `basename ${pidfile}`:"
    if kill -0 $pid 2>/dev/null; then
      echo "Autossh running (pid=${pid})"
    else
      echo "Autossh not running, cleaning up old pidfile ${pidfile}"
      rm "${pidfile}"
    fi
  done
}

# Queries status of running tunnels
status() {
  if [[ "${autossh_available}" = true ]]; then
    autossh_status
  else
    for socket in $(find ${tunnel_dir} -not -type d); do
      echo "Checking `basename ${socket}`:"
      ssh -S "${socket}" -O check localhost
    done
  fi
}

# Starts the tunnel using autossh
autossh_tunnel() {
  local host="$1"
  local pidfile="${pid_dir}/${host}"
  shift

  if [[ -e "${pidfile}" ]]; then
    echo "Tunnel for ${host} is already started." >&2
  else
    echo "${host} tunnel starting..." >&2

    export AUTOSSH_PIDFILE="${pidfile}"
    autossh ${base_ssh_options} -M 0 \
      "$@" \
      ${host}
  fi
}

# Outputs the ssh switches for the specified tunnel.
ssh_switches() {
  local host="$1"
  local socket="${tunnel_dir}/${host}"
  shift

  if [[ -e "${socket}" ]]; then
    echo "Tunnel for ${host} is already started." >&2
  else
    echo "${host} tunnel starting..." >&2

    echo ${base_ssh_options} -o ExitOnForwardFailure=yes \
      -M -S"${tunnel_dir}/${host}" \
      "$@" \
      ${host}
  fi
}

tunnel() {
  if [[ "${autossh_available}" = true ]]; then
    autossh_tunnel "$@"
  else
    ssh_switches "$@"
  fi
}

# Opens the user tunnels file
edit() {
  if [[ -z "${EDITOR}" ]]; then
    echo "Your EDITOR environment variable is not set, please run \`sub check\` for insubuctions." >&2
    exit 1
  fi

  if [[ ! -e "${user_tunnels_file}" ]]; then
    cat > "${user_tunnels_file}" <<SH
# This is an example tunnel
#
# tunnel SOMEHOST \\
#     -L 8001:localhost:8080
#
SH
  fi

  ${EDITOR} "${user_tunnels_file}"
}

# Provide sub completions
if [ "$1" = "--complete" ]; then
  echo start
  echo stop
  echo restart
  echo status
  echo edit
  exit
fi

# environment initialization
if command -v autossh >/dev/null 2>&1; then
  autossh_available=true
fi

tunnel_dir="${HOME}/.sub/tunnels"
pid_dir="${HOME}/.sub/autosshpids"
base_ssh_options="-f -nNT -o ControlPath=/tmp -o ControlMaster=no"
user_tunnels_file="${HOME}/.sub/user_tunnels.sh"

case $1 in
  start)
    start
    ;;
  stop)
    stop
    ;;
  restart)
    stop
    start
    ;;
  status)
    status
    ;;
  edit)
    edit
    ;;
  *)
    exec sub help tunnels
esac