#!/bin/bash
set -Eeuo pipefail

STATE_DIR="/var/lib/k8s-mgmt-pod"
GENERATED_PASSWORDS_FILE="${STATE_DIR}/generated-passwords.env"
SSH_DYNAMIC_CONFIG="/etc/ssh/sshd_config.d/20-k8s-mgmt-pod-auth.conf"
PASSWORD_AUTH_REQUIRED=0
SSHD_PID=""
WEBSSH2_PID=""
NGINX_PID=""
BOOTSTRAP_KUBECONFIG="${BOOTSTRAP_KUBECONFIG:-/bootstrap/kubeconfig/config}"
ENABLE_BOOTSTRAP_KUBECONFIG="${ENABLE_BOOTSTRAP_KUBECONFIG:-false}"

ensure_host_keys() {
  if [[ ! -f "${STATE_DIR}/ssh_host_rsa_key" ]]; then
    ssh-keygen -q -t rsa -b 3072 -N '' -f "${STATE_DIR}/ssh_host_rsa_key"
  fi

  if [[ ! -f "${STATE_DIR}/ssh_host_ecdsa_key" ]]; then
    ssh-keygen -q -t ecdsa -b 256 -N '' -f "${STATE_DIR}/ssh_host_ecdsa_key"
  fi

  if [[ ! -f "${STATE_DIR}/ssh_host_ed25519_key" ]]; then
    ssh-keygen -q -t ed25519 -N '' -f "${STATE_DIR}/ssh_host_ed25519_key"
  fi

  chmod 0600 "${STATE_DIR}/ssh_host_"*_key
}

mkdir -p "${STATE_DIR}"
touch "${GENERATED_PASSWORDS_FILE}"
chmod 0600 "${GENERATED_PASSWORDS_FILE}"

log() {
  printf '[k8s-mgmt-pod] %s\n' "$*"
}

cleanup() {
  local exit_code=$?

  if [[ -n "${WEBSSH2_PID}" ]] && kill -0 "${WEBSSH2_PID}" 2>/dev/null; then
    kill "${WEBSSH2_PID}" 2>/dev/null || true
  fi

  if [[ -n "${SSHD_PID}" ]] && kill -0 "${SSHD_PID}" 2>/dev/null; then
    kill "${SSHD_PID}" 2>/dev/null || true
  fi

  if [[ -n "${NGINX_PID}" ]] && kill -0 "${NGINX_PID}" 2>/dev/null; then
    kill "${NGINX_PID}" 2>/dev/null || true
  fi

  wait || true
  exit "${exit_code}"
}

trap cleanup EXIT INT TERM

store_generated_password() {
  local username="$1"
  local password="$2"

  if grep -q "^${username}=" "${GENERATED_PASSWORDS_FILE}"; then
    sed -i "s|^${username}=.*|${username}=${password}|" "${GENERATED_PASSWORDS_FILE}"
  else
    printf '%s=%s\n' "${username}" "${password}" >> "${GENERATED_PASSWORDS_FILE}"
  fi
}

read_generated_password() {
  local username="$1"

  # Return an empty string if no entry exists; never fail under `set -euo pipefail`.
  awk -F '=' -v username="${username}" '$1 == username { value = substr($0, index($0, "=") + 1) } END { if (value != "") print value }' "${GENERATED_PASSWORDS_FILE}"
}

write_authorized_keys() {
  local username="$1"
  local pubkey_content="$2"
  local home_dir="/home/${username}"
  local ssh_dir="${home_dir}/.ssh"
  local auth_keys="${ssh_dir}/authorized_keys"

  mkdir -p "${ssh_dir}"
  chmod 0700 "${ssh_dir}"
  printf '%s\n' "${pubkey_content}" > "${auth_keys}"
  chmod 0600 "${auth_keys}"
  chown -R "${username}:${username}" "${ssh_dir}"
}

install_kubeconfig_for_user() {
  local username="$1"
  local home_dir="/home/${username}"
  local kube_dir="${home_dir}/.kube"
  local kube_config="${kube_dir}/config"

  mkdir -p "${kube_dir}"
  install -m 0600 -o "${username}" -g "${username}" "${BOOTSTRAP_KUBECONFIG}" "${kube_config}"
}

