#!/usr/bin/env bash
# user_manage.sh
# Interactive script to create/delete sudo users and manage SSH keys.
# - Generates ed25519 keypair per user and saves them to /root/created_user_keys/<username>-<timestamp>/
# - Installs public key into user's authorized_keys and sets permissions
# - Adds user to sudo/wheel group
# - Allows deleting user and cleaning keys/home
#
# Usage: sudo ./user_manage.sh

set -euo pipefail
IFS=$'\n\t'

BASE_KEY_DIR="/root/created_user_keys"

# UI helpers — send user messages to stderr so functions can return values on stdout
function step() { printf "\n\033[1;34m==> %s\033[0m\n" "$1" >&2; }
function ok()   { printf "   \033[1;32m[OK]\033[0m %s\n" "$1" >&2; }
function warn() { printf "   \033[1;33m[WARN]\033[0m %s\n" "$1" >&2; }
function fail() { printf "   \033[1;31m[FAIL]\033[0m %s\n" "$1" >&2; }

if [ "$(id -u)" -ne 0 ]; then
  fail "This script must be run as root. Use: sudo $0"
  exit 2
fi

mkdir -p "$BASE_KEY_DIR"
chmod 700 "$BASE_KEY_DIR"

# Detect distro and sudo group
SUDO_GROUP="sudo"
if getent group sudo >/dev/null 2>&1; then
  SUDO_GROUP="sudo"
elif getent group wheel >/dev/null 2>&1; then
  SUDO_GROUP="wheel"
fi

# Helper to create user (common tasks)
create_user_common() {
  local USERNAME="$1"
  if id "$USERNAME" &>/dev/null; then
    fail "User $USERNAME already exists."
    return 1
  fi

  if command -v adduser >/dev/null 2>&1 && { [ -f /etc/debian_version ] || command -v apt-get >/dev/null 2>&1 || [ -f /etc/lsb-release ] ; }; then
    step "Creating user '$USERNAME' (adduser)"
    adduser --disabled-password --gecos "" "$USERNAME"
  else
    step "Creating user '$USERNAME' (useradd fallback)"
    useradd -m -s /bin/bash "$USERNAME"
  fi

  ok "User $USERNAME created."

  # add to sudo/wheel
  usermod -aG "$SUDO_GROUP" "$USERNAME"
  ok "Added $USERNAME to $SUDO_GROUP group."
}

# generate_keypair: prints only the key directory on stdout (so callers can capture it)
generate_keypair() {
  local USERNAME="$1"
  local DATESTAMP
  DATESTAMP="$(date '+%Y%m%d_%H%M%S')"
  local KEY_DIR="$BASE_KEY_DIR/$USERNAME-$DATESTAMP"
  mkdir -p "$KEY_DIR"
  chmod 700 "$KEY_DIR"

  # generate ed25519 keypair WITHOUT passphrase; to add passphrase remove -N ""
  ssh-keygen -t ed25519 -f "$KEY_DIR/id_ed25519" -N "" -C "$USERNAME@$(hostname -f)" >/dev/null 2>&1

  chmod 600 "$KEY_DIR/id_ed25519"
  chmod 644 "$KEY_DIR/id_ed25519.pub"

  # status message to stderr
  ok "Generated ed25519 keypair for $USERNAME in $KEY_DIR"

  # echo only the path on stdout so the caller can capture it safely
  echo "$KEY_DIR"
}

# Install public key into user's authorized_keys
install_pubkey_for_user() {
  local USERNAME="$1"
  local PUBKEY_LINE="$2"

  local USER_HOME
  USER_HOME=$(eval echo "~$USERNAME")
  local SSH_DIR="$USER_HOME/.ssh"
  local AUTH_KEYS="$SSH_DIR/authorized_keys"

  mkdir -p "$SSH_DIR"
  chmod 700 "$SSH_DIR"
  chown "$USERNAME:$USERNAME" "$SSH_DIR"

  echo "$PUBKEY_LINE" > "$AUTH_KEYS"
  chmod 600 "$AUTH_KEYS"
  chown "$USERNAME:$USERNAME" "$AUTH_KEYS"

  ok "Installed public key to $AUTH_KEYS (owner: $USERNAME)"
}

