#!/usr/bin/env bash
set -euo pipefail

# jumpuserctl - manage local users for an SSH jump host with per-user sshd drop-ins
#
# Features:
# - add/activate/deactivate/delete/list
# - per-user sshd drop-in in /etc/ssh/sshd_config.d/<user>.conf
# - mapping file: /etc/ssh/ssh-autoforward.map
# - automatically ensure a per-user SSH keypair exists (default: ed25519)
# - optionally push the public key to the target host and append to authorized_keys
#
# Notes on pushing keys:
# - Pushing the key requires that *some* authentication to the target works (password, existing admin key, etc.).
# - This script intentionally uses regular ssh (interactive) instead of insecure automation tricks.

SSHD_D_DIR="/etc/ssh/sshd_config.d"
FORWARD_CMD="/usr/share/devop-tools/ssh-autoforward.sh"
MAP_FILE="/etc/ssh/ssh-autoforward.map"

DEFAULT_LIST_PATTERN='^s[0-9]+$'
DEFAULT_KEY_TYPE="ed25519"
DEFAULT_KEY_BITS="4096" # only used for RSA
DEFAULT_KEY_NAME="id_ed25519" # may be changed depending on key type

require_root() {
  if [[ "${EUID}" -ne 0 ]]; then
    echo "ERROR: must be run as root" >&2
    exit 1
  fi
}

usage() {
  cat <<EOF
Usage:
  jumpuserctl add <user> <target_host> <target_user> [--uid <uid>] [--shell <shell>] [--key-type <type>] [--push-key] [--push-as <user>] [--sudo-on-target]
  jumpuserctl activate <user> [--key-type <type>] [--push-key] [--push-as <user>] [--sudo-on-target]
  jumpuserctl deactivate <user>
  jumpuserctl delete <user> [--keep-home]
  jumpuserctl list [--pattern <regex>]

Defaults:
  list pattern : ${DEFAULT_LIST_PATTERN}
  key type     : ${DEFAULT_KEY_TYPE}

Key push options:
  --push-key         Push the local user's public key to the target host and append it to authorized_keys.
                     This uses 'ssh' and may prompt for a password if no prior auth exists.
  --push-as <user>   Connect as a different remote login user (e.g. 'root' or 'admin').
                     If set, the key will be installed into the *target user* account using:
                       - 'sudo -u <target_user>' if --sudo-on-target is given
                       - otherwise, direct file append under ~<target_user> (requires sufficient permissions)
  --sudo-on-target   Use sudo on the target host when --push-as is used.

EOF
}

die() { echo "ERROR: $*" >&2; exit 1; }

valid_user() {
  local u="$1"
  # Allow common Unix usernames (not restricted to sNN).
  [[ "$u" =~ ^[a-z_][a-z0-9_-]{0,30}$ ]]
}

ensure_dirs() {
  [[ -d "${SSHD_D_DIR}" ]] || die "Missing directory: ${SSHD_D_DIR}"
  [[ -x "${FORWARD_CMD}" ]] || die "Missing/Not executable: ${FORWARD_CMD}"

  if [[ ! -f "${MAP_FILE}" ]]; then
    install -m 0644 -o root -g root /dev/null "${MAP_FILE}"
  else
    chmod 0644 "${MAP_FILE}" || true
  fi
}

user_exists() { id "$1" &>/dev/null; }

conf_path() {
  local u="$1"
  echo "${SSHD_D_DIR}/${u}.conf"
}

write_active_conf() {
  local u="$1"
  local p
  p="$(conf_path "$u")"
  cat > "${p}" <<EOF
Match User ${u}
        ForceCommand ${FORWARD_CMD}
        PermitTTY yes
        AllowTcpForwarding no
        X11Forwarding no
EOF
  chmod 0644 "${p}"
}

write_deactivated_conf() {
  local u="$1"
  local p
  p="$(conf_path "$u")"
  cat > "${p}" <<EOF
Match User ${u}
        AuthorizedKeysFile /dev/null
        ForceCommand /bin/false
        PermitTTY no
        AllowTcpForwarding no
        X11Forwarding no
EOF
  chmod 0644 "${p}"
}

