Update ssh add/agent script

The script used to simply add identities at start up and perform minor
persistence of loaded identites.

It now invokes an agent if not available in order to load identities, and
ensures persistence across more useful scenarios, like `sudo su` and tmux after
that.

Does not cover all combinations, though, such as `sudo tmux` (no `su` means no
updating the root auth socket link).

Script messiness is not as bad as a lot of other ssh-agent scripts, but could
still do with some refactoring/cleaning up. For now, at least it mostly works.

zsh/custom/ssh-add.zsh

@@ -1,19 +1,70 @@
-# Add the first user identity if the local ssh-agent doesn't already have
-# something.
+# Expected location of the ssh-agent authentication socket
 #
-# We silence error output for when we're sshing to a machine that we don't want
-# to use ssh-agent forwarding for.
-if [[ `ssh-add -L | grep "^ssh-" | wc -l` -eq 0 ]]; then
-  ssh-add 2> /dev/null
+# `SSH_AUTH_SOCK` points to the socket through which users interact with
+# ssh-agent loaded identities.
+#
+# Having this socket linked to a predictable place means a given user can expect
+# loaded identities to persist
+#  - in any shell instance
+#  - in tmux, screen and similar multiplexers
+#  - through `sudo su` to root (BUT NOT with '-' or plain `su`)
+#
+# This setup helps to persist ssh agent-forwarding across remotes.
+[[ ! -z $SUDO_USER ]] && name=$SUDO_USER || name=$USER
+SOCK="/tmp/ssh-agent-$name-tmux"
+
+# Link up an existing socket to the expected location if it is not already.
+linksocket() {
+  if [[ ! -z $SSH_AUTH_SOCK && $SSH_AUTH_SOCK != $SOCK ]]; then
+    ln -sf $SSH_AUTH_SOCK $SOCK &> /dev/null
+  fi
+  export SSH_AUTH_SOCK=$SOCK
+}
+
+# Find an ssh-agent at the expected location, or start a new one.
+#
+# `ssh-add` returning a code of 2 means this user has no usable ssh-agent.
+#
+# If a new agent is started, remember to link the new socket.
+#
+# Note: ssh-add return codes:
+#   0 = success,
+#   1 = specified command fails (e.g., no keys with ssh-add -l)
+#   2 = unable to contact the authentication agent
+
+# Try to reach an existing socket.
+linksocket
+ssh-add -l &> /dev/null
+agent_reached=$?
+
+# For `sudo su`: a user cannot access a socket with another user's symlink,
+# so try to create a personal symlink to the same socket.
+# TODO Find a cleaner, simpler way to decide to create symlinks.
+if [[ $agent_reached -eq 2 ]]; then
+  SOCK2="/tmp/ssh-agent-$USER-tmux"
+  ln -sf `readlink $SOCK` $SOCK2
+  export SSH_AUTH_SOCK=$SOCK2
+  ssh-add -l &> /dev/null
+  agent_reached=$?
 fi
 
-# Put ssh authentication socket in predicatable place for ssh-agent forwarding.
+# Nothing worked. If there are any identities available, create a new agent.
 #
-# This allows us to overcome stale sockets in tmux.
-SOCK="/tmp/ssh-agent-$USER-tmux"
-if test $SSH_AUTH_SOCK && [ $SSH_AUTH_SOCK != $SOCK ]
-then
-    rm -f /tmp/ssh-agent-$USER-tmux
-    ln -sf $SSH_AUTH_SOCK $SOCK
-    export SSH_AUTH_SOCK=$SOCK
+# Defaults: ~/.ssh/id_rsa ~/.ssh/id_dsa ~/.ssh/id_ecdsa ~/.ssh/identity
+identities_exist=false
+default_files=(~/.ssh/id_rsa ~/.ssh/id_dsa ~/.ssh/id_ecdsa ~/.ssh/identity)
+for file in $default_files; do
+  [[ -f $file ]] && identities_exist=true && break
+done
+
+if [[ $identities_exist == true && $agent_reached -eq 2 ]]; then
+  echo "Starting a new SSH Agent."
+  eval `ssh-agent` &> /dev/null
+  linksocket
+fi
+
+# Add any identities with the default file names if the agent does not have
+# anything.
+if [[ $identities_exist == true && `ssh-add -L | grep "^ssh-" | wc -l` -eq 0 ]]; then
+  ssh-add
 fi