#!/bin/bash
# First-time VPS bootstrap from your workstation:
#   password login → deploy user + SSH key → test key → disable password auth → apt upgrade
#   → optional ~/.ssh/config alias on this machine
#
# Run locally (not on the VPS):
#   wget -O bootstrap-vps-ssh.sh \
#     https://gitea.dialer.work/swissdatabase/rubix-deploy/raw/branch/main/bootstrap-vps-ssh.sh
#   chmod +x bootstrap-vps-ssh.sh
#   ./bootstrap-vps-ssh.sh
#
# Needs: ssh, optional sshpass (apt install sshpass)

set -euo pipefail

SSH_CONFIG="${HOME}/.ssh/config"
SSH_DIR="${HOME}/.ssh"
KNOWN_HOSTS_FILE="${HOME}/.ssh/known_hosts"
MARKER_BEGIN="# --- RUBIX-VPS-BEGIN"
MARKER_END="# --- RUBIX-VPS-END"

prompt() {
  local var_name="$1"
  local text="$2"
  local default="${3:-}"
  local value=""
  if [[ -n "${default}" ]]; then
    read -r -p "${text} [${default}]: " value
    value="${value:-${default}}"
  else
    read -r -p "${text}: " value
  fi
  printf -v "${var_name}" '%s' "${value}"
}

prompt_secret() {
  local var_name="$1"
  local text="$2"
  local value=""
  read -r -s -p "${text}: " value
  echo ""
  printf -v "${var_name}" '%s' "${value}"
}

prompt_yes_no() {
  local var_name="$1"
  local text="$2"
  local default="${3:-y}"
  local hint="Y/n"
  [[ "${default}" == "y" ]] && hint="Y/n" || hint="y/N"
  local answer=""
  read -r -p "${text} [${hint}]: " answer
  answer="${answer:-${default}}"
  case "${answer}" in
    y|Y|yes|YES) printf -v "${var_name}" '%s' "yes" ;;
    *) printf -v "${var_name}" '%s' "no" ;;
  esac
}

have_sshpass() {
  command -v sshpass >/dev/null 2>&1
}

# Install sshpass on the machine running this script (your PC), not on the VPS.
ensure_local_sshpass() {
  if have_sshpass; then
    return 0
  fi

  echo "[rubix-vps] sshpass not found — installing locally (one sudo prompt) ..."

  local run_as=(bash -c)
  if [[ "$(id -u)" -ne 0 ]]; then
    if ! command -v sudo >/dev/null 2>&1; then
      echo "[rubix-vps] WARN: need sudo to install sshpass (run: sudo apt install sshpass)." >&2
      return 1
    fi
    run_as=(sudo bash -c)
  fi

  if command -v apt-get >/dev/null 2>&1; then
    "${run_as[@]}" 'apt-get update -qq && apt-get install -y sshpass'
  elif command -v apt >/dev/null 2>&1; then
    "${run_as[@]}" 'apt update -qq && apt install -y sshpass'
  elif command -v dnf >/dev/null 2>&1; then
    "${run_as[@]}" 'dnf install -y sshpass'
  elif command -v yum >/dev/null 2>&1; then
    "${run_as[@]}" 'yum install -y sshpass'
  elif command -v pacman >/dev/null 2>&1; then
    "${run_as[@]}" 'pacman -Sy --noconfirm sshpass'
  elif command -v zypper >/dev/null 2>&1; then
    "${run_as[@]}" 'zypper install -y sshpass'
  else
    echo "[rubix-vps] WARN: no supported package manager for auto-install of sshpass." >&2
    return 1
  fi

  if have_sshpass; then
    echo "[rubix-vps] sshpass installed."
    return 0
  fi
  echo "[rubix-vps] WARN: sshpass install failed." >&2
  return 1
}

# Stop ssh-agent from offering keys (avoids "Too many authentication failures" on password login).
password_ssh_common_opts() {
  printf '%s\n' \
    StrictHostKeyChecking=accept-new \
    IdentitiesOnly=yes \
    IdentityAgent=none \
    PubkeyAuthentication=no \
    PreferredAuthentications=password,keyboard-interactive \
    NumberOfPasswordPrompts=3
}