install_bootstrap_kubeconfig() {
  if [[ "${ENABLE_BOOTSTRAP_KUBECONFIG}" != "true" ]]; then
    return
  fi

  if [[ ! -f "${BOOTSTRAP_KUBECONFIG}" ]]; then
    log "Skipping kubeconfig bootstrap because ${BOOTSTRAP_KUBECONFIG} is not present."
    return
  fi

  install_kubeconfig_for_user admin
  install_kubeconfig_for_user user
  log "Installed mounted kubeconfig for admin and user."
}

remove_authorized_keys() {
  local username="$1"
  local auth_keys="/home/${username}/.ssh/authorized_keys"

  rm -f "${auth_keys}"
}

set_login_password() {
  local username="$1"
  local password="$2"

  printf '%s:%s\n' "${username}" "${password}" | chpasswd
  passwd -u "${username}" >/dev/null 2>&1 || true
}

generate_password() {
  openssl rand -base64 24 | tr -d '\n'
}

configure_user_auth() {
  local username="$1"
  local pubkey_var_name="$2"
  local password_var_name="$3"
  local pubkey_value="${!pubkey_var_name:-}"
  local password_value="${!password_var_name:-}"
  local generated_password=""
  local has_pubkey=0

  if [[ -n "${pubkey_value}" ]]; then
    write_authorized_keys "${username}" "${pubkey_value}"
    has_pubkey=1
    log "Configured SSH public key authentication for ${username}."
  else
    remove_authorized_keys "${username}"
  fi

  if [[ -n "${password_value}" ]]; then
    PASSWORD_AUTH_REQUIRED=1
    set_login_password "${username}" "${password_value}"
    log "Configured password authentication for ${username} from environment."
    return
  fi

  if [[ "${has_pubkey}" -eq 1 ]]; then
    passwd -l "${username}" >/dev/null 2>&1 || true
    return
  fi

  PASSWORD_AUTH_REQUIRED=1

  generated_password="$(read_generated_password "${username}")"
  if [[ -z "${generated_password}" ]]; then
    generated_password="$(generate_password)"
    store_generated_password "${username}" "${generated_password}"
    log "Generated random password for ${username}: ${generated_password}"
    log "Store this password securely now; it is only printed when first generated."
  fi

  set_login_password "${username}" "${generated_password}"
}

render_sshd_auth_config() {
  local password_value="no"
  local keyboard_interactive_value="no"

  if [[ "${PASSWORD_AUTH_REQUIRED}" -eq 1 ]]; then
    password_value="yes"
    keyboard_interactive_value="yes"
    export WEBSSH2_AUTH_ALLOWED="password,publickey,keyboard-interactive"
  else
    export WEBSSH2_AUTH_ALLOWED="publickey"
  fi

  cat > "${SSH_DYNAMIC_CONFIG}" <<EOF
PasswordAuthentication ${password_value}
KbdInteractiveAuthentication ${keyboard_interactive_value}
EOF
}

prepare_runtime() {
  mkdir -p /var/run/sshd
  ensure_host_keys
  install_bootstrap_kubeconfig
  configure_user_auth admin ADMIN_SSH_PUBKEY ADMIN_PASSWORD
  configure_user_auth user USER_SSH_PUBKEY USER_PASSWORD
  render_sshd_auth_config
  /usr/sbin/sshd -t
}

start_services() {
  /usr/sbin/sshd -D -e &
  SSHD_PID=$!

  NODE_ENV=production npm --prefix /opt/webssh2 start &
  WEBSSH2_PID=$!

  nginx -g 'daemon off;' &
  NGINX_PID=$!

  log "sshd started on port 22."
  log "WebSSH2 started on port ${WEBSSH2_LISTEN_PORT}."
  log "nginx redirect/proxy started on port ${WEB_SSH_PORT}."

  wait -n "${SSHD_PID}" "${WEBSSH2_PID}" "${NGINX_PID}"
}

# Only public key material is accepted via environment variables; never log secret key data.
prepare_runtime
start_services