#!/bin/bash
# -*- Mode:Shell-script; Coding:us-ascii-unix; fill-column:158 -*-
################################################################################################################################################################
##
# @file      ssh-agent-go.sh
# @author    Mitch Richling <https://www.mitchr.me/>
# @brief     Way to start up and use just one ssh-agent per host.@EOL
# @keywords  ssh ssh-agent key socket
# @Std       bash
# @copyright 
#  @parblock
#  Copyright (c) 2016, Mitchell Jay Richling <http://www.mitchr.me> All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
#
#  1. Redistributions of source code must retain the above copyright notice, this list of conditions, and the following disclaimer.
#
#  2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions, and the following disclaimer in the documentation
#     and/or other materials provided with the distribution.
#
#  3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without
#     specific prior written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
#  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
#  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR
#  TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#  @endparblock
# @filedetails
#
#  This script is designed to make it easy to manage a single ssh-agent per host from relatively primitive command line environments.  In normal operation, it
#  will start a new ssh-agent, or find an existing one.  Normally the output of the ssh-agent command is mimicked, so the output may be exec'ed to setup the
#  shell environment.  If the environment variable $SAGDEBUG is set, then various debugging will be printed as well -- which may be exec'ed too!  Several
#  options can change the default behavior:
#
#    -addKey  -- Add all keys to the agent if none currently loaded
#    -addKeys -- Dump all keys, and add ALL keys to the agent!
#    -find    -- Just print info about running agents
#    -kill    -- Kill running ssh-agent
#    -badA    -- Find agents not managed by this script
#    -q       -- Quite -- only print important stuff
#    -noEnv   -- Find agent without SSH_AUTH_SOCK and SSH_AGENT_PID
#    -clean   -- Remove all keys from the agent
#    -ttl val -- Set the key TTL to val
#
#  It is a simple matter to make use of this script in a shell or window manager startup file.  My preference is to embed the script into a few shell aliases:
#
#   # start/reuse ssh-agent, add key if required, set vars
#   alias sag='eval `ssh-agent-go.sh` ; ssh-agent-go.sh -q -addKey '
#   
#   # find current agent, kill it, unset vars
#   alias sak='eval `ssh-agent-go.sh -kill` '
#   
#   # Fire up a new ssh-gent or use an existing one.
#   alias ssh='eval `ssh-agent-go.sh | grep -v Agent` ; ssh-agent-go.sh -addKey -q -addKey ; ssh '
#   alias scp='eval `ssh-agent-go.sh | grep -v Agent` ; ssh-agent-go.sh -addKey -q -addKey ; scp '
#
#  Random Notes:
#
#  * Several other products provide similar functionality including the gentoo "Keychain" package, the OS X keychain, and the GNOME keychain.  These other
#    products are far less difficult to use, but don't quite work the way I want them to work and they are not universally available..
#
#  * TTL (Time To Live) philosophy: I think a 14400s token time to live is a great choice!  Only two or three passwords are required in a normal working day,
#    and we avoid tokens living for 8 or 24 hours after logout.
#
#  * This script requires several binaries: awk, egrep, grep, tail, Berkley ps, ssh-agent, sed, test, date, mkdir, chmod, echo, ssh-add, true, scutil, id
#
#  * On Cygwin this script requires procps
#
################################################################################################################################################################

#---------------------------------------------------------------------------------------------------------------------------------------------------------------
# Configuration options
dKTTL=14400         # Default Key TTL (86400=1d, 28800=8h, 14400=4h)
baseDir='/tmp'      # Directory in which to create socket dirs
subDirPfx='mjrssha' # Prefix part of socket dir names
sfilePfx='agent.'   # Prefix part of socket file names

