#!/usr/bin/env bash
# ssh-close: close SSH master connections (ssh -O exit).
# With no hosts, every host ssh-open recorded. Closing a jump host also
# closes the recorded hosts that jump through it, those first, so none is
# left with a dead transport.
set -euo pipefail

case ${1:-} in
  -h | --help)
    echo "usage: ssh-close [HOST...]   (no HOST: every host ssh-open opened)"
    exit 0
    ;;
esac

state="${XDG_RUNTIME_DIR:-$HOME/.cache}/devkit-ssh/hosts"
recorded=()
[ ! -s "$state" ] || mapfile -t recorded <"$state"

declare -A target=()
if [ $# -eq 0 ]; then
  for h in "${recorded[@]}"; do target[$h]=1; done
else
  for h in "$@"; do target[$h]=1; done
  # Add recorded hosts that jump through a target, until nothing changes.
  changed=1
  while [ "$changed" = 1 ]; do
    changed=0
    for h in "${recorded[@]}"; do
      [ -z "${target[$h]:-}" ] || continue
      jump="$(ssh -G "$h" 2>/dev/null | awk '$1 == "proxyjump" { print $2; exit }')"
      IFS=, read -ra jumps <<<"${jump:-}"
      for j in "${jumps[@]}"; do
        if [ -n "${target[$j]:-}" ]; then target[$h]=1; changed=1; fi
      done
    done
  done
fi
[ ${#target[@]} -gt 0 ] || { echo "ssh-close: no hosts opened with ssh-open" >&2; exit 0; }

# ssh-open records a jump host before the hosts behind it, so closing in
# reverse record order closes those first; unrecorded targets go first.
order=()
for h in "$@"; do
  printf '%s\n' "${recorded[@]}" | grep -qxF "$h" || order+=("$h")
done
for ((i = ${#recorded[@]} - 1; i >= 0; i--)); do
  [ -z "${target[${recorded[i]}]:-}" ] || order+=("${recorded[i]}")
done

for host in "${order[@]}"; do
  if ssh -O check "$host" 2>/dev/null; then
    ssh -O exit "$host" 2>/dev/null && echo "$host closed"
  else
    echo "$host not open"
  fi
done

# Forget the closed hosts.
if [ ${#recorded[@]} -gt 0 ]; then
  keep=()
  for h in "${recorded[@]}"; do [ -n "${target[$h]:-}" ] || keep+=("$h"); done
  if [ ${#keep[@]} -gt 0 ]; then printf '%s\n' "${keep[@]}" >"$state"; else : >"$state"; fi
fi