k3s-setup.sh

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

usage() {
    echo "Usage: $0 --rpi-ips <ip1,ip2,...> --rpi-user <username> --node-name-scheme <scheme> --sleep-duration <seconds>"
    exit 1
}

if [ "$#" -lt 8 ]; then
    usage
fi

while [ "$#" -gt 0 ]; do
    case "$1" in
        --rpi-ips)
            RPI_IPS=(${2//,/ })
            shift 2
            ;;
        --rpi-user)
            RPI_USER="$2"
            shift 2
            ;;
        --node-name-scheme)
            NODE_NAME_SCHEME="$2"
            shift 2
            ;;
        --sleep-duration)
            SLEEP_DURATION="$2"
            shift 2
            ;;
        *)
            usage
            ;;
    esac
done

HOSTNAME_SUFFIX=".local"

enable_container_features() {
    for ip in "${RPI_IPS[@]}"; do
        echo "Enabling container features on $ip"
        ssh "$RPI_USER@$ip" "sudo sh -c 'grep -qxF \"cgroup_enable=cpuset cgroup_memory=1 cgroup_enable=memory\" /boot/firmware/cmdline.txt || sed -i \"s/\$/ cgroup_enable=cpuset cgroup_memory=1 cgroup_enable=memory/\" /boot/firmware/cmdline.txt'"
    done
}

setup_sudo_for_user() {
    for ip in "${RPI_IPS[@]}"; do
        echo "Setting up sudo permissions for $RPI_USER on $ip"
        ssh "$RPI_USER@$ip" "echo '$RPI_USER ALL=(ALL) NOPASSWD: ALL' | sudo tee /etc/sudoers.d/k3sup-$RPI_USER"
        ssh "$RPI_USER@$ip" "sudo chmod 0440 /etc/sudoers.d/k3sup-$RPI_USER"
    done
}

set_unique_hostnames() {
    counter=1
    for ip in "${RPI_IPS[@]}"; do
        new_hostname="${NODE_NAME_SCHEME}$counter${HOSTNAME_SUFFIX}"
        echo "Setting hostname to $new_hostname on $ip"
        ssh "$RPI_USER@$ip" "sudo hostnamectl set-hostname $new_hostname"
        ((counter++))
    done
}

reboot_raspberry_pis() {
    for ip in "${RPI_IPS[@]}"; do
        echo "Rebooting $ip"
        ssh "$RPI_USER@$ip" "sudo reboot"
    done
}

check_container_features() {
    all_enabled=true
    for ip in "${RPI_IPS[@]}"; do
        echo "Checking container features on $ip"
        if ssh "$RPI_USER@$ip" "grep -q 'cgroup_enable=cpuset cgroup_memory=1 cgroup_enable=memory' /boot/firmware/cmdline.txt"; then
            echo "Container features are enabled on $ip"
        else
            echo "Container features are NOT enabled on $ip"
            all_enabled=false
        fi
    done
    if [ "$all_enabled" = false ]; then
        return 1
    fi
    return 0
}

install_k3s_server() {
    RPI_1_IP="${RPI_IPS[0]}"
    echo "Installing k3s on server node $RPI_1_IP"
    ssh "$RPI_USER@$RPI_1_IP" "curl -sfL https://get.k3s.io | sh -"
    ssh "$RPI_USER@$RPI_1_IP" "sudo systemctl status k3s"
    echo "Waiting for k3s to start..."
    sleep 30
    echo "Fetching node token..."
    ssh "$RPI_USER@$RPI_1_IP" "sudo cat /var/lib/rancher/k3s/server/node-token"
}

bootstrap_k3s() {
    RPI_1_IP="${RPI_IPS[0]}"
    k3sup install --ip "$RPI_1_IP" --user "$RPI_USER" --ssh-key ~/.ssh/id_ed25519 --no-extras
    for ip in "${RPI_IPS[@]:1}"; do
        echo "Joining worker node $ip to the cluster"
        k3sup join --ip "$ip" --server-ip "$RPI_1_IP" --user "$RPI_USER" --ssh-key ~/.ssh/id_ed25519
    done
}

enable_container_features
setup_sudo_for_user
set_unique_hostnames
reboot_raspberry_pis

echo "Waiting for Raspberry Pis to reboot..."
sleep "$SLEEP_DURATION"

if check_container_features; then
    install_k3s_server
    bootstrap_k3s
else
    echo "Container features are not enabled on all devices. Aborting k3s installation."
    exit 1
fi