#!/bin/bash
#
# Copyright 2025 Andrew 'Diddymus' Rolfe
# Released under the MIT license.
# SPDX-License-Identifier: MIT
#
# Setup a Kubernetes control-plane or worker node. The control-plane node is
# untainted by default so that it is also a worker node. This allows for a
# simple single node cluster if desired.  Uses flannel for pod network.
# Installs Kubernetes dashboard and a private docker registry on the
# control-plane node.
#
# This script has been tested with the following Debian Trixie net install images:
#
#   https://cdimage.debian.org/debian-cd/13.2.0/amd64/iso-cd/debian-13.2.0-amd64-netinst.iso
#
#   https://cdimage.debian.org/debian-cd/13.2.0/arm64/iso-cd/debian-13.2.0-arm64-netinst.iso
#
# All output will be  automatically logged to the file specified in $LOGFILE
# (default: ./node-setup.log).
#
# A separate log file will be created for creating the control-plane and worker
# nodes (./cluster-create.log or ./cluster-join.log).
#
# For control-plane nodes (CNODE_IP not set) the following environment
# variables can be specified:
#
#   USERNAME name of EXISTING user to setup for kubectl and SSH
#     SSH_PUB public SSH key to add for login as USERNAME
#   DOCKER_USER docker user to setup (default: admin)
#   DOCKER_PASS docker user's password to set (default: random)
#   KUBERNETES_VER version of Kubernetes to install (default: 1.34)
#   PAUSE_VER version of pause image to use (default: 3.10.1)
#   HELM_VER version of helm to install (default: 4.0.0)
#   NO_DASHBOARD if set to anything the Kubernetes dashboard is not installed
#   NODE_IP force this host to use the specified IP address (default: hostname -i)
#   CNODE_IP force control-plane to use the specified IP address (default: hostname -i)
#   ARCH architecture, one of: amd64, x86_64, aarch64, arm64 (default: unmae -m)
#
# For a control-plane node with multiple interfaces a specific IP address can
# be forced by setting both NODE_IP and CNODE_IP to the same IP address to be
# used.
#
# For worker nodes (CNODE_IP set) the following environment variables can be
# specified:
#
#   CNODE_IP [REQUIRED] the ip address of the control-plane node
#   DOCKER_PASS [REQUIRED] must be same as used or generated for control-plane
#   DOCKER_USER must match user used for control-plane if default not used
#   USERNAME name of EXISTING user to setup SSH only for
#     SSH_PUB public SSH key to add for login as USERNAME
#   NODE_IP force this host to use the specified IP address (default: hostname -i)
#   ARCH architecture for install amd64, x86_64, aarch64 or arm64 (default: unmae -m)
#
# For a worker node with multiple interfaces a specific IP address can be
# forced by setting both NODE_IP to the IP address to be used.

NODE_SETUP_VER="v0.0.5"
LOGFILE="./node-setup.log"
RND_PASS=$(openssl rand -base64 32 | tr -d '/+=' | cut -c1-20)

# Configuration that can be overridden by environment variables
USERNAME="${USERNAME}"                    # Normal user to setup
SSH_PUB="${SSH_PUB}"                      # Public ssh key for remote login as USERNAME
DOCKER_USER="${DOCKER_USER:-admin}"       # Docker registry user
DOCKER_PASS="${DOCKER_PASS:-${RND_PASS}}" # Password for docker registry user
CNODE_IP="${CNODE_IP:-$(hostname -i)}"    # IP address of control-plane node
JOIN_TOKEN="${JOIN_TOKEN}"                # Cluster joining token for worker nodes
JOIN_HASH="${JOIN_HASH}"                  # Cluster joining hash for worker nodes
HELM_VER="${HELM_VER:-4.0.0}"             # Version of Helm for dashboard install
KUBERNETES_VER="${KUBERNETES_VER:-1.34}"  # Version of kubernetes to pin install at
PAUSE_VER="${PAUSE_VER:-3.10.1}"          # Version of pause image to use
ARCH="${ARCH:-$(uname -m)}"               # Currently x86_64 and arm64 supported

