#!/bin/zsh
# Load YubiKey PIV SSH keys into the macOS launchd-managed ssh-agent.
# The default health check verifies the PIV AUTH key with ssh-add -T so
# listed-but-stale PKCS#11 keys are reloaded instead of being trusted.
set -u

PKCS11_LIB="${PKCS11_LIB:-/usr/local/lib/opensc-pkcs11.so}"
PIV_KEY_PATTERN="${PIV_KEY_PATTERN:-PIV AUTH pubkey|SIGN pubkey|KEY MAN pubkey|CARD AUTH pubkey}"
PIV_AUTH_PATTERN="${PIV_AUTH_PATTERN:-PIV AUTH pubkey}"
VERIFY_MODE="auth"
TMP_DIRS=()
SCRIPT_NAME="${0:t}"

usage() {
  # Print the small CLI surface; normal operation intentionally has few knobs.
  cat <<EOF
Usage: ${SCRIPT_NAME} [--no-verify] [--help]

Safely load YubiKey PIV keys into the macOS launchd ssh-agent.

Options:
  --no-verify  Only check whether PIV keys are listed in the agent.
  --help       Show this help.

Environment:
  PKCS11_LIB     PKCS#11 provider path. Default: /usr/local/lib/opensc-pkcs11.so
EOF
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --no-verify)
      VERIFY_MODE="none"
      ;;
    --help|-h)
      usage
      exit 0
      ;;
    *)
      print -u2 -- "load-yubikey-piv-ssh: unknown option: $1"
      usage >&2
      exit 2
      ;;
  esac
  shift
done

cleanup() {
  # Remove temporary public-key files created for ssh-add -T verification.
  local dir
  for dir in "${TMP_DIRS[@]}"; do
    [[ -n "${dir}" && -d "${dir}" ]] && rm -rf -- "${dir}"
  done
}
trap cleanup EXIT INT TERM

die() {
  # Report a fatal script error with a stable prefix for shell output.
  print -u2 -- "load-yubikey-piv-ssh: $*"
  exit 1
}

note() {
  # Emit status messages to stderr so stdout stays clean for callers.
  print -u2 -- "load-yubikey-piv-ssh: $*"
}

agent_socket() {
  # Return the launchd-published ssh-agent socket if it exists.
  local sock
  sock="$(launchctl getenv SSH_AUTH_SOCK 2>/dev/null || true)"
  [[ -n "${sock}" && -S "${sock}" ]] || return 1
  print -r -- "${sock}"
}

use_launchd_agent() {
  # Force this shell to use Apple's launchd-managed agent socket and ignore
  # stale SSH_AGENT_PID values from manually started agents.
  local sock
  sock="$(agent_socket)" || die "could not find launchd SSH_AUTH_SOCK"
  export SSH_AUTH_SOCK="${sock}"
  unset SSH_AGENT_PID
}

agent_has_piv_keys() {
  # Check for any listed YubiKey PIV identities in the active agent.
  ssh-add -l 2>/dev/null | egrep -q "${PIV_KEY_PATTERN}"
}

loaded_piv_public_keys() {
  # Select loaded public keys for either health verification or no-verify mode.
  local pattern
  case "${VERIFY_MODE}" in
    auth) pattern="${PIV_AUTH_PATTERN}" ;;
    none) pattern="${PIV_KEY_PATTERN}" ;;
    *) die "internal error: unknown verify mode: ${VERIFY_MODE}" ;;
  esac

  ssh-add -L 2>/dev/null | egrep "${pattern}" || true
}

verify_loaded_piv_keys() {
  # Export the loaded PIV AUTH public key to a temp file and ask ssh-add to
  # perform a local sign-and-verify operation through the agent.
  [[ "${VERIFY_MODE}" == "none" ]] && return 0

  local tmp key_file line count
  local -a key_files

  tmp="$(mktemp -d "${TMPDIR:-/tmp}/load-yubikey-piv-ssh.XXXXXX")" ||
    die "failed to create temporary directory"
  TMP_DIRS+=("${tmp}")

  count=0
  while IFS= read -r line; do
    [[ -n "${line}" ]] || continue
    count=$((count + 1))
    key_file="${tmp}/piv-${count}.pub"
    print -r -- "${line}" >| "${key_file}" ||
      die "failed to write temporary public key"
    chmod 600 "${key_file}" || die "failed to chmod temporary public key"
    key_files+=("${key_file}")
  done < <(loaded_piv_public_keys)

  if [[ "${count}" -eq 0 ]]; then
    note "no matching PIV AUTH public key found for verification"
    return 1
  fi

  SSH_ASKPASS_REQUIRE=never ssh-add -T "${key_files[@]}" >/dev/null 2>&1
}