#---------------------------------------------------------------------------------------------------------------------------------------------------------------
# process arguments
KEY_TTL=$dKTTL  # 86400=1d, 28800=8h, 14400=4h
ADD_KEYS=N      # Add key to agent if required
ADD_ALL_KEYS=N  # Dump all keys, and add all keys to the agent
JUST_FIND=N     # Print info for running agent, don't start new one
FND_BAD=''      # Find bad agents.
KILL_AGENT=N    # Kill off the agent -- must have JUST_FIND=Y too!
CLN_AGENT=N     # Clean keys from running agent -- must have JUST_FIND=Y too!
LOWPRT=N        # Don't print out the variables
USEENV=Y        # Use variables to find ssh-agent
while test -n "$*" ; do
  case "$1" in
    -ttl     ) KEY_TTL="$2"; shift                                     ; shift  ;; # 
    -addKey  ) ADD_KEYS=Y                                              ; shift  ;; # 
    -addKeys ) ADD_KEYS=Y; ADD_ALL_KEYS=Y                              ; shift  ;; # 
    -find    ) JUST_FIND=Y                                             ; shift  ;; # 
    -badA    ) FND_BAD=Y                                               ; shift  ;; # 
    -kill    ) LOWPRT=Y; JUST_FIND=Y; KILL_AGENT=Y                     ; shift  ;; # 
    -clean   ) LOWPRT=Y; JUST_FIND=Y; CLN_AGENT=Y                      ; shift  ;; # 
    -q       ) LOWPRT=Y                                                ; shift  ;; # 
    -noEnv   ) USEENV=N                                                ; shift  ;; # 
    *        ) test -n "$SAGDEBUG" && echo echo "ERROR: Bad arg: '$1'" ; shift  ;; # 
  esac
done

#---------------------------------------------------------------------------------------------------------------------------------------------------------------
# Find the ps command, and figure out what arguments we use
if [ -e '/usr/bin/procps' ] ; then
  PSCMD='/usr/bin/procps -wwwwAF'
else
  if [ -x '/usr/ucb/ps' ] ; then
    PSCMD='/usr/ucb/ps auxwwww'
  else
    PSCMD='/bin/ps auxwwww'
  fi
fi

#---------------------------------------------------------------------------------------------------------------------------------------------------------------
# Make sure key TTL is OK, and if not then set to the default
if ! echo "$KEY_TTL" | egrep '^[0-9][0-9]*$' >/dev/null ; then
  KEY_TTL=$dKTTL
fi

#---------------------------------------------------------------------------------------------------------------------------------------------------------------
# Get User Name
if test -z "$USER" ; then
  if test -n "$LOGNAME" ; then
    test -n "$SAGDEBUG" && echo echo "INFO: Getting USER from LOGNAME"
    USER=$LOGNAME
  else
    test -n "$SAGDEBUG" && echo echo "INFO: Getting USER from 'id'"
    USER=`id -un`
  fi
fi
if test -z "$USER" ; then
  test -n "$SAGDEBUG" && echo echo "ERROR: Could not figure out user name"
  exit
fi
test -n "$SAGDEBUG" && echo echo "INFO: USER='$USER'"

#---------------------------------------------------------------------------------------------------------------------------------------------------------------
# Validate value of $SSH_AGENT_PID if we are using the environment
BAD_SSH_AGENT_PID=Y
if test "$USEENV" = "Y" -a -n "$SSH_AGENT_PID" ; then
  test -n "$SAGDEBUG" && echo echo "INFO: Found SSH_AGENT_PID environment variable: '$SSH_AGENT_PID'"
  if test -n "`$PSCMD 2>/dev/null | grep " ssh-agent .* -a $baseDir/$subDirPfx" | grep $USER | grep -v grep | awk '{print $2}' | grep ^$SSH_AGENT_PID\$`" ; then
    test -n "$SAGDEBUG" && echo echo "INFO: Found running ssh-agent with PID in SSH_AGENT_PID"
    BAD_SSH_AGENT_PID=N
  else
    test -n "$SAGDEBUG" && echo echo "WARNING: No running ssh-agent with PID in SSH_AGENT_PID"
  fi
fi

