#!/bin/bash

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

set -e

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

# Checked while parsing, before anything is installed or opened: an empty or
# option-shaped username otherwise falls through to the interactive menu having
# already changed the machine, and `--gh-keys --help` would take --help as the
# username and set the server up on its way to failing.
require_github_user() {
  if [[ -z $1 || $1 == -* ]]; then
    echo "omarchy-setup-security-sshd: --gh-keys needs a GitHub username." >&2
    exit 2
  fi
}

while (( $# > 0 )); do
  case "$1" in
  --key=*) KEY="${1#--key=}" ;;
  --gh-keys=*)
    GITHUB_USER="${1#--gh-keys=}"
    require_github_user "$GITHUB_USER"
    ;;
  --gh-keys)
    shift
    GITHUB_USER="${1:-}"
    require_github_user "$GITHUB_USER"
    ;;
  -h | --help)
    echo "Usage: omarchy-setup-security-sshd [--key=<public-key>] [--gh-keys <github-username>]"
    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)."
    echo
    echo "Passing --key or --gh-keys skips the prompts, so the command can run"
    echo "unattended from a script or a fresh machine's first login."
    exit 0
    ;;
  *)
    echo "omarchy-setup-security-sshd: unknown option '$1'. Try --help." >&2
    exit 2
    ;;
  esac
  shift
done

if [[ -n $KEY && -n $GITHUB_USER ]]; then
  echo "omarchy-setup-security-sshd: pass either --key or --gh-keys, not both." >&2
  exit 2
fi

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

open_firewall() {
  if omarchy-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 "omarchy-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="$1" keys added=0

  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
}

prompt_for_github_user() {
  local username

  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

  authorize_keys_from_github "$username"
}

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
}

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
elif [[ -n $GITHUB_USER ]]; then
  authorize_keys_from_github "$GITHUB_USER"
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") prompt_for_github_user ;;
  "Paste key manually") authorize_pasted_key ;;
  *) exit 1 ;;
  esac
fi

echo -e "\e[32m\nPerfect! The SSH server is running and your key is authorized.\e[0m"
echo "You can now connect with: ssh $USER@$(hostname)"