persistent

#!/bin/sh
# This script was previously known as 'smart-ssh' and 'persistent-ssh'
iam="${0##*/}"
help() {
  if [ "$1" = install ]; then help_install; fi
  cat <</help
Persistent SSH connection for Linux systems that use NetworkManager
Usage: $iam [ssh|xpra] [OPTIONS] DESTINATION [COMMAND]
       $iam pump [--install|PATTERN]
       $iam help [install|run]
       source <($iam --completion)

COMMAND is not recommended unless you're content letting it restart

See \`$iam help install\` for how to auto-pump on network restarts (like
waking from sleep) plus tab completion configuration and instructions on
how to properly disengage a connection and prevent looping right back into
a new ssh/xpra session.

Part of net-scripts: https://github.com/adamhotep/net-scripts
persistent 4.2.20231215.1 Copyright 2008+ Adam Katz, GPLv2+/BSD 2-clause
/help
  exit
}

help_install() {
  local PAGER="${PAGER:-less}"
  # $message_text is a variable so we can count its lines for the pager.
  # Note that it switches from single to double quotes to get variables' values
  message_text='DEPENDENCIES

Currently, '"$iam"' supports SSH and xpra commands. This guide assumes you
know about SSH. Xpra ("screen for X11") is documented at https://xpra.org
and can be installed with on Debian/Ubuntu/Mint systems with `apt install xpra`
or `dnf install xpra` on Fedora/Red Hat/Alma systems.

'"$iam"' has the ability to "pump" (restart) its connections on demand
(see AUTO-PUMP below), which requires `pgrep`. Install with `apt install procps`or `dnf install procps`.

Auto-pumping requires NetworkManager, which most Linux systems use to manage
their network interfaces. Install this with `apt install network-manager`
or `dnf install NetworkManager` but beware: that might be a big move.


AUTO-PUMP

To install the auto-pump mechanism, run `'"$iam"' pump --install`.
This adds a hook to NetworkManager that reestablishes your `'"$iam"'`
connections with each network startup rather than waiting for timeouts,
which is particularly useful in restarting connections upon waking from sleep.

This uses `sudo` to add a simple dispatcher script for NetworkManager to call
when the network is reconnected (see `man 8 NetworkManager-dispatcher`), to live
at `'"$pumper"'` or, for older systems,
`'"$old_pumper"'`.


TAB COMPLETION

You can configure tab completion for '"$iam"' arguments.
Just run `source <('"$iam"' --completion)` in either bash or zsh.

To install this in your ~/.bashrc, set it up to load on your first completion:

    _persistent() { source <("'"$iam"'" --completion); _persistent "$@"; }
    complete -F _persistent "'"$iam"'"

To install this in your ~/.zshrc, use the same function with compdef:

    _persistent() { source <("'"$iam"'" --completion); _persistent "$@"; }
    compdef _persistent "'"$iam"'"


PROPER EXIT STATUS FOR RESTARTING CONNECTIONS

'"$iam"' will automatically restart any managed ssh/xpra session unless it
exits with code 42. You need to alter the configuration for each of your
REMOTE hosts to detect a clean exit and convert it from 0 to 42.

These examples support Bash and Z shell (Zsh). Unless you use both shells,
you only need to create/modify files that start with `.bash` or with `.z`

~/.bashrc
    # for use in ~/.bash_logout
    get_epoch() { [ -n "$EPOCHSECONDS" ] && echo "$EPOCHSECONDS" || date +%s; }
    _start="$(get_epoch)"

~/.zshrc or ~/.zlogin
    # for use in ~/.zlogout
    get_epoch() { print -P "%D{%s}"; }
    _start="$(get_epoch)"

~/.bash_logout and ~/.zlogout (at the very top)
    retval=$?	# this must be the first line

~/.bash_logout and ~/.zlogout (at the very bottom)
    # For `'"$iam"' ssh ...` usage. See `'"$iam"' help install`
    # or https://github.com/adamhotep/net-scripts/blob/master/persistent
    _end="$(get_epoch)"
    # If this is a login shell and it returned cleanly and we have a start time
    # and the shell ran for > 5 minutes and the SSH TTY is the current TTY
    if [[ $((SHLVL + retval)) == 1 && $_start -gt 0 \
    && $((_end - _start)) -gt 300 && $SSH_TTY == "$(tty)" ]]; then
      exit 42	# inform parent (`'"$iam"' ssh`) this was a clean exit
    fi
    exit $retval'
  if [ ! -t 1 ] \
  || [ "${LINES:-$(tput lines)}" -gt $(echo "$message_text" |grep -c ^) ]; then
    PAGER=cat
  fi
  echo "$message_text" |$PAGER
  exit
}

# Usage: warn MESSAGE
# Display MESSAGE on standard error
warn() { echo "$*" >&2; }

# Usage: die MESSAGE [ERROR]
# Display MESSAGE on standard error and then exit with ERROR (default=2)
die() { warn "$1"; exit ${2:-2}; }

# Usage: title_say MESSAGE
# Report MESSAGE to the titlebar and standard output
# ASSUMPTION: your terminal supports titles
title_say() { printf "\e]0;%s\e\\ %s\n" "$*" "$*"; }

# Usage: sleep TIME
# As the standard sleep command but never report errors
sleep() { command sleep "$@" 2>/dev/null; }

# ASSUMPTION: the `date` command supports +%s for seconds since epoch
epoch() { date +%s; }

# Usage: we_have COMMAND [COMMAND...]
# silently returns true when we have all COMMANDs
we_have() { command -v "$@" >/dev/null 2>&1; }

# got_net: silently return whether we can resolve & connect to google.com:443
if we_have ncat
  then got_net() { ncat --send-only --recv-only -w 333ms google.com 443; }
  else got_net() { nc -zw1 google.com 443; }
fi

# on_ac_power: silently return whether we're plugged into AC power
if ! we_have on_ac_power; then
  on_ac_power() { return 1; }	# dummy: assume we're NOT plugged in
fi

completion_disclaimer() {
  if [ -t 1 ]; then
    warn "# Good move. It is wise to review \`source\` code before running it."
    warn "# For proper invocation instructions, see \`$iam help install\`."
  fi
}

# automatically detect the shell and provide the correct completion function
completion() {
  case $SHELL in
    ( *bash* )	bash_completion ;;
    ( *zsh* )	zsh_completion ;;
  esac
  die "Couldn't detect shell; try --bash-completion or --zsh-completion"
  exit
}

# print instructions on how to complete in bash
# (won't work for zsh since bashcompinit doesn't implement _command_offset)
bash_completion() {
  completion_disclaimer
  cat <<'  /bash_completion'
    _persistent() {
      local cur prev words cword split root_command
      _init_completion -s || return

      case "${words[1]}" in
        ( *ssh | *ssh2 | *xpra )
          root_command=${words[1]}
          _command_offset 1
          return
          ;;
      esac
      case "$prev" in
        ( pump | --pump ) COMPREPLY=($(compgen -W '--install' -- "$cur")) ;;
        ( help | --help ) COMPREPLY=($(compgen -W 'install run' -- "$cur")) ;;
        ( * ) COMPREPLY=($(compgen -W 'help pump ssh xpra --help
          --completion --bash-completion --zsh-completion' -- "$cur")) ;;
      esac
    }
  /bash_completion
  printf '    complete -F _persistent "%s"\n' "$iam"
  exit
}

# print instructions (for sourcing) on how to complete in zsh
zsh_completion() {
  completion_disclaimer
  cat <<'  /zsh_completion'
    if type xpra >/dev/null && ! type "$_comps[xpra]" >/dev/null; then
      # No xpra completion? Here is a quick-and-dirty version via --help
      _xpra_subcommands=( $(
        xpra --help |awk '/^Options/{exit} $1=="xpra"{print $2}'
      ) )
      _xpra() {
        _describe subcommands _xpra_subcommands
        _gnu_generic xpra
      }
      compdef _xpra xpra
    fi
    _persistent() {
      local line state
      _arguments -C "1: :->cmds" "*::arg:->args"
      case "$state" in
        ( cmds )
          _values "persistent command" \
            "help[get help or installation documentation]" \
            "pump[pump all or matching connections]" \
            "ssh[persistent SSH secure shell]" \
            "xpra[persistent remote X application viewer]"
          _arguments -S \
            '--bash-completion[show bash completion code]' \
            '--completion[auto-detect bash/zsh, show its completion code]' \
            '--help[get help or installation documentation]' \
            '--zsh-completion[show Z shell completion code]'
          ;;
        ( args )
          case $line[1] in
            ( help )
              _arguments -C "1: :->cmds"
              if [ "$state" = "cmds" ]; then
                _values "help" \
                  "run[Basic command-line help (default)]" \
                  "install[How to install auto-pump and remote shell exit cues]"
              fi
              ;;
            ( pump )
              _arguments -C "1: :->cmds" "*::arg:->args"
              case "$state" in ( cmds )
                _values pattern ":strings:pattern:"
                _arguments -S \
                  --install'[Install auto-pump on network configuration]'
                ;;
              esac
              ;;
            ( *ssh | *ssh2 )
              words[1]=(ssh); service=ssh; "$_comps[ssh]" ;;
            ( *xpra )
              words[1]=(xpra); service=xpra; "$_comps[xpra]" ;;
          esac
        ;;
      esac
    }
  /zsh_completion
  printf '    compdef _persistent "%s"\n' "$iam"
  exit
}

# Verify that we have pgrep
need_pgrep() {
  if ! we_have pgrep; then
    warn 'You need `pgrep` from procps to pump connections.'
    die  "See \`$iam help install\` for help installing dependencies."
  fi
}

# Pump: kill children of other instances so they can restart
pump() {
  case "$1" in
    ( install | --install* ) install_pump ;;
  esac
  need_pgrep
  touch "$user_pumpfile"
  if [ -n "$1" ]; then	# pump only matching commands
    { pgrep $iam |xargs -n1 -I@ pgrep -f -P @ "$1" |xargs kill; } 2>/dev/null
  else			# pump all
    { pgrep $iam |xargs -n1 pgrep -P |xargs kill; } 2>/dev/null
  fi
  exit $?
}

install_pump() {
  local sudo=''
  need_pgrep
  if [ ! -d "${pumper%/*}" ]; then
    warn "You seem to lack a '${pumper%/*}' directory."
    warn "Let's try again with the older '${old_pumper%/*}' system."
    pumper="$old_pumper"
    pumper_old=1
  fi
  if [ ! -d "${pumper%/*}" ]; then
    warn "You seem to${pumper_old:+ also} lack a '${pumper%/*}' directory."
    die  "I don't think your system is compatible with this technique :-("
  fi
  if ! touch "$pumper" 2>/dev/null; then
    warn "You need more permissions to install $pumper"
    read -p "Try with sudo [Yn]? " sudo >&2
    if [ "$sudo" = "${sudo#[Nn0]}" ]; then	# user didn't say "no"
      if ! we_have sudo; then die "Please install sudo first."; fi
      sudo=sudo
    else
      die "Aborted."
    fi
  fi
  if [ -s "$pumper" ]; then action="Overwriting existing"; fi
  echo "${action:-Installing} $pumper ..."
  if [ -n "$action" ]; then sleep 2; fi	# 'oops' buffer for Ctrl+C action
  {
    echo '#!/bin/sh'
    echo '# Silently kill children of persistent sessions so they restart'
    echo "# Created by $0"
    echo '# https://github.com/adamhotep/net-scripts/blob/master/persistent'
    if [ -z "$pumper_old" ]; then
      echo 'case "$2" in'
      echo '  ( up | connectivity-change | vpn-up | dhcp?-change ) : act ;;'
      echo '  ( * ) exit ;;'
      echo 'esac'
    fi
    echo "{ pgrep '$iam' |xargs -n1 pgrep -P |xargs kill; } 2>/dev/null"
    echo 'true'
  } | $sudo tee "$pumper" >/dev/null
  $sudo chmod 755 "$pumper"
  exit
}

# Usage: was_pumped [MIN]
# wait for net, return true if user pumped or net/system restarted in last MIN
was_pumped() {
  local min="${1:-1}" wait_time=1
  echo
  while ! got_net 2>/dev/null; do
    if [ $wait_time -gt 4096 ] && ! on_ac_power; then # no power? abort at 1h
      if ! [ "$ret" -gt 1 ] 2>/dev/null; then ret=1; fi
      warn "$iam: 12 tests without network, connection lost >= 68 minutes ago"
      die  "$iam: no AC power detected, exiting with failure code $ret" $ret
    elif [ $wait_time -gt 1 ]; then
      printf '\r%s: no network, waiting for %ds...' "$iam" $wait_time >&2
    fi
    sleep $wait_time
    # 1 2 4 8 16 32 1m 2m 4m 9m 17m 34m 1h 2h 4h ...
    wait_time=$((wait_time * 2))
  done
  if [ $wait_time -gt 1 ]; then echo; fi	# complete 'no network' line
  min=$((min + (wait_time + 30) / 60))	# add wait time in rounded minutes

  # return true if we found a pump request, the system was suspended, or
  # the network was refreshed via NetworkManager within N minutes (default N=1)
  [ -n "$(find "$user_pumpfile" "$suspend_log" "$system_pumpdir" \
            -not -type d -mmin -$min 2>/dev/null)" ]
}

old_pumper="/etc/network/if-up.d/persistent"
pumper="/etc/NetworkManager/dispatcher.d/persistent"
user_pumpfile="/tmp/.persistent-pump-$HOSTNAME-$LOGNAME"

# ASSUMPTION: The most recently modified file in [/var]/run/NetworkManager/
#             is representative of the last network negotiation
#             and/or [/var]/run/suspend.log or else /var/log/suspend.log
#             is touched when resuming from sleep
system_pumpdir="/run/NetworkManager"
suspend_log="/run/suspend.log"
if ! [ -d "$system_pumpdir" ]; then system_pumpdir="/var$system_pumpdir"; fi
if ! [ -d "$suspend_log" ]; then suspend_log="/var$suspend_log"; fi
if ! [ -d "$suspend_log" ]; then suspend_log="/var/log/${suspend_log##*/}"; fi

# get the command (we currently support no options, so that's all merged)
case "$1" in
  ( --completion )		completion ;;
  ( --bash-completion ) 	bash_completion ;;
  ( --zsh-completion )		zsh_completion ;;
  ( help | -h | --help* )	shift; help "$@" ;;
  ( install-pump | --instal* )	install_pump "$@" ;;
  ( pump | --pump )		shift; pump "$@" ;;
  ( pump=* | --pump=* ) 	pump "${1#*=}" ;;
  ( *xpra )			cmd=xpra ;;
  ( *ssh | *ssh2 )		cmd=ssh ;;
  ( * ) 			die "$iam: Invalid command '$1'" ;;
esac

ret=0

# initial self-report (on delay after a restart)
title_say "$iam $*"

while true; do
  "$@"
  ret=$?

  # Clean exit, see `persistent help install` (help_install() above)
  if [ $ret = 42 ]; then
    exit	# stop everything, return that clean exit 0
  fi

  # Print five dots in 1.665 seconds (or 2s if sleep is too basic).
  for i in 1 2 3 4 5; do
    printf .
    sleep 0.333 || sleep 2 && break
  done
  if ! was_pumped 1; then echo ''; exit $ret; fi
  if [ "$cmd" != ssh ]; then sleep 1.5 || sleep 2; fi	# wait for ssh
  title_say "reconnecting $*"
  ( sleep 6; title_say "$iam $*"; ) &
done

exit $ret