# IP address of this node
NODE_IP=${NODE_IP:-$(hostname -i)}
[ "${NODE_IP}" != "${CNODE_IP}" ] || CTRL="1"

# Install Kubernetes dashboard?
[ ! -z $NO_DASHBOARD ] || DASHBOARD="1"

# Colours for nice messages
RED="\x1b[31m"
GREEN="\x1b[32m"
YELLOW="\x1b[33m"
RESET="\x1b[0m"

OK=${GREEN}
INFO=${YELLOW}
ERR=${RED}

# Logging functions
log_ok()   { echo -e "${OK}$1${RESET}"; }
log_info() { echo -e "${INFO}- $1${RESET}"; }
log_err()  { echo -e "${ERR}ERROR: $1${RESET}" >&2; }
log_exit() { log_err "$1"; exit 1; }

# Check architecture. Currently tested on amd64 and arm64.
case "${ARCH,,}" in
	amd64|x86_64)	 ARCH="amd64" ;;
	aarch64|arm64) ARCH="arm64" ;;
	*)             log_exit "Unsupported architecture '${ARCH}'" ;;
esac

# Download URLs
HELM_URL="https://get.helm.sh/helm-v${HELM_VER}-linux-${ARCH}.tar.gz"
K8S_KEY_URL="https://pkgs.k8s.io/core:/stable:/v${KUBERNETES_VER}/deb/Release.key"
FLANNEL_URL="https://github.com/flannel-io/flannel/releases/latest/download/kube-flannel.yml"
METRICS_URL="https://github.com/kubernetes-sigs/metrics-server/releases/latest/download/components.yaml"

# Wait for a Kubernetes resource to be ready
wait_for_ready() {
    local resource=$1
    local timeout=${2:-300}
    local interval=5
    local elapsed=0

    log_info "Waiting for $resource to be ready (timeout: ${timeout}s)..."
    while [ $elapsed -lt $timeout ]; do
        if kubectl wait --for=condition=ready $resource --timeout=5s &> /dev/null; then
            log_info "$resource is ready"
            return 0
        fi
        echo "  still waiting..."
        sleep $interval
        elapsed=$((elapsed + interval))
    done
    log_exit "$resource failed to become ready within ${timeout}s"
}

pre_checks() {
  log_ok "Running node-setup ${NODE_SETUP_VER}"
  log_ok "Architecture is ${ARCH}"
  log_ok "Logging to ${LOGFILE}"

  # Check if running as root
  if [ "$EUID" -ne 0 ]; then
     log_exit "Please run as root"
  fi

  # If not control-plane then JOIN_TOKEN and JOIN_HASH must be set
  if [ ! $CTRL ]; then
    if [ -z $JOIN_TOKEN ]; then
      log_exit "JOIN_TOKEN must be set for worker node setup."
    fi
    if [ -z $JOIN_HASH ]; then
      log_exit "JOIN_HASH must be set for worker node setup."
    fi
  fi

  # For worker nodes make sure we can ping the control-plane node
  if [ ! $CTRL ]; then
    log_ok "Checking network for control-panel node connectivity..."
    if ping -c 1 ${CNODE_IP} &> /dev/null; then
      log_info "Contacted control-panel node as ${CNODE_IP} successfully."
    else
      log_exit "Cannot ping control-panel node as ${CNODE_IP}"
    fi
  fi

  # Setting up control-plane or worker node?
  if [ $CTRL ]; then
    log_ok "Setting up ${HOSTNAME} with control-plane and docker."
  else
    log_ok "Setting up ${HOSTNAME} as worker node."
  fi
}