# Simple sudo verification
verify_sudo() {
  local USERNAME="$1"
  if su - "$USERNAME" -c 'sudo -l' >/tmp/sudo_check 2>&1; then
    echo "Output of 'sudo -l' for $USERNAME:" >&2
    sed -n '1,200p' /tmp/sudo_check >&2 || true
    ok "Sudo appears configured for $USERNAME"
  else
    warn "Sudo check produced output (inspect below):" >&2
    sed -n '1,200p' /tmp/sudo_check >&2 || true
    warn "If problems, ensure $USERNAME is in $SUDO_GROUP or update /etc/sudoers with visudo." >&2
  fi
  rm -f /tmp/sudo_check
}

# Delete user and cleanup
delete_user() {
  local USERNAME="$1"
  if ! id "$USERNAME" &>/dev/null; then
    fail "User $USERNAME does not exist."
    return 1
  fi

  read -rp "Are you sure you want to DELETE user '$USERNAME' and remove their home dir and saved keys? This is irreversible. [type 'yes' to confirm]: " CONF
  if [ "$CONF" != "yes" ]; then
    warn "Aborted deletion."
    return 2
  fi

  step "Killing processes for $USERNAME"
  pkill -u "$USERNAME" || true
  ok "Killed processes (if any)."

  step "Removing user and home directory"
  if userdel -r "$USERNAME" >/dev/null 2>&1; then
    ok "userdel -r succeeded for $USERNAME"
  else
    warn "userdel -r failed or partially failed; attempting manual cleanup of home directory"
    rm -rf "$(eval echo "~$USERNAME")" || true
    ok "Manual home cleanup attempted"
  fi

  step "Cleaning created key folders for $USERNAME in $BASE_KEY_DIR"
  find "$BASE_KEY_DIR" -maxdepth 1 -type d -name "$USERNAME*" -exec rm -rf {} \; || true
  ok "Removed saved key directories matching $USERNAME* in $BASE_KEY_DIR"

  ok "User $USERNAME deletion finished."
  return 0
}

# MAIN MENU
while true; do
  cat <<'MENU' >&2

