#!/usr/bin/env bash
#
# sshuser -- manage SSH users on a hardened box (Alpine, Debian, or Alma).
#
# Two roles, matching harden-jumphost.sh:
#   admin   -> group ssh-admins,  interactive shell  (full shell)
#   jumper  -> group ssh-jumpers, nologin shell       (ProxyJump only)
#
# The interactive shell and nologin paths, and the user-management commands,
# differ per distro -- this script detects the OS and adapts (it's installed
# standalone, so it can't share scripts/oslib.sh).
#
# Two modes:
#   - TUI (gum)  : run with no command, or any command with missing args
#   - CLI flags  : full automation, suitable for shell scripts and CI
#
# Install:
#   install -m 0755 sshuser.sh /usr/local/bin/sshuser
#   gum is only needed for TUI mode (apk add gum / apt install gum / dnf install gum)
#
# Usage:
#   sshuser                                              # interactive TUI
#   sshuser add -u alice -r jumper -f alice.pub
#   sshuser add -u bob   -r admin  -k "ssh-ed25519 AAA..."
#   sshuser edit -u alice --add-key "ssh-ed25519 BBB..."
#   sshuser edit -u alice --remove-key "comment-substring"
#   sshuser edit -u alice --role admin
#   sshuser remove -u alice -y
#   sshuser list
#   sshuser list -r jumper
#   sshuser show -u alice

set -euo pipefail

ADMIN_GROUP="ssh-admins"
JUMPER_GROUP="ssh-jumpers"

# ---------------------------------------------------------------------------
# OS detection + per-distro user-management primitives. (Standalone install,
# so we can't source oslib.sh -- this is the minimal slice we need.)
# ---------------------------------------------------------------------------
_OS_FAMILY=alpine
if [[ -r /etc/os-release ]]; then
    _osid="$(. /etc/os-release 2>/dev/null && echo "${ID:-}")"
    _oslike="$(. /etc/os-release 2>/dev/null && echo "${ID_LIKE:-}")"
    case " ${_osid} ${_oslike} " in
        *" debian "*|*" ubuntu "*)               _OS_FAMILY=debian ;;
        *" rhel "*|*" fedora "*|*" centos "*|*" almalinux "*|*" rocky "*) _OS_FAMILY=rhel ;;
        *" alpine "*)                            _OS_FAMILY=alpine ;;
    esac
fi

# Admin (interactive) shell: ash on Alpine, bash elsewhere.
case "$_OS_FAMILY" in
    alpine) ADMIN_SHELL="/bin/ash" ;;
    *)      ADMIN_SHELL="/bin/bash" ;;
esac
[[ -x "$ADMIN_SHELL" ]] || ADMIN_SHELL="/bin/sh"
# Nologin path: /usr/sbin/nologin on Debian, /sbin/nologin on Alpine/Alma.
JUMPER_SHELL="/sbin/nologin"
[[ -x "$JUMPER_SHELL" ]] || for _p in /usr/sbin/nologin /sbin/nologin; do
    [[ -x "$_p" ]] && { JUMPER_SHELL="$_p"; break; }
done

user_create()      { case "$_OS_FAMILY" in alpine) adduser -D -s "$2" -g "" "$1";; *) useradd -m -s "$2" "$1";; esac; }
user_join_group()  { case "$_OS_FAMILY" in alpine) adduser "$1" "$2";; *) usermod -aG "$2" "$1";; esac; }
user_leave_group() { case "$_OS_FAMILY" in alpine) deluser "$1" "$2" 2>/dev/null || true;; *) gpasswd -d "$1" "$2" 2>/dev/null || true;; esac; }
user_delete()      { case "$_OS_FAMILY" in alpine) deluser --remove-home "$1" 2>/dev/null || deluser "$1";; *) userdel -r "$1" 2>/dev/null || userdel "$1";; esac; }
set_user_shell()   { usermod -s "$2" "$1" 2>/dev/null || chsh -s "$2" "$1" 2>/dev/null \
                       || sed -i "s|^\($1:.*:\)[^:]*$|\1$2|" /etc/passwd; }

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
err()  { printf '\033[1;31m[x]\033[0m %s\n' "$*" >&2; exit 1; }
warn() { printf '\033[1;33m[!]\033[0m %s\n' "$*" >&2; }
log()  { printf '\033[1;32m[+]\033[0m %s\n' "$*"; }