setup_ssh() {
  log_ok "Setting up remote SSH access"

  if [ -z ${USERNAME} ]; then
    log_info "USERNAME not set, aborting SSH setup."
    return
  fi

  if [ -z "${SSH_PUB}" ]; then
    log_info "SSH_PUB not set, aborting SSH setup."
    return
  fi

  if id "${USERNAME}" &> /dev/null; then
    log_info "User ${USERNAME} exists."
  else
    log_info "User ${USERNAME} does not exist, aborting SSH setup."
    return
  fi

  local ssh_dir="/home/${USERNAME}/.ssh"

  # Check for .ssh
  if [ -d ${ssh_dir} ]; then
    log_info "SSH dir ${ssh_dir} already exists"
  else
    log_info "creating SSH dir ${ssh_dir}"
    mkdir -p ${ssh_dir}
    chmod 0700 ${ssh_dir}
    chown ${USERNAME}:${USERNAME} ${ssh_dir}
  fi

  # Check for .ssh/authorized_keys
  if [ -f ${ssh_dir}/authorized_keys ]; then
    log_info "SSH authorized keys file ${ssh_dir}/authorized_keys already exists"

    # Check for public key
    if grep -q "${SSH_PUB}" ${ssh_dir}/authorized_keys; then
      log_info "SSH authorized keys file already contains piblic key"
    else
      log_info "adding public key to ${ssh_dir}/authorized_keys"
      echo "${SSH_PUB}" >> /home/${USERNAME}/.ssh/authorized_keys
    fi
  else
    log_info "creating SSH authorized keys file ${ssh_dir}/authorized_keys"
    echo "${SSH_PUB}" >> ${ssh_dir}/authorized_keys
    chmod 0600 ${ssh_dir}/authorized_keys
    chown ${USERNAME}:${USERNAME} ${ssh_dir}/authorized_keys
  fi

}

# Install a package if it is missing
install_pkg() {
  local name=$1

  if dpkg-query -s ${name} &> /dev/null ; then
    log_info "package ${name} already installed"
  else
    log_info "installing package ${name}"
    apt-get install -q -q -y --no-install-recommends ${name} &> /dev/null || log_err "failed to install ${name}"
  fi
}

# Install needed packages if not detected
install_packages() {
  log_ok "Checking packages"

  install_pkg ed
  install_pkg jq
  install_pkg docker.io
  install_pkg docker-buildx
  install_pkg docker-cli
  install_pkg docker-registry
  install_pkg apache2-utils
  install_pkg curl
  install_pkg gpg
  install_pkg ca-certificates
  install_pkg apt-transport-https
}

# Install Helm if not detected
install_helm() {
  log_ok "Checking if helm installed"

  if [ ! -z `command -v helm` ]; then
    log_info "found helm $(helm version --short)"
    return
  fi

  log_info "installing helm v${HELM_VER}"

  [ -d ./downloads ] || mkdir ./downloads
  cd downloads
  curl -sO ${HELM_URL}
  tar -zxf helm-v${HELM_VER}-linux-${ARCH}.tar.gz
  cp linux-${ARCH}/helm /usr/local/bin/helm
  cd ..

  log_info "installed helm $(helm version --short)"
}

load_modules() {
  log_ok "Loading kernel modules"

  local config="/etc/modules-load.d/containerd.conf"

  if [ -f ${config} ]; then
    log_info "module config file ${config} exists"
    return
  fi

  log_info "creating module config file ${config}"
  cat > ${config} <<EOT
overlay
br_netfilter
EOT

  modprobe overlay
  modprobe br_netfilter
}

# Configure network bridge
config_networking() {
  log_ok "Configuring network"

  local config="/etc/sysctl.d/99-kubernetes-k8s.conf"

  if [ -f ${config} ]; then
    log_info "network configuration already exists ${config}"
    return
  fi

  log_info "creating network configuration ${config}"
  cat > ${config} <<EOT
net.bridge.bridge-nf-call-iptables = 1
net.ipv4.ip_forward = 1
net.bridge.bridge-nf-call-ip6tables = 1
EOT

  log_ok "Reloading configuration"
  sysctl -q --system
}