=== USER MANAGEMENT MENU ===
1) Create new sudo user AND auto-generate SSH keypair (private key saved for copying)
2) Create new sudo user from an EXISTING public key (paste)
3) Delete a user and remove their home + saved keys
4) Show saved key directories (/root/created_user_keys)
5) Exit
MENU

  read -rp "Choose an option [1-5]: " CHOICE

  case "$CHOICE" in
    1)
      read -rp $'\nEnter the username to create (example: alice): ' NEWUSER
      NEWUSER="${NEWUSER:-}"
      if [ -z "$NEWUSER" ]; then fail "No username entered."; continue; fi

      read -rp "Create user '$NEWUSER' and generate keypair? [y/N]: " CONF
      CONF=${CONF,,}
      if [[ "$CONF" != "y" && "$CONF" != "yes" ]]; then warn "Cancelled."; continue; fi

      step "Creating user and setting up keys for $NEWUSER"
      if ! create_user_common "$NEWUSER"; then warn "User creation failed. Aborting."; continue; fi

      # generate keys (capture only the path)
      KEY_DIR="$(generate_keypair "$NEWUSER")"
      # read public key
      PUBKEY_LINE="$(cat "$KEY_DIR/id_ed25519.pub" | tr -d '\r\n')"
      install_pubkey_for_user "$NEWUSER" "$PUBKEY_LINE"

      read -rp "Lock local password for $NEWUSER (disallow password login for this account)? [y/N]: " LOCKCHOICE
      LOCKCHOICE=${LOCKCHOICE,,}
      if [[ "$LOCKCHOICE" == "y" || "$LOCKCHOICE" == "yes" ]]; then
        passwd -l "$NEWUSER" >/dev/null 2>&1 || true
        ok "Password locked for $NEWUSER (key-only)."
      else
        ok "Password not locked; you may set one with: sudo passwd $NEWUSER"
      fi

      step "Verifying sudo for $NEWUSER"
      verify_sudo "$NEWUSER"

      step "Result & next steps"
      echo " - Keys saved in: $KEY_DIR" >&2
      echo " - Files to copy to the user (deliver PRIVATE key securely):" >&2
      echo "     $KEY_DIR/id_ed25519        (private - deliver securely)" >&2
      echo "     $KEY_DIR/id_ed25519.pub    (public - for records)" >&2
      echo " - Installed public key location on server:" >&2
      echo "     /home/$NEWUSER/.ssh/authorized_keys" >&2
      echo "" >&2
      echo "Example: to copy private key to your workstation (run locally on workstation):" >&2
      echo "  scp root@server:$KEY_DIR/id_ed25519 /local/path/" >&2
      ok "Creation complete for $NEWUSER."
      ;;

    2)
      read -rp $'\nEnter the username to create (example: bob): ' NEWUSER
      if [ -z "$NEWUSER" ]; then fail "No username entered."; continue; fi
      echo "Paste the PUBLIC SSH key (single line), then press Enter:" >&2
      read -r PKEY
      PKEY="${PKEY:-}"
      if [ -z "$PKEY" ]; then fail "No public key provided."; continue; fi

      read -rp "Create user '$NEWUSER' and install pasted key? [y/N]: " CONF
      CONF=${CONF,,}
      if [[ "$CONF" != "y" && "$CONF" != "yes" ]]; then warn "Cancelled."; continue; fi

      step "Creating user and installing provided public key"
      if ! create_user_common "$NEWUSER"; then warn "User creation failed. Aborting."; continue; fi

      DATESTAMP="$(date '+%Y%m%d_%H%M%S')"
      KEY_DIR="$BASE_KEY_DIR/$NEWUSER-$DATESTAMP"
      mkdir -p "$KEY_DIR"
      echo "$PKEY" > "$KEY_DIR/from_user.pub"
      chmod 600 "$KEY_DIR/from_user.pub"
      ok "Saved provided public key to $KEY_DIR/from_user.pub"

      install_pubkey_for_user "$NEWUSER" "$PKEY"

      read -rp "Lock local password for $NEWUSER (disallow password login for this account)? [y/N]: " LOCKCHOICE
      LOCKCHOICE=${LOCKCHOICE,,}
      if [[ "$LOCKCHOICE" == "y" || "$LOCKCHOICE" == "yes" ]]; then
        passwd -l "$NEWUSER" >/dev/null 2>&1 || true
        ok "Password locked for $NEWUSER (key-only)."
      else
        ok "Password not locked; you may set one with: sudo passwd $NEWUSER"
      fi

      step "Verifying sudo for $NEWUSER"
      verify_sudo "$NEWUSER"

      step "Result & next steps"
      echo " - Stored provided public key: $KEY_DIR/from_user.pub" >&2
      ok "Creation complete for $NEWUSER."
      ;;

    3)
      read -rp $'\nEnter the username to DELETE: ' DELUSER
      if [ -z "$DELUSER" ]; then fail "No username entered."; continue; fi
      delete_user "$DELUSER" || { warn "Delete encountered problems or was aborted."; continue; }
      ;;

    4)
      step "Saved key directories:"
      if [ -d "$BASE_KEY_DIR" ]; then
        ls -la "$BASE_KEY_DIR" >&2 || true
      else
        echo "(no saved keys yet)" >&2
      fi
      echo "" >&2
      echo "To download a private key to your workstation, use scp. Example:" >&2
      echo "  scp root@server:$BASE_KEY_DIR/username-YYYYMMDD_HHMMSS/id_ed25519 /local/path/" >&2
      ;;

    5)
      echo "Exiting." >&2
      exit 0
      ;;

    *)
      warn "Invalid choice."
      ;;
  esac

  echo
  read -rp "Press Enter to return to menu..." dummy
done