#!/usr/bin/env zsh
################################################################################
#
# This file is autoloaded by .zshrc, and actually loaded when executed
#
################################################################################
# swagent fixes ssh-/gpg-agent, and makes sure you're connected to one

#### SWAGENT

emulate -L zsh
setopt errexit nounset # be strict
autoload -Uz say && say

local swagfile=${XDG_RUNTIME_DIR}/ssh-swagent
local swagsock=${XDG_RUNTIME_DIR}/ssh-swagsock

: ${SSH_AGENT=}
: ${SSH_AGENT_PID=}
: ${SSH_AUTH_SOCK=}
: ${SSH_CLIENT=}
: ${SSH_CONNECTION=}


main() {

  case ${1:-ssh} in
    ssh|ssh-agent) swag-ssh;;
    gpg|gpg-agent) swag-gpg;;
    *) say error "swagent takes *one* argument - ssh or gpg."
  esac

}

swag-ssh() {
  (( $+commands[ssh-agent] )) || { say error "Could not find binary for ssh-agent in \$PATH"; return }

  # First check if we're already connected, in which case we don't have to do anything.
  if swagcheck "Already connected"; then
    # If we're connected, and ssh-agent is forwarded, and this is a
    # login-shell, record the details in $swagfile so we can reconnect later
    if [[ $SSH_AGENT == "forwarded" && -o login ]]; then
      print -- "export SSH_AUTH_SOCK=$SSH_AUTH_SOCK" | tee $swagfile
    fi

    # Return if we're already connected
    return 0
  fi

  # Check if details in $swagfile are valid, and connect if they are
  if [[ -f $swagfile ]]; then
    source $swagfile > /dev/null
    swagcheck "Reconnected" && return 0 || rm -f $swagfile
  fi

  # If we got here, the $swagfile was stale or non-existant
  # Launch new ssh-agent
  eval $(ssh-agent -s | sed 's/^echo/#echo/' | tee $swagfile) > /dev/null
  swagcheck "Started ssh-agent" && return

  # If we got here, something failed
  unset SSH_AGENT
  say error "Something failed"
  return 1
}

swagcheck() {
  if [[ -w $SSH_AUTH_SOCK && -n $SSH_AGENT_PID ]] \
    && ps -p $SSH_AGENT_PID &> /dev/null; then
    say info "$@ - %F{green}local%f ssh-agent pid $SSH_AGENT_PID"
  elif [[ -w $SSH_AUTH_SOCK && -n $SSH_CLIENT && -n $SSH_CONNECTION ]]; then
    say info "$@ - %F{yellow}forwarded%f ssh-agent"
    SSH_AGENT=forwarded
  else
    return 1
  fi

  if [[ ${SSH_AUTH_SOCK:A} != ${swagsock:A} ]]; then
    ln -sf $SSH_AUTH_SOCK $swagsock
    SSH_AUTH_SOCK=$swagsock
  fi

  return 0
}

swag-gpg() {
  (( $+commands[gpg-agent] && $+commands[ssh-add] )) || return
  export GPG_TTY=$TTY
  export SSH_AGENT=gpg-agent
  export SSH_AUTH_SOCK="/run/user/$UID/gnupg/S.gpg-agent.ssh"

  gpg-connect-agent updatestartuptty /bye

  # See if agent works
  ssh-add -l &>/dev/null
  if (( $? == 2 )); then
    say error "Something is wrong with your gpg-agent. Check it manually."
    say warn "SSH_AUTH_SOCK is set to ${SSH_AUTH_SOCK}, and gpg-agent has this pid:"
    pgrep gpg-agent
    exit 1
  fi
  # If we got here, it works
  say info "Attached to gpg-agent (socket: ${SSH_AUTH_SOCK})"
  return 0
}


main "${@:-}"

## END OF FILE #################################################################