have_gum() { command -v gum >/dev/null 2>&1; }

require_root() {
    [[ $EUID -eq 0 ]] || err "Run as root."
}

require_groups() {
    getent group "$ADMIN_GROUP"  >/dev/null || err "Group $ADMIN_GROUP missing -- run harden-jumphost.sh first."
    getent group "$JUMPER_GROUP" >/dev/null || err "Group $JUMPER_GROUP missing -- run harden-jumphost.sh first."
}

confirm() {
    # confirm "Prompt" -- returns 0 if yes, 1 if no.
    local prompt="$1"
    if [[ "${YES:-0}" == "1" ]]; then return 0; fi
    if have_gum; then
        gum confirm "$prompt"
    else
        read -r -p "$prompt [y/N] " ans
        [[ "${ans,,}" == "y" || "${ans,,}" == "yes" ]]
    fi
}

ask() {
    # ask "Prompt" [default] -- prints user input.
    local prompt="$1"
    local default="${2:-}"
    if have_gum; then
        gum input --prompt "$prompt: " --placeholder "$default"
    else
        local ans
        read -r -p "$prompt${default:+ [$default]}: " ans
        echo "${ans:-$default}"
    fi
}

choose() {
    # choose "Prompt" opt1 opt2 ... -- prints chosen.
    local prompt="$1"; shift
    if have_gum; then
        gum choose --header "$prompt" "$@"
    else
        printf '%s\n' "$prompt" >&2
        local i=1; for o in "$@"; do printf '  %d) %s\n' "$i" "$o" >&2; ((i++)); done
        local n; read -r -p "Choice: " n
        echo "${!n}"
    fi
}

valid_pubkey() {
    # Quick sanity check: starts with a known algo prefix.
    [[ "$1" =~ ^(ssh-ed25519|ssh-rsa|sk-ssh-ed25519@openssh\.com|ecdsa-sha2-)[[:space:]]+[A-Za-z0-9+/=]+ ]]
}

resolve_role() {
    # Map role -> group + shell. echoes "GROUP SHELL".
    case "$1" in
        admin)  echo "$ADMIN_GROUP $ADMIN_SHELL" ;;
        jumper) echo "$JUMPER_GROUP $JUMPER_SHELL" ;;
        *) err "Unknown role: $1 (must be admin|jumper)" ;;
    esac
}

ssh_dir_setup() {
    # Ensure ~user/.ssh exists with correct perms and ownership.
    local user="$1"
    local home; home=$(getent passwd "$user" | cut -d: -f6)
    [[ -n "$home" ]] || err "User $user has no home directory."
    install -d -m 0700 -o "$user" -g "$user" "$home/.ssh"
    install -m 0600 -o "$user" -g "$user" /dev/null "$home/.ssh/authorized_keys" 2>/dev/null || \
        touch "$home/.ssh/authorized_keys"
    chown "$user:$user" "$home/.ssh/authorized_keys"
    chmod 600 "$home/.ssh/authorized_keys"
    echo "$home/.ssh/authorized_keys"
}

# ---------------------------------------------------------------------------
# Commands
# ---------------------------------------------------------------------------

cmd_add() {
    local user="${USER_ARG:-}" role="${ROLE_ARG:-}" key="${KEY_ARG:-}"
    [[ -z "$user" ]] && user=$(ask "Username")
    [[ -z "$user" ]] && err "Username required."
    if id "$user" >/dev/null 2>&1; then
        err "User $user already exists. Use 'sshuser edit' instead."
    fi
    [[ -z "$role" ]] && role=$(choose "Role" admin jumper)
    read -r group shell <<<"$(resolve_role "$role")"

    if [[ -z "$key" ]]; then
        if [[ -n "${KEY_FILE_ARG:-}" ]]; then
            [[ -f "$KEY_FILE_ARG" ]] || err "Key file not found: $KEY_FILE_ARG"
            key=$(<"$KEY_FILE_ARG")
        elif have_gum; then
            key=$(gum input --prompt "SSH public key (paste full line, blank to skip): " --width 200)
        else
            read -r -p "SSH public key (paste full line, blank to skip): " key
        fi
    fi
    if [[ -n "$key" ]]; then
        valid_pubkey "$key" || err "That doesn't look like a valid SSH public key."
    fi

    confirm "Create user $user as $role (group $group, shell $shell)?" || { warn "Aborted."; exit 1; }

    log "Creating $user..."
    user_create "$user" "$shell"
    user_join_group "$user" "$group"

    if [[ -n "$key" ]]; then
        local ak; ak=$(ssh_dir_setup "$user")
        if grep -qxF "$key" "$ak" 2>/dev/null; then
            warn "Key already present, skipping."
        else
            echo "$key" >> "$ak"
            log "Added SSH key to $ak"
        fi
    else
        warn "No SSH key added. User cannot log in until 'sshuser edit -u $user --add-key ...' is run."
    fi
    log "Done."
}