ssh_password() {
  local target="$1"
  shift
  local -a base=(ssh -p "${VPS_PORT}")
  local opt
  while IFS= read -r opt; do
    base+=(-o "${opt}")
  done < <(password_ssh_common_opts)
  if [[ -n "${ROOT_PASSWORD:-}" ]] && have_sshpass; then
    sshpass -p "${ROOT_PASSWORD}" "${base[@]}" "${target}" "$@"
  else
    "${base[@]}" "${target}" "$@"
  fi
}

scp_password() {
  local src="$1"
  local dest="$2"
  local -a base=(scp -P "${VPS_PORT}")
  local opt
  while IFS= read -r opt; do
    base+=(-o "${opt}")
  done < <(password_ssh_common_opts)
  if [[ -n "${ROOT_PASSWORD:-}" ]] && have_sshpass; then
    sshpass -p "${ROOT_PASSWORD}" "${base[@]}" "${src}" "${dest}"
  else
    "${base[@]}" "${src}" "${dest}"
  fi
}

ssh_key() {
  local user="$1"
  shift
  ssh -o BatchMode=yes \
    -o StrictHostKeyChecking=accept-new \
    -o PasswordAuthentication=no \
    -o IdentitiesOnly=yes \
    -i "${PRIVATE_KEY_FILE}" \
    -p "${VPS_PORT}" \
    "${user}@${VPS_HOST}" \
    "$@"
}

test_key_login() {
  local user="$1"
  ssh_key "${user}" "echo ok" >/dev/null 2>&1
}

cleanup_known_host_entries() {
  local host="$1"
  local port="$2"
  [[ -f "${KNOWN_HOSTS_FILE}" ]] || return 0

  # Remove stale host keys so reinstalled/reprovisioned VPS can be reached without manual known_hosts edits.
  ssh-keygen -f "${KNOWN_HOSTS_FILE}" -R "${host}" >/dev/null 2>&1 || true
  ssh-keygen -f "${KNOWN_HOSTS_FILE}" -R "[${host}]:${port}" >/dev/null 2>&1 || true
}

remote_env() {
  printf 'DEPLOY_USER=%q KEEP_ROOT=%q INSTALL_ROOT_KEY=%q' \
    "${DEPLOY_USER}" "${KEEP_ROOT}" "${INSTALL_ROOT_KEY}"
}

update_local_ssh_config() {
  local alias="$1"
  local host="$2"
  local port="$3"
  local user="$4"
  local key_file="$5"

  mkdir -p "${SSH_DIR}"
  chmod 700 "${SSH_DIR}"
  touch "${SSH_CONFIG}"
  chmod 600 "${SSH_CONFIG}"

  local block
  block="$(cat <<EOF
${MARKER_BEGIN} ${alias} ---
Host ${alias}
    HostName ${host}
    User ${user}
    Port ${port}
    IdentityFile ${key_file}
    IdentitiesOnly yes
${MARKER_END} ${alias} ---
EOF
)"

  local tmp
  tmp="$(mktemp)"
  if [[ -f "${SSH_CONFIG}" ]]; then
    awk -v begin="${MARKER_BEGIN} ${alias} ---" -v end="${MARKER_END} ${alias} ---" '
      $0 == begin { skip=1; next }
      $0 == end { skip=0; next }
      !skip { print }
    ' "${SSH_CONFIG}" > "${tmp}" || true
  fi
  {
    cat "${tmp}"
    echo ""
    printf '%s\n' "${block}"
  } > "${SSH_CONFIG}.new"
  mv "${SSH_CONFIG}.new" "${SSH_CONFIG}"
  rm -f "${tmp}"
  echo "[rubix-vps] wrote Host ${alias} → ${SSH_CONFIG}"
}