#---------------------------------------------------------------------------------------------------------------------------------------------------------------
# Validate value of $SSH_AUTH_SOCK if we are using the environment
BAD_SSH_AUTH_SOCK=Y
if test "$USEENV" = "Y" -a -n "$SSH_AUTH_SOCK" ; then
  if test -n "`$PSCMD 2>/dev/null | grep " ssh-agent .* -a $baseDir/$subDirPfx" | grep $USER | grep $SSH_AUTH_SOCK | grep -v grep`" ; then
    test -n "$SAGDEBUG" && echo echo "INFO: Found valid SSH_AUTH_SOCK environment variable value: '$SSH_AUTH_SOCK'"
    BAD_SSH_AUTH_SOCK=N
  else
    test -n "$SAGDEBUG" && echo echo "WARNING:: Found bad SSH_AUTH_SOCK environment variable value: '$SSH_AUTH_SOCK'"
  fi
  if test "$BAD_SSH_AUTH_SOCK" = "N" -a -r "$SSH_AUTH_SOCK" ; then
    test -n "$SAGDEBUG" && echo echo "INFO: Socket file in SSH_AUTH_SOCK exists and is readable: '$SSH_AUTH_SOCK'"
  else
    BAD_SSH_AUTH_SOCK=Y
    test -n "$SAGDEBUG" && echo echo "WARNING: Socket file in SSH_AUTH_SOCK doesn't exist: '$SSH_AUTH_SOCK'"
  fi
fi

#---------------------------------------------------------------------------------------------------------------------------------------------------------------
# Find our agents, Report if we find, or start up a new one
if test "$USEENV" = "N" -o "$BAD_SSH_AUTH_SOCK" = "Y" -o "$BAD_SSH_AGENT_PID" = "Y" ; then
  test -n "$SAGDEBUG" && echo echo "INFO: Looking for ssh-agent PID the hard way"
  saPID=''
  test -n "$SAGDEBUG" && echo echo "INFO: Looking for ssh-agent PID with 'ps'"
  saPID=`$PSCMD 2>/dev/null | grep " ssh-agent .* -a $baseDir/$subDirPfx" | grep $USER | grep -v grep | awk '{print $2}' | tail -1`
  test -n "$SAGDEBUG" && echo echo "INFO: Looking for ssh-agent AUTH_SOCK the hard way"
  sshAgentSocFile=''
  test -n "$SAGDEBUG" && echo echo "INFO: Looking for ssh-agent AUTH_SOCK with 'ps'"
  sshAgentSocFile=`$PSCMD 2>/dev/null | grep " ssh-agent .* -a $baseDir/$subDirPfx" | grep $USER | grep -v grep | tail -1 | sed "s/^.*-a //"`
  test -n "$SAGDEBUG" && echo echo "INFO: Search result: Candidate SSH_AGENT_PID: $saPID"
  test -n "$SAGDEBUG" && echo echo "INFO: Search result: Candidate SSH_AUTH_SOCK: $sshAgentSocFile"
  # Now we start up a new ssh-agent or report on what we found.
  if test -z "$saPID" -o -z "$sshAgentSocFile" ; then
    test -n "$SAGDEBUG" && echo echo "INFO: No valid ssh-agent found"
    if test "$JUST_FIND" = "Y"; then
      if test "$LOWPRT" = "N"; then
        echo "unset SSH_AUTH_SOCK;"
        echo "unset SSH_AGENT_PID;"
      fi
    else
      test -n "$SAGDEBUG" && echo echo "INFO: Starting new ssh-agent"
      ddate=`date '+%Y%m%d%H%M%S'`
      keyDir="$baseDir/$subDirPfx$ddate$RANDOM"
      sshAgentSocFile="$keyDir/$sfilePfx$ddate"
      mkdir $keyDir
      chmod u=rwx $keyDir
      ssh-agent -t $KEY_TTL -s -a $sshAgentSocFile
    fi
  else
    test -n "$SAGDEBUG" && echo echo "INFO: Found existing ssh-agent"
    if test "$LOWPRT" = "N"; then
      echo "SSH_AUTH_SOCK=$sshAgentSocFile; export SSH_AUTH_SOCK;"
      echo "SSH_AGENT_PID=$saPID; export SSH_AGENT_PID;"
      echo "echo Agent pid $saPID;"
    fi
  fi
