#!/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

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

swagent() {
  case ${1:-ssh} in
    ssh|ssh-agent)  _swag-ssh ${@:2};;
    gpg|gpg-agent)  _swag-gpg;;
    restart)        _swag-restart ${2:-ssh};;
    env)            env | grep -iE "^SSH|^GPG|^GNUPG";;
    -w)             _swag-ssh $@;;
    *)              say -r1 error "swagent takes one argument - ssh or gpg. Or restart, with ssh or gpg as a second argument."
  esac
}

_swag-restart() {
  local agent=${${1:-ssh}%-agent}
  if [[ $1 == (ssh|gpg) ]]; then
    pkill $agent && say "Killed ${agent}-agent"
    _swag-$agent
  else
    say -r1 error "swagent restart takes one argument - ssh or gpg."
  fi
}

_swag-ssh() {
  typeset -gi WRITE=0
  local swagfile=${XDG_RUNTIME_DIR}/ssh-swagent
  local swagsock=${XDG_RUNTIME_DIR}/ssh-swagsock

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

  [[ ${1:-} == '-w' ]] && WRITE=1
  # First check if we're already connected, in which case we don't have to do anything.
  if _swag-check "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
    _swag-check "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
  _swag-check "Started ssh-agent" && return

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

_swag-check() {
  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 (( WRITE )); then
    print -r "SSH_AUTH_SOCK=$SSH_AUTH_SOCK" > $swagfile
    [[ -n $SSH_AGENT_PID ]] && print -r "SSH_AGENT_PID=$SSH_AGENT_PID" >> $swagfile
    say info "Wrote new swagfile":
    swagent env
  fi

  if [[ ${SSH_AUTH_SOCK:A} != ${swagsock:A} ]]; then
    ln -sf $SSH_AUTH_SOCK $swagsock
    export 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
}

swagent "${@:-}"

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