print_key_summary() {
  local alias="$1"
  echo ""
  echo "[rubix-vps] SSH key for Host alias \"${alias}\":"
  echo "  Private key : ${PRIVATE_KEY_FILE}"
  echo "  Public key  : ${PUBLIC_KEY_FILE}"
  echo "  ssh config  : Host ${alias} → IdentityFile ${PRIVATE_KEY_FILE}"
  echo ""
  echo "  authorized_keys line:"
  sed 's/^/    /' "${PUBLIC_KEY_FILE}"
  echo ""
}

pick_ssh_key() {
  local alias="$1"
  local choice=""
  local key_path="${SSH_DIR}/${alias}.ed25519"
  echo ""
  echo "SSH key for ~/.ssh/config Host \"${alias}\":"
  echo "  Naming: ${key_path} (+ ${key_path}.pub)"
  echo ""
  echo "  1) Use existing public key"
  echo "  2) Generate new ed25519 key pair at ${key_path}"
  read -r -p "Choice [2]: " choice
  choice="${choice:-2}"

  case "${choice}" in
    1)
      local default_pub="${SSH_DIR}/id_ed25519.pub"
      prompt PUBLIC_KEY_FILE "Path to public key (.pub)" "${default_pub}"
      if [[ ! -f "${PUBLIC_KEY_FILE}" ]]; then
        echo "Missing: ${PUBLIC_KEY_FILE}" >&2
        exit 1
      fi
      local priv="${PUBLIC_KEY_FILE%.pub}"
      if [[ -f "${priv}" ]]; then
        PRIVATE_KEY_FILE="${priv}"
      else
        prompt PRIVATE_KEY_FILE "Path to matching private key" "${SSH_DIR}/id_ed25519"
      fi
      ;;
    *)
      if [[ -f "${key_path}" ]]; then
        prompt_yes_no OVERWRITE "Key ${key_path} exists. Overwrite?" "n"
        if [[ "${OVERWRITE}" != "yes" ]]; then
          echo "Aborted." >&2
          exit 1
        fi
        rm -f "${key_path}" "${key_path}.pub"
      fi
      local comment="rubix-vps-${alias}-$(date +%Y%m%d)"
      echo ""
      echo "[rubix-vps] Generating ${key_path} (comment: ${comment}) ..."
      ssh-keygen -t ed25519 -f "${key_path}" -C "${comment}" -N ""
      PRIVATE_KEY_FILE="${key_path}"
      PUBLIC_KEY_FILE="${key_path}.pub"
      ;;
  esac

  if [[ ! -f "${PRIVATE_KEY_FILE}" ]]; then
    echo "Private key not found: ${PRIVATE_KEY_FILE}" >&2
    exit 1
  fi
  chmod 600 "${PRIVATE_KEY_FILE}" 2>/dev/null || true
  print_key_summary "${alias}"
}

remote_install_users_and_keys() {
  ssh_password "${INITIAL_USER}@${VPS_HOST}" \
    "$(remote_env)" bash -s <<'REMOTE'
set -euo pipefail

if [[ "$(id -u)" -ne 0 ]]; then
  echo "Remote: must run as root." >&2
  exit 1
fi

PUBKEY="$(tr -d '\r\n' < /tmp/rubix-bootstrap.pub)"

if ! id "${DEPLOY_USER}" >/dev/null 2>&1; then
  useradd -m -s /bin/bash "${DEPLOY_USER}"
  echo "Created user ${DEPLOY_USER}"
else
  echo "User ${DEPLOY_USER} already exists"
fi

if getent group sudo >/dev/null 2>&1; then
  usermod -aG sudo "${DEPLOY_USER}"
elif getent group wheel >/dev/null 2>&1; then
  usermod -aG wheel "${DEPLOY_USER}"
fi

install_key_for_user() {
  local u="$1"
  local home_dir
  home_dir="$(getent passwd "${u}" | cut -d: -f6)"
  local ssh_dir="${home_dir}/.ssh"
  mkdir -p "${ssh_dir}"
  chmod 700 "${ssh_dir}"
  touch "${ssh_dir}/authorized_keys"
  if ! grep -qxF "${PUBKEY}" "${ssh_dir}/authorized_keys" 2>/dev/null; then
    echo "${PUBKEY}" >> "${ssh_dir}/authorized_keys"
  fi
  chown -R "${u}:${u}" "${ssh_dir}"
  chmod 600 "${ssh_dir}/authorized_keys"
}

install_key_for_user "${DEPLOY_USER}"

if [[ "${INSTALL_ROOT_KEY}" == "yes" ]]; then
  install_key_for_user root
fi

echo "Keys installed for ${DEPLOY_USER} (and root if requested)."
REMOTE
}