reload_sshd() {
  if command -v systemctl >/dev/null 2>&1; then
    systemctl reload ssh 2>/dev/null || systemctl reload sshd 2>/dev/null || true
  fi
}

map_get() {
  local u="$1"
  awk -v u="$u" '
    $1 ~ /^#/ {next}
    NF < 3 {next}
    $1 == u {print $2, $3; exit}
  ' "${MAP_FILE}" 2>/dev/null || true
}

map_set() {
  local u="$1" h="$2" tu="$3"
  map_del "$u"
  printf "%-12s %-28s %s\n" "$u" "$h" "$tu" >> "${MAP_FILE}"
  chmod 0640 "${MAP_FILE}" || true
}

map_del() {
  local u="$1"
  if [[ -f "${MAP_FILE}" ]]; then
    awk -v u="$u" '
      $1 == u && $1 !~ /^#/ {next}
      {print}
    ' "${MAP_FILE}" > "${MAP_FILE}.tmp"
    mv "${MAP_FILE}.tmp" "${MAP_FILE}"
    chmod 0640 "${MAP_FILE}" || true
  fi
}

# Determine key filename based on key type.
key_filename_for_type() {
  local key_type="$1"
  case "${key_type}" in
    ed25519) echo "id_ed25519" ;;
    rsa)     echo "id_rsa" ;;
    ecdsa)   echo "id_ecdsa" ;;
    *)       die "Unsupported key type: ${key_type} (use ed25519|rsa|ecdsa)" ;;
  esac
}

ensure_user_ssh_dir() {
  local u="$1"
  local home="/home/${u}"
  install -d -m 0700 -o "${u}" -g "${u}" "${home}/.ssh"
  touch "${home}/.ssh/authorized_keys"
  chown "${u}:${u}" "${home}/.ssh/authorized_keys"
  chmod 0600 "${home}/.ssh/authorized_keys"
}

ensure_keypair() {
  local u="$1"
  local key_type="$2"
  local home="/home/${u}"

  local key_name
  key_name="$(key_filename_for_type "${key_type}")"

  local key_path="${home}/.ssh/${key_name}"
  local pub_path="${key_path}.pub"

  ensure_user_ssh_dir "${u}"

  if [[ -s "${key_path}" && -s "${pub_path}" ]]; then
    return 0
  fi

  # Generate a keypair without passphrase (typical for automation jump users).
  # If you want passphrases, remove -N "" and handle agents accordingly.
  local gen_args=()
  gen_args+=(-t "${key_type}")
  if [[ "${key_type}" == "rsa" ]]; then
    gen_args+=(-b "${DEFAULT_KEY_BITS}")
  fi
  gen_args+=(-f "${key_path}" -N "" -C "${u}@$(hostname -f 2>/dev/null || hostname)")

  # Run ssh-keygen as the target user so file ownership is correct.
  sudo -u "${u}" /usr/bin/ssh-keygen "${gen_args[@]}" >/dev/null

  # Enforce permissions
  chown "${u}:${u}" "${key_path}" "${pub_path}"
  chmod 0600 "${key_path}"
  chmod 0644 "${pub_path}"
}