setup_containerd() {
  log_ok "Setting up containerd"

  install_pkg containerd

  local config="/etc/containerd/config.toml"

  if [ -f ${config} ]; then
    log_info "backing up existing ${config}"
    cp ${config} ${config}.bak
  fi

  log_info "creating new ${config}"
  containerd config default > ${config}

  ed -qs /etc/containerd/config.toml <<EOT
1,\$s/SystemdCgroup = false/SystemdCgroup = true
1,\$s/pause:.*"/pause:${PAUSE_VER}"
w
q
EOT

  log_ok "Starting containerd"
  systemctl enable --now containerd
  systemctl restart containerd
}

# Install kubeadmin, kubelet and kubectl
install_kubernetes() {
  log_ok "Installing kubernetes v${KUBERNETES_VER}"

  local keyring="/etc/apt/keyrings/kubernetes-apt-keyring.gpg"

  if [ -f ${keyring} ]; then
    log_info "keyring exists ${keyring}"
  else
    log_info "creating keyring ${keyring}"

    [ -d ./downloads ] || mkdir ./downloads
    cd downloads
    curl -fsSLO ${K8S_KEY_URL}
    gpg --dearmor -o ${keyring} ./Release.key
    echo "deb [signed-by=${keyring}] https://pkgs.k8s.io/core:/stable:/v${KUBERNETES_VER}/deb/ /" > /etc/apt/sources.list.d/kubernetes.list
    cd ..
  fi

  log_ok "Installing kubernetes packages"
  apt-get update -q -q &> /dev/null
  install_pkg kubelet
  install_pkg kubeadm
  install_pkg kubectl
  apt-mark hold kubelet kubeadm kubectl

  log_ok "Starting kubelet"
  systemctl enable --now kubelet
  systemctl restart kubelet
}

# Symlink cni drivers
symlink_cni_drivers() {
  log_ok "Symlinking CNI drivers"

  if [ -h /usr/lib/cni ]; then
    log_info "CNI drivers already symlinked /usr/lib/cni"
    return
  fi

  log_info "creating CNI driver symlink /usr/lib/cni"
  ln -s /opt/cni/bin /usr/lib/cni
}

# Create cluster
create_cluster() {
  log_ok "Creating cluster ${HOSTNAME}, please wait..."
  kubeadm init --pod-network-cidr=10.244.0.0/16 --control-plane-endpoint=${HOSTNAME} &> ./cluster-create.log
}

# Join cluster
join_cluster() {
  log_ok "Joining cluster ${CNODE_IP}"
  kubeadm join --token ${JOIN_TOKEN} ${CNODE_IP}:6443 --discovery-token-ca-cert-hash sha256:${JOIN_HASH} &> ./cluster-join.log
}

# Configure kubectl for local user
config_kubectl() {
  log_ok "Configuring kubectl for local user"

  # Export kubernetes for root to use in script
  export KUBECONFIG=/etc/kubernetes/admin.conf

  if [ -z $USERNAME ]; then
    log_info "USERNAME not set, aborting kubectl setup."
    return
  fi

  local kube_dir="/home/${USERNAME}/.kube"

  if [ -d ${kube_dir} ]; then
    log_info "directory exists ${kube_dir}"
  else
    log_info "creating directory ${kube_dir}"
    mkdir -p ${kube_dir}
    chown ${USERNAME}:${USERNAME} ${kube_dir}
  fi

  if [ -d ${kube_dir}/config ]; then
    log_info "backing up existing file ${kube_dir}/config"
    cp ${kube_dir}/config ${kube_dir}/config.bak
  fi

  log_info "creating file ${kube_dir}/config"
  cp /etc/kubernetes/admin.conf ${kube_dir}/config
  chown ${USERNAME}:${USERNAME} ${kube_dir}/config
}

# Install pod network
install_pod_network() {
  log_ok "Installing pod network"

  [ -d ./downloads ] || mkdir ./downloads
  cd downloads
  curl -LsO ${FLANNEL_URL}
  kubectl create -f ./kube-flannel.yml
  cd ..

  wait_for_ready "pod -n kube-flannel -l app=flannel" 120
}