remote_harden_and_upgrade() {
  ssh_password "${INITIAL_USER}@${VPS_HOST}" \
    "$(remote_env)" bash -s <<'REMOTE'
set -euo pipefail

SSHD="/etc/ssh/sshd_config"
cp -a "${SSHD}" "${SSHD}.rubix-backup-$(date +%Y%m%d%H%M%S)"

set_sshd_option() {
  local key="$1"
  local value="$2"
  if grep -qE "^[#[:space:]]*${key}[[:space:]]" "${SSHD}"; then
    sed -i "s/^[#[:space:]]*${key}[[:space:]].*/${key} ${value}/" "${SSHD}"
  else
    echo "${key} ${value}" >> "${SSHD}"
  fi
}

set_sshd_option PasswordAuthentication no
set_sshd_option KbdInteractiveAuthentication no
set_sshd_option ChallengeResponseAuthentication no
set_sshd_option PubkeyAuthentication yes
set_sshd_option UsePAM no

if [[ "${KEEP_ROOT}" == "yes" ]]; then
  set_sshd_option PermitRootLogin prohibit-password
else
  set_sshd_option PermitRootLogin no
fi

if sshd -t 2>/dev/null || /usr/sbin/sshd -t 2>/dev/null; then
  :
else
  latest="$(ls -t ${SSHD}.rubix-backup-* 2>/dev/null | head -1)"
  [[ -n "${latest}" ]] && cp -a "${latest}" "${SSHD}"
  echo "sshd config test failed — restored backup" >&2
  exit 1
fi

if systemctl list-units --type=service --all 2>/dev/null | grep -q 'sshd\.service'; then
  systemctl reload sshd
elif systemctl list-units --type=service --all 2>/dev/null | grep -q 'ssh\.service'; then
  systemctl reload ssh
else
  service ssh reload 2>/dev/null || service sshd reload
fi

if command -v apt-get >/dev/null 2>&1; then
  export DEBIAN_FRONTEND=noninteractive
  apt-get update -y
  apt-get upgrade -y
  apt-get install -y locales
  if [[ -f /etc/locale.gen ]]; then
    sed -i '/en_US.UTF-8 UTF-8/s/^# *//' /etc/locale.gen
    grep -q '^en_US.UTF-8 UTF-8' /etc/locale.gen || echo 'en_US.UTF-8 UTF-8' >> /etc/locale.gen
    locale-gen en_US.UTF-8 2>/dev/null || locale-gen
  fi
  update-locale LANG=en_US.UTF-8 LC_CTYPE=en_US.UTF-8 LC_MESSAGES=en_US.UTF-8 2>/dev/null || true
  mkdir -p /var/lib/cloud/instance
  touch /var/lib/cloud/instance/locale-check.skip 2>/dev/null || true
elif command -v dnf >/dev/null 2>&1; then
  dnf -y upgrade
elif command -v yum >/dev/null 2>&1; then
  yum -y update
fi

rm -f /tmp/rubix-bootstrap.pub
echo "SSH hardened and packages upgraded."
REMOTE
}