else
  test -n "$SAGDEBUG" && echo echo "INFO: Found running ssh-agent via environment variables!"
  sshAgentSocFile=$SSH_AUTH_SOCK
  saPID=$SSH_AGENT_PID
  if test "$LOWPRT" = "N"; then
    echo "SSH_AUTH_SOCK=$sshAgentSocFile; export SSH_AUTH_SOCK;"
    echo "SSH_AGENT_PID=$saPID; export SSH_AGENT_PID;"
    echo "echo Agent pid $saPID;"
  fi
fi

#---------------------------------------------------------------------------------------------------------------------------------------------------------------
# Add keys to the agent if we were requested to do so
if test "$ADD_KEYS" = "Y" ; then
  if test -n "$sshAgentSocFile" ; then
    export SSH_AUTH_SOCK=$sshAgentSocFile
    if test "$ADD_ALL_KEYS" = "Y" ; then
      test -n "$SAGDEBUG" && echo echo "INFO: Dumping all keys from the ssh-agent"
      ssh-add -D
    fi
    if ssh-add -l >/dev/null ; then
      test -n "$SAGDEBUG" && echo echo "INFO: The agent has keys already.  No need to add more."
    else
      KEYFILES=''
      for pat in 'id_dsa_*_*.prv' 'id_dsa_*_*.priv' 'id_dsa' 'id_rsa_*_*.prv' 'id_rsa_*_*.priv' 'id_rsa'; do
        if ls "$HOME/.ssh/"$pat >/dev/null 2>&1 ; then
          test -n "$SAGDEBUG" && echo echo "INFO: Found key files to add to ssh-agent: $HOME/.ssh/$pat"
          ssh-add "$HOME/.ssh/"$pat
          KEYFILES='FOUND'
        fi
      done
      if test -z "$KEYFILES" ; then
        test -n "$SAGDEBUG" && echo echo "INFO: Did not find any key files to add to ssh-agent"
      fi
    fi
  fi
fi

#---------------------------------------------------------------------------------------------------------------------------------------------------------------
# Kill ssh-agent if we were requested to do so
if test "$KILL_AGENT" = "Y" ; then
  if test -n "$saPID" ; then
    test -n "$SAGDEBUG" && echo echo "INFO: Killing agent. PID: '$saPID'"
    export SSH_AGENT_PID=$saPID
    ssh-agent -k
  fi
fi

#---------------------------------------------------------------------------------------------------------------------------------------------------------------
# Clean all keys from ssh-agent if we were requested to do so
if test "$CLN_AGENT" = "Y" ; then
  if test -n "$sshAgentSocFile" ; then
    test -n "$SAGDEBUG" && echo echo "INFO: Cleaning all keys out of ssh-agent."
    export SSH_AUTH_SOCK=$sshAgentSocFile
    ssh-add -D
  fi
fi

#---------------------------------------------------------------------------------------------------------------------------------------------------------------
# If $FND_BAD is non-NULL, then we look for agents without our special socket files
if test -n "$FND_BAD" ; then
  test -n "$SAGDEBUG" && echo echo "INFO: Looking for bad agents"
  badAgents=`$PSCMD 2>/dev/null | grep 'ssh-agent' | grep $USER | grep -v 'ssh-agent-go' | grep -v " ssh-agent .* -a $baseDir/$subDirPfx" | grep -v grep | sed 's/^/#   /'`
  if test -n "$badAgents" ; then
    echo
    echo "#++++++++++++++++"
    echo "#BAD AGENTS FOUND"
    echo "#^^^^^^^^^^^^^^^^"
    echo $badAgents
  fi
fi