From 5a9d9553cd6861d80cc958132ab5af40fe661ac4 Mon Sep 17 00:00:00 2001
From: Simone Gaiarin <simgunz@gmail.com>
Date: Thu, 26 Aug 2021 11:42:17 +0200
Subject: [PATCH] feat(ssh-agent): allow using external helper to ask for
 passwords (#7631)

---
 plugins/ssh-agent/README.md            |   9 ++
 plugins/ssh-agent/ssh-agent.plugin.zsh | 128 ++++++++++++++-----------
 2 files changed, 83 insertions(+), 54 deletions(-)

diff --git a/plugins/ssh-agent/README.md b/plugins/ssh-agent/README.md
index 8765a9c7..f46e8bf6 100644
--- a/plugins/ssh-agent/README.md
+++ b/plugins/ssh-agent/README.md
@@ -55,6 +55,15 @@ ssh-add -K -c -a /run/user/1000/ssh-auth <identities>
 
 For valid `ssh-add` arguments run `ssh-add --help` or `man ssh-add`.
 
+----
+
+To set an **external helper** to ask for the passwords and possibly store
+them in the system keychain use the `helper` style. For example:
+
+```zsh
+zstyle :omz:plugins:ssh-agent helper ksshaskpass
+```
+
 ## Credits
 
 Based on code from Joseph M. Reagle: https://www.cygwin.com/ml/cygwin/2001-06/msg00537.html
diff --git a/plugins/ssh-agent/ssh-agent.plugin.zsh b/plugins/ssh-agent/ssh-agent.plugin.zsh
index d45406f6..9bc2e8a2 100644
--- a/plugins/ssh-agent/ssh-agent.plugin.zsh
+++ b/plugins/ssh-agent/ssh-agent.plugin.zsh
@@ -1,56 +1,76 @@
 typeset _agent_forwarding _ssh_env_cache
 
 function _start_agent() {
-	local lifetime
-	zstyle -s :omz:plugins:ssh-agent lifetime lifetime
+  local lifetime
+  zstyle -s :omz:plugins:ssh-agent lifetime lifetime
 
-	# start ssh-agent and setup environment
-	echo Starting ssh-agent...
-	ssh-agent -s ${lifetime:+-t} ${lifetime} | sed 's/^echo/#echo/' >! $_ssh_env_cache
-	chmod 600 $_ssh_env_cache
-	. $_ssh_env_cache > /dev/null
+  # start ssh-agent and setup environment
+  echo Starting ssh-agent...
+  ssh-agent -s ${lifetime:+-t} ${lifetime} | sed 's/^echo/#echo/' >! $_ssh_env_cache
+  chmod 600 $_ssh_env_cache
+  . $_ssh_env_cache > /dev/null
 }
 
 function _add_identities() {
-	local id line sig lines
-	local -a identities loaded_sigs loaded_ids not_loaded
-	zstyle -a :omz:plugins:ssh-agent identities identities
+  local id line sig lines
+  local -a identities loaded_sigs loaded_ids not_loaded
+  zstyle -a :omz:plugins:ssh-agent identities identities
 
-	# check for .ssh folder presence
-	if [[ ! -d $HOME/.ssh ]]; then
-		return
-	fi
+  # check for .ssh folder presence
+  if [[ ! -d $HOME/.ssh ]]; then
+    return
+  fi
 
-	# add default keys if no identities were set up via zstyle
-	# this is to mimic the call to ssh-add with no identities
-	if [[ ${#identities} -eq 0 ]]; then
-		# key list found on `ssh-add` man page's DESCRIPTION section
-		for id in id_rsa id_dsa id_ecdsa id_ed25519 identity; do
-			# check if file exists
-			[[ -f "$HOME/.ssh/$id" ]] && identities+=$id
-		done
-	fi
+  # add default keys if no identities were set up via zstyle
+  # this is to mimic the call to ssh-add with no identities
+  if [[ ${#identities} -eq 0 ]]; then
+    # key list found on `ssh-add` man page's DESCRIPTION section
+    for id in id_rsa id_dsa id_ecdsa id_ed25519 identity; do
+      # check if file exists
+      [[ -f "$HOME/.ssh/$id" ]] && identities+=($id)
+    done
+  fi
 
-	# get list of loaded identities' signatures and filenames
-	if lines=$(ssh-add -l); then
-		for line in ${(f)lines}; do
-			loaded_sigs+=${${(z)line}[2]}
-			loaded_ids+=${${(z)line}[3]}
-		done
-	fi
+  # get list of loaded identities' signatures and filenames
+  if lines=$(ssh-add -l); then
+    for line in ${(f)lines}; do
+      loaded_sigs+=${${(z)line}[2]}
+      loaded_ids+=${${(z)line}[3]}
+    done
+  fi
 
-	# add identities if not already loaded
-	for id in $identities; do
-		# check for filename match, otherwise try for signature match
-		if [[ ${loaded_ids[(I)$HOME/.ssh/$id]} -le 0 ]]; then
-			sig="$(ssh-keygen -lf "$HOME/.ssh/$id" | awk '{print $2}')"
-			[[ ${loaded_sigs[(I)$sig]} -le 0 ]] && not_loaded+="$HOME/.ssh/$id"
-		fi
-	done
+  # add identities if not already loaded
+  for id in $identities; do
+    # check for filename match, otherwise try for signature match
+    if [[ ${loaded_ids[(I)$HOME/.ssh/$id]} -le 0 ]]; then
+      sig="$(ssh-keygen -lf "$HOME/.ssh/$id" | awk '{print $2}')"
+      [[ ${loaded_sigs[(I)$sig]} -le 0 ]] && not_loaded+=("$HOME/.ssh/$id")
+    fi
+  done
 
-	local args
-	zstyle -a :omz:plugins:ssh-agent ssh-add-args args
-	[[ -n "$not_loaded" ]] && ssh-add "${args[@]}" ${^not_loaded}
+  # abort if no identities need to be loaded
+  if [[ ${#not_loaded} -eq 0 ]]; then
+    return
+  fi
+
+  # pass extra arguments to ssh-add
+  local args
+  zstyle -a :omz:plugins:ssh-agent ssh-add-args args
+
+  # use user specified helper to ask for password (ksshaskpass, etc)
+  local helper
+  zstyle -s :omz:plugins:ssh-agent helper helper
+
+  if [[ -n "$helper" ]]; then
+    if [[ -z "${commands[$helper]}" ]]; then
+      echo "ssh-agent: the helper '$helper' has not been found."
+    else
+      SSH_ASKPASS="$helper" ssh-add "${args[@]}" ${^not_loaded} < /dev/null
+      return $?
+    fi
+  fi
+
+  ssh-add "${args[@]}" ${^not_loaded}
 }
 
 # Get the filename to store/lookup the environment from
@@ -60,21 +80,21 @@ _ssh_env_cache="$HOME/.ssh/environment-$SHORT_HOST"
 zstyle -b :omz:plugins:ssh-agent agent-forwarding _agent_forwarding
 
 if [[ $_agent_forwarding == "yes" && -n "$SSH_AUTH_SOCK" ]]; then
-	# Add a nifty symlink for screen/tmux if agent forwarding
-	[[ -L $SSH_AUTH_SOCK ]] || ln -sf "$SSH_AUTH_SOCK" /tmp/ssh-agent-$USERNAME-screen
+  # Add a nifty symlink for screen/tmux if agent forwarding
+  [[ -L $SSH_AUTH_SOCK ]] || ln -sf "$SSH_AUTH_SOCK" /tmp/ssh-agent-$USERNAME-screen
 elif [[ -f "$_ssh_env_cache" ]]; then
-	# Source SSH settings, if applicable
-	. $_ssh_env_cache > /dev/null
-	if [[ $USERNAME == "root" ]]; then
-		FILTER="ax"
-	else
-		FILTER="x"
-	fi
-	ps $FILTER | grep ssh-agent | grep -q $SSH_AGENT_PID || {
-		_start_agent
-	}
+  # Source SSH settings, if applicable
+  . $_ssh_env_cache > /dev/null
+  if [[ $USERNAME == "root" ]]; then
+    FILTER="ax"
+  else
+    FILTER="x"
+  fi
+  ps $FILTER | grep ssh-agent | grep -q $SSH_AGENT_PID || {
+    _start_agent
+  }
 else
-	_start_agent
+  _start_agent
 fi
 
 _add_identities