#!/bin/bash
# =============================================================================
# machine-setup — ssh keys and ssh hardening
# =============================================================================
#
# Definitions only, like the other lib/ files.
#
# ── Why the original's hardening did not work, and could not be seen not to ──
#
# It sed'd /etc/ssh/sshd_config directly. Two things make that wrong on a modern
# Ubuntu, and both fail silently:
#
#   Ubuntu's sshd_config has `Include /etc/ssh/sshd_config.d/*.conf` on line 12,
#   and sshd takes the FIRST value it obtains for a keyword — not the last. Cloud
#   images ship 50-cloud-init.conf containing `PasswordAuthentication yes`, which
#   is read before anything further down the main file. So the sed edits a line
#   sshd never reaches, the script reports "SSH hardened", and password login is
#   still on.
#
#   It also sed'd ChallengeResponseAuthentication, which OpenSSH renamed to
#   KbdInteractiveAuthentication in 8.7. On 24.04 the old name appears nowhere in
#   the file, so that substitution matched nothing at all.
#
# So the settings go in a drop-in named to sort FIRST — 01- beats 50-cloud-init —
# which is the only placement that actually wins under first-value-wins.
#
# ── And the reason it is dangerous ──
#
# Step 8 of the original could warn-and-skip (no ssh-keys.zip, or an unrecognised
# menu choice, since its case had no default arm) and still mark itself done.
# Step 9 then disabled password authentication and root login regardless. No key,
# no password, no root: locked out at the next disconnect, on a machine that may
# be in a datacentre. Nothing here disables password authentication without first
# confirming a usable key is in place.

[[ -n "${MACHINE_SETUP_SSH_LOADED:-}" ]] && return 0
MACHINE_SETUP_SSH_LOADED=1

SSHD_DROPIN=/etc/ssh/sshd_config.d/01-machine-setup.conf

# -----------------------------------------------------------------------------
# Keys
# -----------------------------------------------------------------------------

user_ssh_dir() { echo "${USER_HOME}/.ssh"; }
user_authorized_keys() { echo "${USER_HOME}/.ssh/authorized_keys"; }

# How many usable keys the account can log in with.
#
# Counted by asking ssh-keygen to parse the file rather than by counting lines:
# comments, blanks and a half-pasted key all look like lines, and "there is a
# file" is not the same fact as "there is a key that works".
authorized_key_count() {
  local file
  file="$(user_authorized_keys)"
  [[ -r "$file" ]] || return 0
  ssh-keygen -l -f "$file" 2>/dev/null | grep -c . || true
}

has_authorized_key() { (($(authorized_key_count) > 0)); }

# Everything about ~/.ssh that has to be true for sshd to use it at all. sshd
# ignores an authorized_keys file that is group- or world-writable, and does so
# silently from the client's point of view — the login just fails.
fix_ssh_permissions() {
  local dir
  dir="$(user_ssh_dir)"
  [[ -d "$dir" ]] || install -d -m 0700 -o "$USERNAME" -g "$(user_group)" "$dir"
  chmod 700 "$dir"
  [[ -f "$dir/authorized_keys" ]] && chmod 600 "$dir/authorized_keys"
  find "$dir" -maxdepth 1 -type f -name 'id_*' ! -name '*.pub' -exec chmod 600 {} +
  chown -R "${USERNAME}:$(user_group)" "$dir"
  # Returns 0 whatever happens. This is an optional improvement, and a
  # function that ends on a failing command is fatal under `set -e` when it
  # is called as a plain command — which would abort the remaining sections
  # over something the run could simply report. The caller checks the outcome.
  return 0
}

# Add a public key, once. Appending blindly is how authorized_keys ends up with
# the same key four times after four runs.
add_authorized_key() {
  local key="$1" file
  file="$(user_authorized_keys)"

  # Validated before it is stored. A truncated paste or a private key pasted by
  # mistake would otherwise sit there looking like a key and never work.
  if ! ssh-keygen -l -f /dev/stdin <<<"$key" >/dev/null 2>&1; then
    warn "that does not parse as an ssh public key — nothing added"
    return 1
  fi

  install -d -m 0700 -o "$USERNAME" -g "$(user_group)" "$(user_ssh_dir)"
  touch "$file"

  # Compare on the key body, not the whole line: the trailing comment differs
  # between machines and is not part of the identity.
  local body
  body="$(awk '{print $2}' <<<"$key")"
  if [[ -n "$body" ]] && grep -qF "$body" "$file" 2>/dev/null; then
    info "     that key is already authorised"
    return 0
  fi

  printf '%s\n' "$key" >>"$file"
  fix_ssh_permissions
}

# Generate a keypair for the account and authorise it.
generate_user_key() {
  local comment="$1" key
  key="$(user_ssh_dir)/id_ed25519"

  install -d -m 0700 -o "$USERNAME" -g "$(user_group)" "$(user_ssh_dir)"
  sudo -u "$USERNAME" ssh-keygen -t ed25519 -C "$comment" -f "$key" -N "" >/dev/null
  add_authorized_key "$(cat "${key}.pub")"
}

# -----------------------------------------------------------------------------
# Hardening
# -----------------------------------------------------------------------------

# What sshd actually resolves a setting to, across the main file and every
# drop-in. The only honest way to report the current state: reading the config
# files tells you what is written, not what wins.
sshd_effective() { sshd -T 2>/dev/null | awk -v k="${1,,}" 'tolower($1) == k { print $2; exit }'; }

# Write the drop-in, verify it, and only then reload.
#
# Returns non-zero without touching the running daemon if the result would not
# parse — the alternative is a config that sshd refuses, at which point it will
# not come back after a restart and the machine has no ssh at all.
harden_sshd() {
  local backup=""

  [[ -f "$SSHD_DROPIN" ]] && backup="$(mktemp)" && cp "$SSHD_DROPIN" "$backup"

  install -d -m 0755 /etc/ssh/sshd_config.d
  cat >"$SSHD_DROPIN" <<'EOF'
# Written by machine-setup.
#
# Named 01- deliberately: sshd uses the FIRST value it obtains for a keyword, and
# Ubuntu includes this directory from the top of sshd_config. A file sorting
# after 50-cloud-init.conf would be read too late to override it.
PasswordAuthentication no
KbdInteractiveAuthentication no
PermitRootLogin no
PubkeyAuthentication yes
EOF
  chmod 644 "$SSHD_DROPIN"

  if ! sshd -t 2>/dev/null; then
    warn "sshd rejected the new configuration — reverting, nothing changed"
    if [[ -n "$backup" ]]; then cp "$backup" "$SSHD_DROPIN"; else rm -f "$SSHD_DROPIN"; fi
    [[ -n "$backup" ]] && rm -f "$backup"
    return 1
  fi
  [[ -n "$backup" ]] && rm -f "$backup"

  # Reload rather than restart: existing sessions keep their sshd, so the
  # connection this is being run over is not the thing being experimented on.
  systemctl reload ssh 2>/dev/null || systemctl reload sshd 2>/dev/null || systemctl restart ssh
}