push_key_to_target() {
  local u="$1"
  local push_as="$2"          # empty means connect as target_user
  local sudo_on_target="$3"   # "yes" or "no"

  local mg
  mg="$(map_get "${u}")"
  [[ -n "${mg}" ]] || die "No mapping for user ${u} in ${MAP_FILE}"

  local target_host target_user
  target_host="$(awk '{print $1}' <<<"${mg}")"
  target_user="$(awk '{print $2}' <<<"${mg}")"

  local home="/home/${u}"

  # Prefer the key type file that exists, otherwise assume ed25519.
  local pub_candidates=("${home}/.ssh/id_ed25519.pub" "${home}/.ssh/id_rsa.pub" "${home}/.ssh/id_ecdsa.pub")
  local pub_path=""
  for p in "${pub_candidates[@]}"; do
    if [[ -s "${p}" ]]; then
      pub_path="${p}"
      break
    fi
  done
  [[ -n "${pub_path}" ]] || die "No public key found for ${u}. Run add/activate to generate it."

  local pub_key
  pub_key="$(cat "${pub_path}")"

  # Build remote login user
  local remote_login="${target_user}@${target_host}"
  if [[ -n "${push_as}" ]]; then
    remote_login="${push_as}@${target_host}"
  fi

  # Remote command:
  # - Ensure ~/.ssh exists and authorized_keys exists
  # - Append key only if not already present (exact line match)
  # - If we connect as a different user, optionally use sudo -u <target_user> to install key
  local remote_cmd=""

  if [[ -n "${push_as}" ]]; then
    if [[ "${sudo_on_target}" == "yes" ]]; then
      remote_cmd=$(
        cat <<EOF
set -e
sudo -u '${target_user}' sh -c "umask 077; mkdir -p ~/.ssh; touch ~/.ssh/authorized_keys"
sudo -u '${target_user}' sh -c "grep -qxF '${pub_key}' ~/.ssh/authorized_keys || printf '%s\n' '${pub_key}' >> ~/.ssh/authorized_keys"
EOF
      )
    else
      # This requires that push_as has permissions to write into target_user's home.
      # Many systems won't allow that unless push_as is root.
      remote_cmd=$(
        cat <<EOF
set -e
TARGET_HOME=\$(getent passwd '${target_user}' | cut -d: -f6)
[ -n "\${TARGET_HOME}" ] || exit 1
umask 077
mkdir -p "\${TARGET_HOME}/.ssh"
touch "\${TARGET_HOME}/.ssh/authorized_keys"
grep -qxF '${pub_key}' "\${TARGET_HOME}/.ssh/authorized_keys" || printf '%s\n' '${pub_key}' >> "\${TARGET_HOME}/.ssh/authorized_keys"
EOF
      )
    fi
  else
    # Connect directly as target_user
    remote_cmd=$(
      cat <<EOF
set -e
umask 077
mkdir -p ~/.ssh
touch ~/.ssh/authorized_keys
grep -qxF '${pub_key}' ~/.ssh/authorized_keys || printf '%s\n' '${pub_key}' >> ~/.ssh/authorized_keys
EOF
    )
  fi

  echo "INFO: Pushing key for local user '${u}' to '${remote_login}' (target account: ${target_user}@${target_host})" >&2
  /usr/bin/ssh -o StrictHostKeyChecking=accept-new "${remote_login}" "${remote_cmd}"

  echo "OK: key installed on target for ${target_user}@${target_host}"
}

cmd_add() {
  local u="$1" h="$2" tu="$3"
  shift 3

  local uid=""
  local shell="/bin/sh"
  local key_type="${DEFAULT_KEY_TYPE}"
  local do_push="no"
  local push_as=""
  local sudo_on_target="no"

  while [[ $# -gt 0 ]]; do
    case "$1" in
      --uid) uid="${2:-}"; shift 2;;
      --shell) shell="${2:-}"; shift 2;;
      --key-type) key_type="${2:-}"; shift 2;;
      --push-key) do_push="yes"; shift 1;;
      --push-as) push_as="${2:-}"; shift 2;;
      --sudo-on-target) sudo_on_target="yes"; shift 1;;
      *) die "Unknown option: $1";;
    esac
  done

  valid_user "$u" || die "Invalid username: $u"

  if user_exists "$u"; then
    die "User already exists: $u"
  fi

  local args=(--create-home --home-dir "/home/${u}" --shell "${shell}" --comment "")
  if [[ -n "${uid}" ]]; then
    [[ "${uid}" =~ ^[0-9]+$ ]] || die "--uid must be numeric"
    args+=(--uid "${uid}")
  fi

  useradd "${args[@]}" "${u}"

  map_set "${u}" "${h}" "${tu}"
  write_active_conf "${u}"

  ensure_keypair "${u}" "${key_type}"

  if [[ "${do_push}" == "yes" ]]; then
    push_key_to_target "${u}" "${push_as}" "${sudo_on_target}"
  fi

  reload_sshd
  echo "OK: added user ${u} -> ${tu}@${h}"
}