cmd_edit() {
    local user="${USER_ARG:-}"
    [[ -z "$user" ]] && user=$(ask "Username to edit")
    id "$user" >/dev/null 2>&1 || err "No such user: $user"

    # If no specific edit flag, prompt for action.
    local has_action=0
    [[ -n "${ADD_KEY_ARG:-}" ]] && has_action=1
    [[ -n "${REMOVE_KEY_ARG:-}" ]] && has_action=1
    [[ -n "${ROLE_ARG:-}" ]] && has_action=1
    [[ -n "${SHELL_ARG:-}" ]] && has_action=1

    if [[ $has_action -eq 0 ]]; then
        local action
        action=$(choose "What to do for $user?" \
            "add ssh key" "remove ssh key" "change role" "change shell" "cancel")
        case "$action" in
            "add ssh key")    ADD_KEY_ARG=$(ask "Paste SSH public key") ;;
            "remove ssh key") REMOVE_KEY_ARG=$(ask "Substring of key to remove (comment is fine)") ;;
            "change role")    ROLE_ARG=$(choose "New role" admin jumper) ;;
            "change shell")   SHELL_ARG=$(ask "New shell" "$ADMIN_SHELL") ;;
            *) warn "Cancelled."; return 0 ;;
        esac
    fi

    if [[ -n "${ROLE_ARG:-}" ]]; then
        read -r new_group new_shell <<<"$(resolve_role "$ROLE_ARG")"
        log "Setting $user role to $ROLE_ARG (group $new_group, shell $new_shell)"
        # Remove from the other ssh group, add to the target.
        local other
        [[ "$new_group" == "$ADMIN_GROUP" ]] && other="$JUMPER_GROUP" || other="$ADMIN_GROUP"
        user_leave_group "$user" "$other"
        user_join_group "$user" "$new_group"
        set_user_shell "$user" "$new_shell"
    fi

    if [[ -n "${SHELL_ARG:-}" ]]; then
        log "Setting $user shell to $SHELL_ARG"
        set_user_shell "$user" "$SHELL_ARG"
    fi

    if [[ -n "${ADD_KEY_ARG:-}" ]]; then
        valid_pubkey "$ADD_KEY_ARG" || err "Invalid pubkey."
        local ak; ak=$(ssh_dir_setup "$user")
        if grep -qxF "$ADD_KEY_ARG" "$ak" 2>/dev/null; then
            warn "Key already present."
        else
            echo "$ADD_KEY_ARG" >> "$ak"
            log "Added key to $ak"
        fi
    fi

    if [[ -n "${REMOVE_KEY_ARG:-}" ]]; then
        local ak; ak=$(ssh_dir_setup "$user")
        local before after
        before=$(wc -l <"$ak")
        # Remove any line containing the substring (safe escape for sed).
        local pat; pat=$(printf '%s\n' "$REMOVE_KEY_ARG" | sed 's/[][\.*^$/]/\\&/g')
        sed -i "/$pat/d" "$ak"
        after=$(wc -l <"$ak")
        log "Removed $((before - after)) key line(s) from $ak"
    fi
}

cmd_remove() {
    local user="${USER_ARG:-}"
    [[ -z "$user" ]] && user=$(ask "Username to remove")
    id "$user" >/dev/null 2>&1 || err "No such user: $user"
    [[ "$user" == "root" ]] && err "Refusing to remove root."

    confirm "DELETE user $user and their home directory?" || { warn "Aborted."; exit 1; }
    log "Deleting $user..."
    user_delete "$user"
    log "Done."
}