main() {
  echo "=== Rubix VPS first-time SSH bootstrap (run on your PC) ==="
  echo ""

  ensure_local_sshpass || true

  prompt VPS_HOST "VPS hostname or IP"
  prompt VPS_PORT "SSH port" "22"
  prompt INITIAL_USER "First-login user (usually root)" "root"
  prompt DEPLOY_USER "Deploy/sudo user to create" "deployer"
  if [[ "${DEPLOY_USER}" == "root" ]]; then
    echo "WARN: Using root as deploy user — prefer a sudo user (e.g. deployer) for daily SSH." >&2
  fi
  prompt SSH_ALIAS "Local ~/.ssh/config Host alias (e.g. rubix-server-5)" "${DEPLOY_USER}-${VPS_HOST}"

  prompt_yes_no KEEP_ROOT "Keep root SSH login (key-only, no password)" "yes"
  prompt_yes_no INSTALL_ROOT_KEY "Install same public key for root" "${KEEP_ROOT}"

  if have_sshpass; then
    prompt_secret ROOT_PASSWORD "Initial ${INITIAL_USER} password (entered once via sshpass)"
  else
    echo ""
    echo "[rubix-vps] sshpass unavailable — ssh/scp will prompt for the VPS password each time."
    echo ""
    prompt_secret ROOT_PASSWORD "Initial ${INITIAL_USER} password (empty = type when ssh asks)"
  fi

  mkdir -p "${SSH_DIR}"
  chmod 700 "${SSH_DIR}"
  pick_ssh_key "${SSH_ALIAS}"

  if [[ ! -f "${PUBLIC_KEY_FILE}" ]]; then
    PUBLIC_KEY_FILE="${PRIVATE_KEY_FILE}.pub"
  fi
  if [[ ! -f "${PUBLIC_KEY_FILE}" ]]; then
    echo "Public key not found: ${PUBLIC_KEY_FILE}" >&2
    exit 1
  fi

  echo "[rubix-vps] Cleaning stale known_hosts entries for ${VPS_HOST}:${VPS_PORT} ..."
  cleanup_known_host_entries "${VPS_HOST}" "${VPS_PORT}"

  echo ""
  echo "[rubix-vps] Testing password login ..."
  ssh_password "${INITIAL_USER}@${VPS_HOST}" "echo connected && uname -a"

  echo "[rubix-vps] Uploading public key ..."
  scp_password "${PUBLIC_KEY_FILE}" "${INITIAL_USER}@${VPS_HOST}:/tmp/rubix-bootstrap.pub"

  echo "[rubix-vps] Creating ${DEPLOY_USER} and installing authorized_keys ..."
  remote_install_users_and_keys

  echo "[rubix-vps] Testing key login (before disabling passwords) ..."
  if ! test_key_login "${DEPLOY_USER}"; then
    echo "ERROR: Key login as ${DEPLOY_USER} failed. Password auth still enabled — fix keys and retry." >&2
    exit 1
  fi
  echo "[rubix-vps] Key login OK for ${DEPLOY_USER}"

  if [[ "${INSTALL_ROOT_KEY}" == "yes" && "${DEPLOY_USER}" != "root" ]]; then
    if test_key_login root; then
      echo "[rubix-vps] Key login OK for root"
    else
      echo "[rubix-vps] WARN: root key login failed (deploy user works)" >&2
    fi
  fi

  echo "[rubix-vps] Disabling password auth, reloading sshd, running apt upgrade ..."
  remote_harden_and_upgrade

  echo "[rubix-vps] Testing key login after hardening ..."
  if ! test_key_login "${DEPLOY_USER}"; then
    echo "ERROR: Key login failed after hardening. Use provider console; restore /etc/ssh/sshd_config.rubix-backup-*" >&2
    exit 1
  fi

  prompt_yes_no WRITE_CONFIG "Add Host ${SSH_ALIAS} to ${SSH_CONFIG}" "yes"
  if [[ "${WRITE_CONFIG}" == "yes" ]]; then
    update_local_ssh_config "${SSH_ALIAS}" "${VPS_HOST}" "${VPS_PORT}" "${DEPLOY_USER}" "${PRIVATE_KEY_FILE}"
    echo ""
    echo "Next login:"
    echo "  ssh ${SSH_ALIAS}"
  else
    echo ""
    echo "Manual login:"
    echo "  ssh -i ${PRIVATE_KEY_FILE} -p ${VPS_PORT} ${DEPLOY_USER}@${VPS_HOST}"
  fi

  echo ""
  echo "Done."
}

main "$@"