#!/bin/bash - not strictly necessary, but helps nano with syntax highlighting.

__find_ssh_agent_sock() {
  # Find an *active* ssh agent socket.
  # Returns:	0 = Found an active socket.
  #		1 = Did not find a viable socket.
  local I

  # Search the SSH_AUTH_SOCKS array for a viable socket.
  for ((I = 0; I < ${#SSH_AUTH_SOCKS[@]}; I++)); do
    [[ "${SSH_AUTH_SOCKS[$I]}" =~ ^[[:blank:]]*$ ]] && continue
    SSH_AUTH_SOCK="${SSH_AUTH_SOCKS[$I]}" ssh-add -l >/dev/null 2>&1
    (( $? <= 1 )) && {
      export SSH_AUTH_SOCK="${SSH_AUTH_SOCKS[$I]}"
      return 0
    }
    unset SSH_AUTH_SOCK
  done

  return 1
}

__read_ssh_agents() {
  # Read all the known ssh agent sockets into an array.
  # Returns:	0 = Processed and read the agents file without issue.
  #		1 = Error processing/reading the agents file.
  local ERR I SOCK

  [[ ! -e "$HOME/.ssh/agents" ]] && touch "$HOME/.ssh/agents"

  # Lock the ~/.ssh/agents file.
  if [[ "$PLATFORM" == "Linux" ]]; then
    # Linux has 'flock', thankfully.
    exec 9<"$HOME/.ssh/agents" && flock -E 10 -e -w 0.5 9
    ERR=$?
    if (( ERR == 10 )); then
      printf "\\033[1;31m%s\\033[0m\\n" "Failed to obtain lock on ~/.ssh/agents." >&2
      return 1
    elif (( ERR > 0 )); then
      printf "\\033[1;31m%s\\033[0m\\n" "Flock usage error." >&2
      return 1
    fi

    # Make note of the mtime for use in write_ssh_agents.
    SSH_AGENTS_MTIME="$(stat --format=%.9Y "$HOME/.ssh/agents")"
  elif [[ "$PLATFORM" == "Darwin" ]]; then
    # Do locking the sucky way on Darwin.
    for ((I = 0; I <= 5; I++)); do
      if shlock -p "$$" -f "$HOME/.ssh/agents.lock"; then
        exec 9<"$HOME/.ssh/agents"
        # Make note of the mtime for use in write_ssh_agents.
        SSH_AGENTS_MTIME="$(stat -f %Fm "$HOME/.ssh/agents")"
        ERR=0
        break
      else
        ERR=1
        sleep 0.1
      fi
    done
    (( ERR != 0 )) && {
      printf "\\033[1;31m%s\\033[0m\\n" "Failed to obtain lock on ~/.ssh/agents." >&2
      return 1
    }
  else
    printf "\\033[1;31m%s\\033[0m\\n" "File locking unsupported on '$PLATFORM'." >&2
    return 1
  fi

  # Read the socket list (bash v2+ compliant)
  while read -u 9 -r SOCK; do
    [[ -n "$SOCK" ]] && SSH_AUTH_SOCKS+=("$SOCK")
  done
  ERR=$?

  # Close the file descriptor (which on Linux releases the flock too).
  exec 9<&-

  # On Darwin, release the lock on the file.
  rm -f  "$HOME/.ssh/agents.lock"

  # Remove the . in the mtime.
  SSH_AGENTS_MTIME="${SSH_AGENTS_MTIME/\.}"

  # Error out if the data couldn't be read.
  (( ERR != 0 )) && {
    printf "\\033[1;31m%s\\033[0m\\n" "Failed to read ssh-agent socket list." >&2
    unset SSH_AUTH_SOCKS SSH_AGENTS_MTIME
    return 1
  }

  return 0
}

__ssh_agent_prompt_command() {
  # If necessary, find and activate a new ssh agent socket before each prompt is displayed.
  # Returns:	0 = All is good.
  #		1 = An error occured.
  local ERR

  if [[ -z "$SSH_AUTH_SOCK" ]]; then
    ERR=2
  else
    ssh-add -l >/dev/null 2>&1
    ERR=$?
  fi
  (( ERR == 2 )) && {
    # Read alternative sockets from ~/.ssh/agents.
    __read_ssh_agents || {
      unset SSH_AUTH_SOCK
      return 1
    }

    # Find a new socket to use.
    if __find_ssh_agent_sock; then
      printf "\\033[1;33m%s\\033[0m\\n" "Connected to existing ssh-agent socket."
      sleep 0.5
    else
      # Start a new agent.
      eval "$(ssh-agent -s 2>/dev/null | grep -v 'echo'; printf "%s" "ERR=${PIPESTATUS[0]}")"
      (( ERR > 0 )) && {
        printf "\\033[1;31m%s\\033[0m\\n" "Failed to start new ssh-agent - continuing with no agent."
        sleep 0.5
        return 1
      }
      printf "\\033[1;32m%s\\033[0m\\n" "Started new ssh-agent."
      __write_ssh_agents
      (( $? == 2 )) && __read_ssh_agents && __write_ssh_agents
      sleep 0.5
    fi
  }

  return 0
}

__write_ssh_agents() {
  # Write all unique ssh agent sockets into the ~/.ssh/agents file.
  # Returns:	0 = Processed and wrote the agents file without issue.
  #		1 = Error processing/writing the agents file.
  #		2 = The SSH_AUTH_SOCKS array may be out of date as the agents file'a mtime has changed.
  local ERR I J MTIME SOCKS

  # Add the current agent socket to the sockets array.
  SSH_AUTH_SOCKS=("$SSH_AUTH_SOCK" "${SSH_AUTH_SOCKS[@]}")

  # Remove any duplicates from SSH_AUTH_SOCKS.
  for ((I = 0; I < ${#SSH_AUTH_SOCKS[@]}; I++)); do
    [[ -z "${SSH_AUTH_SOCKS[$I]}" ]] && continue
    for ((J = 0; J < ${#SOCKS[@]}; J++)); do
      [[ "${SSH_AUTH_SOCKS[$I]}" == "${SOCKS[$J]}" ]] && continue 2
    done
    # Only add the socket if it's still viable.
    SSH_AUTH_SOCK="${SSH_AUTH_SOCKS[$I]}" ssh-add -l >/dev/null 2>&1
    (( $? <= 1 )) && SOCKS+=("${SSH_AUTH_SOCKS[$I]}")
  done

  # Lock the ~/.ssh/agents file.
  if [[ "$PLATFORM" == "Linux" ]]; then
    # Make sure SSH_AUTH_SOCKS has the most up to date data.
    MTIME="$(stat --format=%.9Y "$HOME/.ssh/agents")"
    (( ${MTIME/\.} > SSH_AGENTS_MTIME )) && return 2

    # Lock the agents file.
    exec 9<"$HOME/.ssh/agents" && flock -E 10 -e -w 0.5 9
    ERR=$?
    if (( ERR == 10 )); then
      printf "\\033[1;31m%s\\033[0m\\n" "Failed to obtain lock on ~/.ssh/agents." >&2
      return 1
    elif (( ERR > 0 )); then
      printf "\\033[1;31m%s\\033[0m\\n" "Flock usage error." >&2
      return 1
    fi
  elif [[ "$(uname -s)" == "Darwin" ]]; then
    # Make sure SSH_AUTH_SOCKS has the most up to date data.
    MTIME="$(stat -f %Fm "$HOME/.ssh/agents")"
    (( ${MTIME/\.} > SSH_AGENTS_MTIME )) && return 2

    # Do locking the sucky way on Darwin.
    for ((I = 0; I <= 5; I++)); do
      if shlock -p "$$" -f "$HOME/.ssh/agents.lock"; then
        exec 9<"$HOME/.ssh/agents"
        ERR=0
        break
      else
        ERR=1
        sleep 0.1
      fi
    done
    (( ERR != 0 )) && {
      printf "\\033[1;31m%s\\033[0m\\n" "Failed to obtain lock on ~/.ssh/agents." >&2
      return 1
    }
  else
    printf "\\033[1;31m%s\\033[0m\\n" "File locking unsupported on '$PLATFORM'." >&2
    return 1
  fi

  # Write the cleaned array to disk.
  ERR=-1
  [[ -n "${SOCKS[*]}" ]] && { printf "%s\\n" "${SOCKS[@]}" >"$HOME/.ssh/agents" 2>/dev/null; ERR=$?; }

  # Release locks.
  exec 9<&-
  rm -f  "$HOME/.ssh/agents.lock"

  # Error out if the data couldn't be written.
  if (( ERR == -1 )); then
    rm -f "$HOME/.ssh/agents" 2>/dev/null
  elif (( ERR >= 1 )); then
    rm -f "$HOME/.ssh/agents" 2>/dev/null
    printf "\\033[1;31m%s\\033[0m\\n" "Failed to write ssh-agent socket list." >&2
    return 1
  fi

  return 0
}

# Add the ssh-agent prompt function call to PROMPT_COMMAND.
PROMPT_COMMAND+="__ssh_agent_prompt_command;"

# Auto start the ssh agent.
__ssh_agent_prompt_command

# Prompt to add keys to ssh-agent with scp/sftp/ssh commands.
hash scp ssh ssh-add >/dev/null 2>&1 && alias scp='_SSHEXEC=scp ssh'
hash sftp ssh ssh-add >/dev/null 2>&1 && alias sftp='_SSHEXEC=sftp ssh'
hash ssh ssh-add >/dev/null 2>&1 && ssh() {
  local ERR

  if [[ -z "$SSH_AUTH_SOCK" ]]; then
    ERR=2
  else
    ssh-add -l >/dev/null 2>&1
    ERR=$?
  fi
  if (( ERR == 1 )); then
    ssh-add
  elif (( ERR == 2 )); then
    __ssh_agent_prompt_command
    ssh-add
  fi
  command "${_SSHEXEC:-ssh}" "$@"
  unset _SSHEXEC
}