#!/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"
K8S_RUNTIME_ENV_PROFILE="/etc/profile.d/30-kubernetes-runtime-env.sh"
PASSWORD_AUTH_REQUIRED=0
SSHD_PID=""
WEBSSH2_PID=""
NGINX_PID=""
BOOTSTRAP_KUBECONFIG="${BOOTSTRAP_KUBECONFIG:-/bootstrap/kubeconfig/config}"
ENABLE_BOOTSTRAP_KUBECONFIG="${ENABLE_BOOTSTRAP_KUBECONFIG:-false}"
KUBERNETES_SECRETS_DIR="${KUBERNETES_SECRETS_DIR:-/var/run/kubernetes.io}"
SSH_PORT="${SSH_PORT:-2222}"

if [[ -z "${WEBSSH2_SSH_PORT:-}" ]]; then
  WEBSSH2_SSH_PORT="${SSH_PORT}"
fi
export WEBSSH2_SSH_PORT

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 source_kubeconfig="$2"
  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}" "${source_kubeconfig}" "${kube_config}"
}

resolve_kubernetes_secrets_dir() {
  local configured_dir="${KUBERNETES_SECRETS_DIR}"
  local fallback_dir="/var/run/secrets/kubernetes.io/serviceaccount"

  if [[ -f "${configured_dir}/token" && -f "${configured_dir}/ca.crt" && -f "${configured_dir}/namespace" ]]; then
    printf '%s\n' "${configured_dir}"
    return 0
  fi

  if [[ "${configured_dir}" != "${fallback_dir}" ]] \
    && [[ -f "${fallback_dir}/token" && -f "${fallback_dir}/ca.crt" && -f "${fallback_dir}/namespace" ]]; then
    printf '%s\n' "${fallback_dir}"
    return 0
  fi

  return 1
}

install_serviceaccount_kubeconfig() {
  local sa_dir=""
  local token_file=""
  local ca_file=""
  local namespace_file=""
  local namespace=""
  local kubernetes_host="${KUBERNETES_SERVICE_HOST:-kubernetes.default.svc}"
  local kubernetes_port="${KUBERNETES_SERVICE_PORT_HTTPS:-443}"
  local server=""
  local generated_kubeconfig="${STATE_DIR}/incluster.kubeconfig"

  if ! sa_dir="$(resolve_kubernetes_secrets_dir)"; then
    log "Skipping in-cluster kubeconfig generation because no service-account mount was found in ${KUBERNETES_SECRETS_DIR} or /var/run/secrets/kubernetes.io/serviceaccount."
    return
  fi

  token_file="${sa_dir}/token"
  ca_file="${sa_dir}/ca.crt"
  namespace_file="${sa_dir}/namespace"
  namespace="$(tr -d '\n' < "${namespace_file}")"
  if [[ -z "${namespace}" ]]; then
    namespace="default"
  fi

  server="https://${kubernetes_host}:${kubernetes_port}"

  cat > "${generated_kubeconfig}" <<EOF
apiVersion: v1
kind: Config
clusters:
- name: in-cluster
  cluster:
    certificate-authority: ${ca_file}
    server: ${server}
contexts:
- name: in-cluster
  context:
    cluster: in-cluster
    user: service-account
    namespace: ${namespace}
current-context: in-cluster
users:
- name: service-account
  user:
    tokenFile: ${token_file}
EOF

  chmod 0600 "${generated_kubeconfig}"
  install_kubeconfig_for_user admin "${generated_kubeconfig}"
  install_kubeconfig_for_user user "${generated_kubeconfig}"
  log "Installed generated in-cluster kubeconfig for admin and user from ${sa_dir}."
}

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 "${BOOTSTRAP_KUBECONFIG}"
  install_kubeconfig_for_user user "${BOOTSTRAP_KUBECONFIG}"
  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
Port ${SSH_PORT}
PasswordAuthentication ${password_value}
KbdInteractiveAuthentication ${keyboard_interactive_value}
EOF
}

persist_kubernetes_login_env() {
  local kubernetes_host="${KUBERNETES_SERVICE_HOST:-kubernetes.default.svc}"
  local kubernetes_https_port="${KUBERNETES_SERVICE_PORT_HTTPS:-443}"
  local kubernetes_port="${KUBERNETES_SERVICE_PORT:-${kubernetes_https_port}}"

  cat > "${K8S_RUNTIME_ENV_PROFILE}" <<EOF
# Generated at container start so SSH login shells see Kubernetes API environment values.
export KUBERNETES_SERVICE_HOST='${kubernetes_host}'
export KUBERNETES_SERVICE_PORT='${kubernetes_port}'
export KUBERNETES_SERVICE_PORT_HTTPS='${kubernetes_https_port}'
EOF

  chmod 0644 "${K8S_RUNTIME_ENV_PROFILE}"
}

prepare_runtime() {
  mkdir -p /var/run/sshd
  ensure_host_keys
  if [[ "${ENABLE_BOOTSTRAP_KUBECONFIG}" == "true" ]]; then
    log "Kubeconfig mode: bootstrap-only (serviceaccount kubeconfig generation disabled)."
    install_bootstrap_kubeconfig
  else
    log "Kubeconfig mode: in-cluster serviceaccount."
    install_serviceaccount_kubeconfig
  fi
  persist_kubernetes_login_env
  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 ${SSH_PORT}."
  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