# Untaint control-plane node so that it is also a worker node.
untaint_control_plane() {
  log_ok "Untainting control-plane node to be a worker node"
  kubectl taint node ${HOSTNAME} node-role.kubernetes.io/control-plane:NoSchedule-
}

# Install metrics server
install_metrics_server() {
  log_ok "Installing metrics server"

  [ -d ./downloads ] || mkdir ./downloads
  cd downloads
  curl -LsO ${METRICS_URL}

  ed -qs ./components.yaml <<EOT
/--metric-resolution
a
        - --kubelet-insecure-tls=true
        - --kubelet-preferred-address-types=InternalIP
.
w
q
EOT

  kubectl create -f ./components.yaml
  cd ..

  wait_for_ready "pod -n kube-system -l k8s-app=metrics-server" 120
}

# Install kubernetes dashboard
install_dashboard() {
  log_ok "Installing kubernetes dashboard"

  helm repo add kubernetes-dashboard https://kubernetes-retired.github.io/dashboard/ &> /dev/null
  helm upgrade --install kubernetes-dashboard kubernetes-dashboard/kubernetes-dashboard --create-namespace --namespace kubernetes-dashboard &> /dev/null

  log_ok "Creating kubernetes dashboard service account"

  [ -d ./downloads ] || mkdir ./downloads
  cd downloads
  cat > ./dashboard-service-account.yaml <<EOT
apiVersion: v1
kind: ServiceAccount
metadata:
  name: admin-user
  namespace: kubernetes-dashboard
---
apiVersion: rbac.authorization.k8s.io/v1
kind: ClusterRoleBinding
metadata:
  name: admin-user
roleRef:
  apiGroup: rbac.authorization.k8s.io
  kind: ClusterRole
  name: cluster-admin
subjects:
- kind: ServiceAccount
  name: admin-user
  namespace: kubernetes-dashboard
---
apiVersion: v1
kind: Secret
metadata:
  name: admin-user
  namespace: kubernetes-dashboard
  annotations:
    kubernetes.io/service-account.name: "admin-user"
type: kubernetes.io/service-account-token
EOT

  kubectl create -f ./dashboard-service-account.yaml -n kubernetes-dashboard
  cd ..

  wait_for_ready "pod -n kubernetes-dashboard -l app=kubernetes-dashboard-kong" 300
}

# Configure docker registry and containerd to use it
config_registry() {
  log_ok "Configuring docker registry"

  local config="/etc/docker/daemon.json"

  if [ -f ${config} ]; then
    log_info "docker config exists ${config}"
  else
    log_info "creating docker config ${config}"

    cat > ${config} <<EOT
{
  "insecure-registries" : [ "${CNODE_IP}:5000" ]
}
EOT

    log_ok "Restarting docker, docker-registry and containerd services"
    systemctl restart docker
    systemctl restart docker-registry
    systemctl restart containerd
  fi

  local htpasswd="/etc/docker/registry/htpasswd"

  log_ok "Creating docker registry user ${DOCKER_USER}"

  if [ -s ${htpasswd} ]; then
    log_info "file already exists ${htpasswd}"
  else
    log_info "creating ${htpasswd} for ${DOCKER_USER}"
    mkdir -p $(dirname ${htpasswd})
    htpasswd -Bcb ${htpasswd} ${DOCKER_USER} ${DOCKER_PASS}
    if id "${DOCKER_USER}" &> /dev/null; then
      log_info "user ${DOCKER_USER} added to docker group"
      usermod -a -G docker ${DOCKER_USER}
    fi
  fi

  log_ok "Restarting docker, docker-registry and containerd services"
  systemctl restart docker
  systemctl restart docker-registry
  systemctl restart containerd

  # Wait for services to settle
  sleep 10s

  log_ok "Configuring containerd to use docker registry"

  log_info "logging into docker registry ${CNODE_IP}:5000"
  docker login -u${DOCKER_USER} -p${DOCKER_PASS} ${CNODE_IP}:5000 \
    || log_exit "Failed to login to docker registry ${CNODE_IP}:5000 as ${DOCKER_USER}"

  local cert="/etc/containerd/certs.d/${CNODE_IP}:5000"
  local token=$(jq -r ".auths[\"${CNODE_IP}:5000\"].auth" ~/.docker/config.json)

  log_info "creating containerd certificate ${cert}/hosts.toml"
  log_info "token for ${DOCKER_USER} is ${token}"
  mkdir -p ${cert}

  cat > ${cert}/hosts.toml <<EOT
server = "http://${CNODE_IP}:5000"

[host."http://${CNODE_IP}:5000"]
  capabilities = ["pull", "resolve"]
  skip_verify = true
  [host."http://${CNODE_IP}:5000".header]
    authorization = "Basic ${token}"
EOT

  log_info "updating containerd config /etc/containerd/config.toml"
  ed -qs /etc/containerd/config.toml <<EOT
/cri"\.registry
+1
s/""/"\/etc\/containerd\/certs.d"/
w
q
EOT

  log_ok "Restarting containerd service"
  systemctl restart containerd
}

