#!/bin/bash

# blob:summary=Set up the OpenSSH server, open the firewall, and authorize an SSH key
# blob:args=[--key=<public-key>]
# blob:examples=blob-setup-security-sshd | blob-setup-security-sshd --key="ssh-ed25519 AAAA... user@host"
# blob:requires-sudo=true

set -e

AUTHORIZED_KEYS="$HOME/.ssh/authorized_keys"
KEY=""

for arg in "$@"; do
  case "$arg" in
  --key=*) KEY="${arg#--key=}" ;;
  -h | --help)
    echo "Usage: blob-setup-security-sshd [--key=<public-key>]"
    echo
    echo "Sets up the OpenSSH server, opens the SSH port in the UFW firewall,"
    echo "and authorizes an SSH key (from GitHub, pasted, or passed via --key)."
    exit 0
    ;;
  *)
    echo "blob-setup-security-sshd: unknown option '$arg'. Try --help." >&2
    exit 2
    ;;
  esac
done

setup_sshd() {
  echo "Installing and starting the OpenSSH server..."
  blob-pkg-add openssh
  sudo systemctl enable --now sshd.service
}

open_firewall() {
  if blob-cmd-missing ufw; then
    echo "UFW is not installed; skipping firewall rule."
    return
  fi

  echo "Opening the SSH port in the firewall (rate limited against brute force)..."
  sudo ufw limit 22/tcp comment "blob-sshd" >/dev/null
  sudo ufw reload >/dev/null
}

valid_key() {
  ssh-keygen -lf /dev/stdin <<<"$1" >/dev/null 2>&1
}

authorize_key() {
  local key="$1"

  if ! valid_key "$key"; then
    echo -e "\e[31mNot a valid SSH public key: $key\e[0m" >&2
    return 1
  fi

  mkdir -p "$HOME/.ssh"
  chmod 700 "$HOME/.ssh"
  touch "$AUTHORIZED_KEYS"
  chmod 600 "$AUTHORIZED_KEYS"

  if grep -qxF "$key" "$AUTHORIZED_KEYS"; then
    echo "Key already authorized: $(ssh-keygen -lf /dev/stdin <<<"$key")"
  else
    echo "$key" >>"$AUTHORIZED_KEYS"
    echo "Authorized key: $(ssh-keygen -lf /dev/stdin <<<"$key")"
  fi
}

authorize_keys_from_github() {
  local username keys added=0

  username=$(gum input --prompt "GitHub username> " --placeholder "dhh") || exit 1
  if [[ -z $username ]]; then
    echo -e "\e[31mNo GitHub username given.\e[0m" >&2
    exit 1
  fi

  echo "Fetching keys from https://github.com/$username.keys..."
  if ! keys=$(curl -fsSL "https://github.com/$username.keys") || [[ -z $keys ]]; then
    echo -e "\e[31mCould not fetch any SSH keys for GitHub user '$username'.\e[0m" >&2
    exit 1
  fi

  while IFS= read -r key; do
    [[ -z $key ]] && continue
    authorize_key "$key" && added=$((added + 1))
  done <<<"$keys"

  if (( added == 0 )); then
    echo -e "\e[31mNo valid SSH keys found for GitHub user '$username'.\e[0m" >&2
    exit 1
  fi
}

authorize_pasted_key() {
  local key

  key=$(gum input --prompt "Public key> " --placeholder "ssh-ed25519 AAAA... user@host") || exit 1
  if [[ -z $key ]]; then
    echo -e "\e[31mNo SSH key given.\e[0m" >&2
    exit 1
  fi

  authorize_key "$key" || exit 1
}

# Only called after a key is authorized. Disabling password authentication
# before then could lock the owner out of the machine.
disable_password_auth() {
  local config=/etc/ssh/sshd_config.d/10-blob-hardening.conf
  local effective_config

  if [[ ! -s $AUTHORIZED_KEYS ]]; then
    echo -e "\e[31mCannot disable SSH password authentication without an authorized key.\e[0m" >&2
    return 1
  fi

  echo "Disabling SSH password authentication, now that a key is authorized..."
  sudo install -Dm644 /dev/stdin "$config" <<'CONF'
# Written by blob-setup-security-sshd once an SSH key was authorized.
# Delete this file and reload sshd to allow password logins again.
PasswordAuthentication no
KbdInteractiveAuthentication no
CONF

  # Validate before reloading: a config sshd rejects would otherwise take the
  # service down on its next restart, potentially stranding a remote owner.
  if ! sudo sshd -t; then
    echo -e "\e[31msshd rejected the hardening config; removing it and leaving passwords on.\e[0m" >&2
    sudo rm -f "$config"
    return 1
  fi

  # Syntax alone is insufficient because sshd uses the first value it reads for
  # these settings. An earlier administrator rule could leave passwords enabled.
  # Match keywords case-insensitively: OpenSSH 9.x dumps them lowercase, 10.x
  # in CamelCase.
  if ! effective_config=$(sudo sshd -T) ||
    ! grep -qixF "passwordauthentication no" <<<"$effective_config" ||
    ! grep -qixF "kbdinteractiveauthentication no" <<<"$effective_config"; then
    echo -e "\e[31msshd did not apply the password-authentication restrictions; removing the ineffective config.\e[0m" >&2
    sudo rm -f "$config"
    return 1
  fi

  # Reload rather than restart so an administrator already connected keeps
  # their session.
  sudo systemctl reload sshd.service
}

echo -e "\e[32mSetting up SSH server access with key-based authentication.\n\e[0m"

setup_sshd
open_firewall

echo
if [[ -n $KEY ]]; then
  authorize_key "$KEY" || exit 1
else
  case $(gum choose "Grab key from GitHub" "Paste key manually" --header "How would you like to add your SSH key?") in
  "Grab key from GitHub") authorize_keys_from_github ;;
  "Paste key manually") authorize_pasted_key ;;
  *) exit 1 ;;
  esac
fi

disable_password_auth

echo -e "\e[32m\nPerfect! The SSH server is running and your key is authorized.\e[0m"
echo "Password logins are off; this machine now accepts authorized keys only."
echo "You can now connect with: ssh $USER@$(hostname)"