cmd_list() {
    local role="${ROLE_ARG:-}"
    local groups=()
    case "$role" in
        admin)  groups=("$ADMIN_GROUP") ;;
        jumper) groups=("$JUMPER_GROUP") ;;
        "")     groups=("$ADMIN_GROUP" "$JUMPER_GROUP") ;;
        *) err "Unknown role: $role" ;;
    esac

    printf '%-20s %-15s %-20s %s\n' "USER" "ROLE" "SHELL" "KEYS"
    printf '%-20s %-15s %-20s %s\n' "----" "----" "-----" "----"
    for g in "${groups[@]}"; do
        local label="admin"
        [[ "$g" == "$JUMPER_GROUP" ]] && label="jumper"
        # getent group returns: groupname:x:gid:user1,user2,...
        local members; members=$(getent group "$g" | awk -F: '{print $4}' | tr ',' ' ')
        for u in $members; do
            [[ -z "$u" ]] && continue
            local home; home=$(getent passwd "$u" | cut -d: -f6)
            local shell; shell=$(getent passwd "$u" | cut -d: -f7)
            local nkeys=0
            [[ -f "$home/.ssh/authorized_keys" ]] && nkeys=$(grep -cv '^\s*$\|^\s*#' "$home/.ssh/authorized_keys" 2>/dev/null || echo 0)
            printf '%-20s %-15s %-20s %s\n' "$u" "$label" "$shell" "$nkeys"
        done
    done
}

cmd_show() {
    local user="${USER_ARG:-}"
    [[ -z "$user" ]] && user=$(ask "Username")
    id "$user" >/dev/null 2>&1 || err "No such user: $user"
    local home; home=$(getent passwd "$user" | cut -d: -f6)
    local shell; shell=$(getent passwd "$user" | cut -d: -f7)
    local groups; groups=$(id -nG "$user" | tr ' ' ',')
    echo "user:    $user"
    echo "home:    $home"
    echo "shell:   $shell"
    echo "groups:  $groups"
    echo "keys:"
    if [[ -f "$home/.ssh/authorized_keys" ]]; then
        awk 'NF && !/^#/ {
            n=split($0, p, " ");
            type=p[1]; comment=(n>=3 ? p[n] : "(no comment)");
            printf "  - %s  %s\n", type, comment
        }' "$home/.ssh/authorized_keys"
    else
        echo "  (none)"
    fi
}

cmd_tui() {
    have_gum || err "TUI mode requires gum (apk/apt/dnf install gum). Or use CLI flags (sshuser --help)."
    local action
    action=$(choose "What do you want to do?" \
        "add user" "edit user" "remove user" "list users" "show user" "quit")
    case "$action" in
        "add user")    cmd_add ;;
        "edit user")   cmd_edit ;;
        "remove user") cmd_remove ;;
        "list users")  cmd_list ;;
        "show user")   cmd_show ;;
        *) ;;
    esac
}

# ---------------------------------------------------------------------------
# Argument parsing
# ---------------------------------------------------------------------------
usage() {
    sed -n '2,30p' "$0" | sed 's/^# \{0,1\}//'
}

CMD=""
USER_ARG=""
ROLE_ARG=""
KEY_ARG=""
KEY_FILE_ARG=""
ADD_KEY_ARG=""
REMOVE_KEY_ARG=""
SHELL_ARG=""
YES=0

if [[ $# -gt 0 && "$1" != -* ]]; then
    CMD="$1"; shift
fi

while [[ $# -gt 0 ]]; do
    case "$1" in
        -u|--user)         USER_ARG="$2"; shift 2 ;;
        -r|--role)         ROLE_ARG="$2"; shift 2 ;;
        -k|--key)          KEY_ARG="$2"; shift 2 ;;
        -f|--key-file)     KEY_FILE_ARG="$2"; shift 2 ;;
        --add-key)         ADD_KEY_ARG="$2"; shift 2 ;;
        --remove-key)      REMOVE_KEY_ARG="$2"; shift 2 ;;
        --shell)           SHELL_ARG="$2"; shift 2 ;;
        -y|--yes)          YES=1; shift ;;
        -h|--help)         usage; exit 0 ;;
        *) err "Unknown flag: $1" ;;
    esac
done

require_root
require_groups

case "$CMD" in
    add)    cmd_add ;;
    edit)   cmd_edit ;;
    remove|rm|del) cmd_remove ;;
    list|ls) cmd_list ;;
    show)   cmd_show ;;
    "")     cmd_tui ;;
    *) err "Unknown command: $CMD (try: add edit remove list show)" ;;
esac