# Expose kubernetes dashboard to network on port 30443
expose_dashboard() {
  log_ok "Exposing kubernetes dashboard to network"

  local namespace="kubernetes-dashboard"
  local proxy_svc="kubernetes-dashboard-kong-proxy"

  kubectl -n ${namespace} patch svc ${proxy_svc} -p '{"spec":{"type":"NodePort"}}'
  kubectl -n ${namespace} patch svc ${proxy_svc} -p '{"spec":{"ports":[{"nodePort":30443, "port":443, "protocol":"TCP", "targetPort":8443, "name":"kong-proxy-tls"}]}}'
}

display_setup() {
  echo   ""
  log_ok "+++ Node ${HOSTNAME} Setup Complete +++"

  [ $CTRL ] || return

  local dashboard_token=$(kubectl get secret admin-user -n kubernetes-dashboard -o jsonpath="{.data.token}" | base64 -d)
  local join_token=$(kubeadm token list -o jsonpath="{.token}")
  local join_hash=$(cat /etc/kubernetes/pki/ca.crt | openssl x509 -pubkey  | openssl rsa -pubin -outform der 2>/dev/null | openssl dgst -sha256 -hex | sed 's/^.* //')

  echo   ""
  log_ok "=== Services ==="
  kubectl get services --all-namespaces
  echo   ""
  log_ok "=== Pods ==="
  kubectl get pods --all-namespaces
  echo   ""
  log_ok "=== Access Information ==="
  if [ $DASHBOARD ]; then
    echo   "Dashboard URL............: https://${CNODE_IP}:30443"
    echo   "Dashboard token..........: ${dashboard_token}"
  else
    echo   "Dashboard URL............: NOT INSTALLED"
  fi
  echo   "Docker registry URL......: ${CNODE_IP}:5000"
  echo   "Docker registry user.....: ${DOCKER_USER}"
  echo   "Docker registry password.: ${DOCKER_PASS}"
  echo   "Cluster join token.......: ${join_token}"
  echo   "Cluster join hash........: ${join_hash}"
  echo   ""
  log_ok "=== kubectl admin.conf ==="
  cat  /etc/kubernetes/admin.conf
  echo   ""
}

# Main routine
{
  pre_checks
  setup_ssh
  install_packages
  load_modules
  config_networking
  setup_containerd
  install_kubernetes
  symlink_cni_drivers
  [ $CTRL ] && create_cluster
  [ $CTRL ] || join_cluster
  [ $CTRL ] && config_kubectl
  [ $CTRL ] && install_pod_network
  [ $CTRL ] && untaint_control_plane
  [ $CTRL ] && install_metrics_server
  config_registry
  [ $CTRL ] && [ $DASHBOARD ] && install_helm
  [ $CTRL ] && [ $DASHBOARD ] && install_dashboard
  [ $CTRL ] && [ $DASHBOARD ] && expose_dashboard
  display_setup
} 2>&1 | tee ${LOGFILE}