agent_has_healthy_piv_keys() {
  # A healthy state requires both listed PIV keys and a successful sign test.
  if ! agent_has_piv_keys; then
    return 1
  fi

  if verify_loaded_piv_keys; then
    return 0
  fi

  return 1
}

agent_responds() {
  # Treat "no identities" as a working agent, but reject broken sockets.
  ssh-add -l >/dev/null 2>&1
  case "$?" in
    0|1) return 0 ;;  # 1 means a reachable agent with no identities.
    *) return 1 ;;
  esac
}

restart_launchd_agent() {
  # Restart only the launchd-managed ssh-agent. macOS may block kickstart under
  # SIP, so fall back to terminating the exact PID reported by launchctl.
  local pid
  local i

  note "restarting launchd ssh-agent"

  if launchctl kickstart -k "gui/${UID}/com.openssh.ssh-agent" >/dev/null 2>&1; then
    sleep 1
    use_launchd_agent
    return 0
  fi

  pid="$(launchctl print "gui/${UID}/com.openssh.ssh-agent" 2>/dev/null |
    awk '/pid = / { print $3; exit }')"

  if [[ -n "${pid}" ]]; then
    note "launchctl kickstart failed; terminating launchd ssh-agent pid ${pid}"
    kill -TERM "${pid}" >/dev/null 2>&1 ||
      die "failed to terminate launchd ssh-agent pid ${pid}"

    for i in {1..10}; do
      kill -0 "${pid}" >/dev/null 2>&1 || break
      sleep 0.2
    done

    if kill -0 "${pid}" >/dev/null 2>&1; then
      note "launchd ssh-agent pid ${pid} did not exit after TERM; sending KILL"
      kill -KILL "${pid}" >/dev/null 2>&1 ||
        die "failed to kill launchd ssh-agent pid ${pid}"
    fi
  else
    note "launchctl kickstart failed and no running launchd ssh-agent pid was found"
  fi

  use_launchd_agent

  if ! agent_responds; then
    die "failed to restart launchd ssh-agent"
  fi
}

load_provider() {
  # Load the OpenSC PKCS#11 provider and force PIN entry on the terminal.
  SSH_ASKPASS_REQUIRE=never ssh-add -s "${PKCS11_LIB}"
}

verify_or_report_stale() {
  # Treat listed PIV keys as healthy only after verification succeeds.
  if agent_has_healthy_piv_keys; then
    case "${VERIFY_MODE}" in
      none) note "PIV keys are already loaded" ;;
      auth) note "PIV AUTH key is already loaded and verified" ;;
    esac
    return 0
  fi

  if agent_has_piv_keys; then
    note "PIV keys are listed but failed verification; treating agent as stale"
  fi

  return 1
}

[[ -r "${PKCS11_LIB}" ]] || die "PKCS#11 library is not readable: ${PKCS11_LIB}"

use_launchd_agent

if verify_or_report_stale; then
  exit 0
fi

if agent_has_piv_keys; then
  restart_launchd_agent
elif ! agent_responds; then
  restart_launchd_agent
fi

if load_provider && agent_has_healthy_piv_keys; then
  case "${VERIFY_MODE}" in
    none) note "PIV keys loaded" ;;
    auth) note "PIV keys loaded and PIV AUTH key verified" ;;
  esac
  exit 0
fi

note "initial load failed; checking agent state"

if verify_or_report_stale; then
  exit 0
fi

restart_launchd_agent

if load_provider && agent_has_healthy_piv_keys; then
  case "${VERIFY_MODE}" in
    none) note "PIV keys loaded after agent restart" ;;
    auth) note "PIV keys loaded after agent restart and PIV AUTH key verified" ;;
  esac
  exit 0
fi

die "failed to load and verify PIV keys"