cmd_activate() {
  local u="$1"
  shift 1

  local key_type="${DEFAULT_KEY_TYPE}"
  local do_push="no"
  local push_as=""
  local sudo_on_target="no"

  while [[ $# -gt 0 ]]; do
    case "$1" in
      --key-type) key_type="${2:-}"; shift 2;;
      --push-key) do_push="yes"; shift 1;;
      --push-as) push_as="${2:-}"; shift 2;;
      --sudo-on-target) sudo_on_target="yes"; shift 1;;
      *) die "Unknown option: $1";;
    esac
  done

  user_exists "$u" || die "User does not exist: $u"

  if [[ -z "$(map_get "$u")" ]]; then
    die "No mapping for user ${u} in ${MAP_FILE}"
  fi

  write_active_conf "$u"
  usermod --shell /bin/sh "$u" >/dev/null 2>&1 || true

  ensure_keypair "${u}" "${key_type}"

  if [[ "${do_push}" == "yes" ]]; then
    push_key_to_target "${u}" "${push_as}" "${sudo_on_target}"
  fi

  reload_sshd
  echo "OK: activated ${u}"
}

cmd_deactivate() {
  local u="$1"
  user_exists "$u" || die "User does not exist: $u"

  write_deactivated_conf "$u"
  usermod --shell /usr/sbin/nologin "$u" >/dev/null 2>&1 || true

  reload_sshd
  echo "OK: deactivated ${u}"
}

cmd_delete() {
  local u="$1"
  shift 1

  local keep_home="no"
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --keep-home) keep_home="yes"; shift 1;;
      *) die "Unknown option: $1";;
    esac
  done

  if user_exists "$u"; then
    if [[ "${keep_home}" == "yes" ]]; then
      userdel "${u}"
    else
      userdel --remove "${u}"
    fi
  fi

  rm -f "$(conf_path "$u")" || true
  map_del "$u"

  reload_sshd
  echo "OK: deleted ${u}"
}

cmd_list() {
  local pattern="${DEFAULT_LIST_PATTERN}"
  if [[ "${1:-}" == "--pattern" ]]; then
    pattern="${2:-}"
  fi

  printf "%-12s %-6s %-6s %-18s %-28s %-12s %s\n" "USER" "UID" "GID" "SHELL" "HOME" "STATE" "TARGET"
  printf "%-12s %-6s %-6s %-18s %-28s %-12s %s\n" "----" "---" "---" "-----" "----" "-----" "------"

  while IFS=: read -r name _ uid gid _ home shell; do
    [[ "${name}" =~ ${pattern} ]] || continue

    local state="unknown"
    local target="-"
    local conf
    conf="$(conf_path "${name}")"

    if [[ -f "${conf}" ]]; then
      if grep -q "ForceCommand ${FORWARD_CMD}" "${conf}" 2>/dev/null; then
        state="active"
      elif grep -q "AuthorizedKeysFile /dev/null" "${conf}" 2>/dev/null; then
        state="disabled"
      else
        state="custom"
      fi
    else
      state="no-conf"
    fi

    local mg
    mg="$(map_get "${name}")"
    if [[ -n "${mg}" ]]; then
      local h tu
      h="$(awk '{print $1}' <<<"${mg}")"
      tu="$(awk '{print $2}' <<<"${mg}")"
      target="${tu}@${h}"
    fi

    printf "%-12s %-6s %-6s %-18s %-28s %-12s %s\n" "${name}" "${uid}" "${gid}" "${shell}" "${home}" "${state}" "${target}"
  done < /etc/passwd
}

main() {
  require_root
  ensure_dirs

  local cmd="${1:-}"
  shift || true

  case "${cmd}" in
    add)
      [[ $# -ge 3 ]] || { usage; exit 1; }
      cmd_add "$@"
      ;;
    activate)
      [[ $# -ge 1 ]] || { usage; exit 1; }
      cmd_activate "$@"
      ;;
    deactivate)
      [[ $# -eq 1 ]] || { usage; exit 1; }
      cmd_deactivate "$1"
      ;;
    delete)
      [[ $# -ge 1 ]] || { usage; exit 1; }
      cmd_delete "$@"
      ;;
    list)
      cmd_list "$@"
      ;;
    -h|--help|"")
      usage
      ;;
    *)
      die "Unknown command: ${cmd}"
      ;;
  esac
}

main "$@"
