screen/tmux wrapper script

S

#!/bin/zsh -f

# TODO: systemd support
# systemd is transitioning to terminating processes when a user session ends,
# so we'll need to start a new scope.
#
# Launching the initial screen/tmux should be prefixed with:
#   systemd-run --scope --user   [$CMD...]
#
# This should go into plexer_new() as a contextual prefix.

: "${S_SCREEN_PREFIX:=s-command}"
: "${S_SCREEN:=default}"
: "${S_SCREEN_AGENTDIR:=${HOME}/.ssh/agents}"

# We explicitly allow this to be set-but-empty
: "${S_SCREEN_sshagent_lifetime=24h}"

progname="${0:t}"
have_screen=0
have_tmux=0

setopt re_match_pcre

function die { print -u2 -r -- "${progname}: $*"; exit 1; }

screen_version="$(screen -v 2>&1 || true)"
if [[ $screen_version == 'Screen version'* ]];
then
	have_screen=1
else
	tmux_version="$(tmux -V 2>&1 || true)"
	if [[ $tmux_version == tmux\ * ]]; then
		have_tmux=1
	fi
fi

unset SHLVL

case $TERM in
(putty|xterm*)
	function xtitle { print >/dev/tty -n -- "\e]2;$*\a"; }
	;;
(*)
	function xtitle { : ; }
	;;
esac

# These are effectively the same thing, as we're implemented now,
# but conceptually the use-cases differ, so we're _exposing_ two
# options which map to the same functionality.  We reserve the right
# to change the functionality of either, as long as it aligns with
# the name.
want_agent=0
want_ifneeded_agent=0
if [[ -n $1 && $1 == -startagent ]]; then
	shift
	want_agent=1
elif [[ -n $1 && $1 == -autoagent ]]; then
	shift
	want_ifneeded_agent=1
fi

# hey, it's a convenience wrapper, right?  Might as well make the now-common
# desire to "attach, but let me create a new window or look" nicer.
if [[ -n $1 && $# -eq 1 && ( $1 == '=' || $1 == '-' ) ]]; then
	set -- -p "$1"
fi

if [[ -n $1 && $1 != -* ]]; then
	S_SCREEN="$1"
	shift
fi

if (( have_screen )); then
	function plexer_set_list__screen {
	  screens=($(
	    screen -ls | sed -n $'s/^[ \t][ \t]*[0-9][0-9]*\.\\([^ \t]*\\)[ \t].*$/\\1/p'
	    ))
	}
	function plexer_reattach {
	  local n="$1"; shift
	  screen -x -r "$n" "$@"
	}
	function plexer_new {
	  local n="$1"; shift
	  screen -S "$n" "$@"
	}
	function plexer_have {
	  local n="$1"; shift
	  plexer_set_list__screen
	  for s in $screens; do
	    if [[ $s == $want ]]; then
	      return 0
	    fi
	  done
	  return 1
	}
	want="$S_SCREEN_PREFIX.$S_SCREEN"

elif (( have_tmux )); then
	function plexer_reattach {
	  local n="$1"; shift
	  tmux attach-session -t "$n" "$@"
	}
	function plexer_new {
	  local n="$1"; shift
	  tmux new-session -s "$n" "$@"
	}
	function plexer_have {
	  local n="$1"; shift
	  tmux has-session -t $n 2>/dev/null
	}
	want="$S_SCREEN_PREFIX-$S_SCREEN"

else
	die "coding error; neither screen nor tmux available"
fi

if zmodload -i zsh/pcre 2>/dev/null ; then
  function pid_from_agentfile {
    settings="$(< "$S_SCREEN_AGENTFILE")"
    if [[ $settings =~ '(?m)^SSH_AGENT_PID=(\d+)' ]]; then
      printf "%s\n" "${match[1]}"
    fi
  }
elif pcregrep --help >/dev/null 2>&1 ; then
  function pid_from_agentfile {
    pcregrep -h -o1 -e '(?m)^SSH_AGENT_PID=(\d+)' -- "$S_SCREEN_AGENTFILE"
  }
else
  function pid_from_agentfile {
    sed -n -e 's/^SSH_AGENT_PID=\([0-9][0-9]*\);.*$/\1/p' < "$S_SCREEN_AGENTFILE"
  }
fi

S_SCREEN_AGENTFILE="$S_SCREEN_AGENTDIR/screen.$want"
if [[ -f "$S_SCREEN_AGENTFILE" ]]; then
  export S_SCREEN_AGENTFILE
fi

function cleanup {
  exited=$1
  xtitle "${LOGNAME}@${HOST}"
  if [[ ! -f "$S_SCREEN_AGENTFILE" ]]; then
    exit $exited
  fi
  pid=$(pid_from_agentfile)
  if [[ -z $pid ]] || ! kill -0 $pid 2>/dev/null; then
    exit $exited
  fi
  if plexer_have $want; then
    exit $exited
  fi

  # At this point, if the screen still existed, we'd have exited
  # So: screen gone, agent is alive
  echo "Cleaning up ssh-agent: pid $pid"
  kill $pid
  rm -f -- "$S_SCREEN_AGENTFILE"
  exit $exited
}

xtitle screen $want ${HOST%%.*}

if plexer_have $want; then
  plexer_reattach $s "$@"
  cleanup $?
fi

if (( want_agent )) || (( want_ifneeded_agent )); then
  mkdir -p -m 0700 $S_SCREEN_AGENTDIR
  agent_args=(-s)
  if [[ -n "${S_SCREEN_sshagent_lifetime:-}" ]]; then
    agent_args+=(-t "${S_SCREEN_sshagent_lifetime:-}")
  fi
  ssh-agent ${agent_args} | sed '/^echo/d' > "$S_SCREEN_AGENTFILE"
  export S_SCREEN_AGENTFILE
fi

plexer_new $want "$@"
cleanup $?

I use this with two shell aliases:

alias S=' command S'
alias S.=' exec command S -autoagent ='