os-dependent/ubuntu/scripts/satlab_install

#!/bin/bash
# Copyright 2025 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

# This script configures a Satlab drone environment.
# It includes robust error handling to ensure a safe and reliable setup.

# Exit immediately if a command exits with a non-zero status.
# Treat unset variables as an error.
# The return value of a pipeline is the status of the last command to exit with a non-zero status.
set -euo pipefail

# --- UI Colors ---
C_YELLOW='\033[1;33m'
C_GREEN='\033[1;32m'
C_RED='\033[0;31m'
B_WHITE='\033[1;37m'
C_NC='\033[0m' # No Color

# --- Default channel ---
CHANNEL="release"

# --- Argument Parsing ---
for arg in "$@"; do
  case $arg in
    --beta)
      CHANNEL="beta"
      shift
      ;;
    *)
      # Ignore unknown arguments or handle them here
      ;;
  esac
done

##
# Centralized error handler. It's triggered by the 'trap' command below.
##
error_handler() {
  local line_number=$1
  local failed_command=$2
  echo -e "\n${C_RED}Error on line ${line_number}: Command failed: '${failed_command}'${C_NC}"
  echo -e "${C_RED}Aborting script execution.${C_NC}"
}

# Trap errors and call the error_handler function
trap 'error_handler ${LINENO} "$BASH_COMMAND"' ERR

# --- Configuration and Function Sourcing ---

# Define the CONFIG path early so it's available to sourced files
CONFIG="/etc/satlab/satlab.conf"

# Attempt to source the config file to load the set_config function and variables.
if [[ -f "$CONFIG" ]]; then
    source "$CONFIG"
fi

# Verify set_config was loaded from the file; if not, define it as a fallback.
# This ensures the script works on a fresh system.
if ! declare -F "set_config" > /dev/null; then
    echo -e "${C_YELLOW}Info: 'set_config' function not found in config file. Using fallback definition.${C_NC}"
    function set_config(){
        local key=$1
        local value=$2
        # Use the global CONFIG variable
        mkdir -p "$(dirname "$CONFIG")"
        touch "$CONFIG"
        if grep -q "^${key}=" "$CONFIG"; then
            # Use | as a delimiter for sed to avoid issues with paths
            sed -i "s|^\(${key}\s*=\s*\).*$|\1${value}|" "$CONFIG"
        else
            echo "${key}=${value}" >> "$CONFIG"
        fi
    }
fi

# --- Pre-flight Checks ---

# Check if running as root
if [[ $EUID -ne 0 ]]; then
   echo -e "${C_RED}This script must be run as root. Please use sudo.${C_NC}"
   exit 1
fi

# Check if running in SSH (optional)
if [[ -z ${SSH_TTY:-} ]]; then
    echo -e "${C_YELLOW}It's recommended to run this script within an SSH session.${C_NC}"
    while true; do
        read -p "Do you want to continue (y/n)?" yn
        case $yn in
            [Yy]* ) break;;
            [Nn]* ) exit;;
            * ) echo -e "${C_YELLOW}Please answer yes or no.${C_NC}";;
        esac
    done
fi

# --- Kernel Version Check ---

KERNEL_FULL=$(uname -r)
KERNEL_MAJOR=$(echo "$KERNEL_FULL" | cut -d. -f1)
KERNEL_MINOR=$(echo "$KERNEL_FULL" | cut -d. -f2)

echo -e "${C_YELLOW}Checking kernel compatibility...${C_NC}"

if [[ "$KERNEL_MAJOR" -gt 6 ]] || [[ "$KERNEL_MAJOR" -eq 6 && "$KERNEL_MINOR" -gt 13 ]]; then
    echo -e "${C_RED}Error: Incompatible kernel version found: ($KERNEL_FULL).${C_NC}"
    echo -e "${C_RED}Satlab requires kernel version 6.13 or lower.${C_NC}"
    echo "Aborting installation."
    exit 1
else
    echo -e "${C_GREEN}Kernel version $KERNEL_FULL is compatible.${C_NC}"
fi

# --- Install Docker ---

TARGET_DOCKER_VERSION="28.3.3"
echo -e "\n${C_YELLOW}Checking Docker installation...${C_NC}"

INSTALL_DOCKER=false
# Safely check if docker exists and is executable
DOCKER_BIN=$(command -v docker 2>/dev/null || true)
if [[ -n "$DOCKER_BIN" && -x "$DOCKER_BIN" ]]; then
    CURRENT_DOCKER_VERSION=$("$DOCKER_BIN" --version 2>/dev/null | awk '{print $3}' | tr -d ',' || true)
    # Check if the current version matches the target version
    if [[ -n "$CURRENT_DOCKER_VERSION" && "$CURRENT_DOCKER_VERSION" == "$TARGET_DOCKER_VERSION"* ]]; then
        echo -e "${C_GREEN}Docker version $TARGET_DOCKER_VERSION is already installed. Skipping installation.${C_NC}"
    else
        echo -e "${C_RED}Error: Found Docker version ${CURRENT_DOCKER_VERSION:-Unknown}, but target version is $TARGET_DOCKER_VERSION.${C_NC}"
        echo -e "${C_YELLOW}Please uninstall the current Docker version manually before continuing with this installation.${C_NC}"
        exit 1
    fi
else
    echo -e "${C_YELLOW}Docker is not installed. Installing version $TARGET_DOCKER_VERSION...${C_NC}"
    INSTALL_DOCKER=true
fi

if [[ "$INSTALL_DOCKER" == true ]]; then
    curl -fsSL https://get.docker.com -o install-docker.sh
    sh install-docker.sh --version "$TARGET_DOCKER_VERSION"
    rm -f install-docker.sh
    echo -e "${C_GREEN}Docker $TARGET_DOCKER_VERSION installed successfully.${C_NC}"
fi

# --- Configure swap memory ---

SWAP_FILE="/swap.img"
DESIRED_GB=32
# Calculate KiB for assertion (32 * 1024 * 1024).
# We'll check if the reported size is "close enough" (within 1% of this value).
TARGET_KIB=$(( DESIRED_GB * 1024 * 1024 ))

echo -e "\n${C_YELLOW}Setting swap memory size to ${DESIRED_GB}GB...${C_NC}"
# Disable current swap if it's active (ignore errors if it isn't)
swapoff "$SWAP_FILE" 2>/dev/null || true
# Resize the file. Using 'fallocate' is fast; if it fails, 'dd' is a fallback.
if ! fallocate -l "${DESIRED_GB}G" "$SWAP_FILE"; then
    echo -e "${C_YELLOW}fallocate failed, attempting dd...${C_NC}"
    dd if=/dev/zero of="$SWAP_FILE" bs=1M count=$(( DESIRED_GB * 1024 ))
fi
# Secure the file permissions (critical for swap security)
chmod 600 "$SWAP_FILE"
# Format and enable
mkswap "$SWAP_FILE"
swapon "$SWAP_FILE"

# Assertion Logic
ACTUAL_KIB=$(awk -v f="$SWAP_FILE" '$1==f {print $3}' /proc/swaps)
if [[ -n "$ACTUAL_KIB" ]]; then
    # Calculate a 1% margin of error for filesystem overhead
    MARGIN=$(( TARGET_KIB / 100 ))
    DIFF=$(( TARGET_KIB - ACTUAL_KIB ))
    # Check if the absolute difference is within the margin
    if [ "${DIFF#-}" -lt "$MARGIN" ]; then
        echo -e "${C_GREEN}Swap resized successfully to ~${DESIRED_GB}GB (${ACTUAL_KIB} KiB).${C_NC}"
    else
        echo -e "${C_RED}Assertion Failed: Swap size is ${ACTUAL_KIB} KiB, expected ~${TARGET_KIB} KiB.${C_NC}"
        exit 1
    fi
else
    echo -e "${C_RED}Error: Swap file $SWAP_FILE is not active.${C_NC}"
    exit 1
fi

echo -e "Using software channel: ${C_YELLOW}${CHANNEL}${C_NC}"

# --- Configure packages ---

echo -e "\n${C_YELLOW}Configuring packages...${C_NC}"
apt-get update -y
apt-get install -y bridge-utils net-tools python3-grpcio python3-netifaces python3-protobuf
systemctl stop unattended-upgrades || true
systemctl disable unattended-upgrades || true
apt-get remove --purge -y unattended-upgrades
apt-mark hold linux-image-generic \
  docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin docker-ce-rootless-extras

# --- Disable auto-mounting partitions (to prevent blockdev busy errors while writing to a USB drive) ---

systemctl stop udisks2.service
systemctl mask udisks2.service

# --- Optional: Install Graphical Environment (Xfce) ---

echo -e "\n${B_WHITE}Optional: Install desktop environment${C_NC}"
read -p "Do you want to install the Xfce desktop environment? (y/N): " install_gui

if [[ "${install_gui,,}" =~ ^(y|yes)$ ]]; then
    echo -e "${C_YELLOW}Installing Xfce and LightDM (this may take a few minutes)...${C_NC}"
    # Pre-configure LightDM to avoid interactive debconf prompts during install
    echo "lightdm shared/default-x-display-manager select lightdm" | debconf-set-selections
    # Install Xfce, the X server, and the display manager
    # --no-install-recommends keeps the installation lightweight
    apt-get install -y xfce4 xfce4-goodies xfce4-session xorg lightdm lightdm-gtk-greeter dbus-x11 chromium-browser --no-install-recommends
    # Force Xfce in AccountsService so it doesn't look for 'ubuntu'
    if [ -f /var/lib/AccountsService/users/satlab ]; then
        sed -i 's/Session=.*/Session=xfce/' /var/lib/AccountsService/users/satlab
    else
        # If the file doesn't exist, create it with xfce as default
        mkdir -p /var/lib/AccountsService/users/
        echo -e "[User]\nSession=xfce" > /var/lib/AccountsService/users/satlab
    fi
    # Create the .dmrc for the user
    echo -e "[Desktop]\nSession=xfce" > /home/satlab/.dmrc
    chown satlab:satlab /home/satlab/.dmrc
    # System-wide override
    mkdir -p /etc/lightdm/lightdm.conf.d/
    echo -e "[Seat:*]\nuser-session=xfce" > /etc/lightdm/lightdm.conf.d/90-xfce.conf
    # Set the system to boot into the graphical target
    systemctl set-default graphical.target
    echo -e "${C_GREEN}Xfce Desktop Environment installed successfully.${C_NC}"
else
    echo -e "Skipping desktop environment installation."
fi

# --- Configure the host ---

echo -e "\n${C_YELLOW}Sourcing local configuration scripts...${C_NC}"
nic_config_script="$(dirname "$0")/satlab_config_nics"
if [[ ! -f "$nic_config_script" ]]; then
    echo -e "${C_RED}Error: Required scripts (satlab_config_nics, satlab_get_keys) not found in the same directory.${C_NC}"
    exit 1
fi
source "$nic_config_script"

echo -e "\n${C_YELLOW}Configuring SSH server...${C_NC}"
sed -i -E 's/^#?Port\s+22\b/Port 2233/' /etc/ssh/sshd_config

echo -e "\n${C_YELLOW}Configuring udev rules for persistent peripheral access...${C_NC}"
cat > /etc/udev/rules.d/99-satlab-peripherals.rules <<'EOF'
# Persistent access for serial devices
KERNEL=="ttyACM*", MODE="0666"
KERNEL=="ttyUSB*", MODE="0666"

# Persistent access for video peripheral devices
KERNEL=="video*", MODE="0666"
EOF

if command -v udevadm >/dev/null 2>&1; then
    udevadm control --reload-rules || echo "Warning: Failed to reload udev rules"
    udevadm trigger || echo "Warning: Failed to trigger udev events"
else
    echo "Warning: udevadm not found, skipping rule reloading"
fi

echo -e "\n${C_YELLOW}Configuring the Docker socket...${C_NC}"
mkdir -p /etc/systemd/system/docker.service.d/
cat > /etc/systemd/system/docker.service.d/override.conf <<'EOT'
[Service]
ExecStart=
ExecStart=/usr/bin/dockerd -H fd:// -H unix:///var/run/docker.sock -H tcp://0.0.0.0:2375
EOT

echo -e "\n${C_YELLOW}Configuring the Satlab hostname...${C_NC}"
if [[ ! -f /etc/machine-id ]]; then
    echo -e "${C_RED}Error: /etc/machine-id not found. Cannot generate hostname.${C_NC}"
    exit 1
fi
host_identifier="5w$(cat /etc/machine-id | sha256sum | head -c 10)"
echo "Setting Satlab hostname: ${host_identifier}"
set_config "HOST_IDENTIFIER" "$host_identifier"

# --- Download images ---

echo -e "\n${C_YELLOW}Pulling Docker images...${C_NC}"
docker pull ${REMOTE_ACCESS_IMAGE}
docker pull ${DHCP_DOCKER_IMAGE}
docker pull ${COMPOSE_DOCKER_IMAGE}

if [[ "${download_keys:-n}" =~ ^[yY]([eE][sS])?$ ]]; then
    echo -e "\n${C_YELLOW}Pulling images via docker-compose...${C_NC}"
    docker rm compose-pull 2> /dev/null || true
    docker run -it --rm --name compose-pull \
        --mount type=volume,source=satlab_keys,target=/home/satlab/keys/ \
        --mount type=bind,source=/var/run/docker.sock,target=/var/run/docker.sock \
        -e LABEL="${LABEL}" \
        -e VOLUME_KEYS_FOLDER="${VOLUME_KEYS_FOLDER}" \
        "${COMPOSE_DOCKER_IMAGE}" pull
    echo -e "${C_GREEN}Docker images pulled successfully.${C_NC}"
else
    echo -e "\n${C_YELLOW}Skipping Docker image pull (partner user).${C_NC}"
fi

# --- Switch to cgroups v1 ---

echo -e "\n${C_YELLOW}Configuring cgroups for legacy compatibility...${C_NC}"
if ! grep -q "systemd.unified_cgroup_hierarchy=0" /etc/default/grub; then
    echo "Adding cgroup kernel parameter to GRUB..."
    sed -i 's/\(GRUB_CMDLINE_LINUX_DEFAULT=".*\)"/\1 systemd.unified_cgroup_hierarchy=0"/' /etc/default/grub
    update-grub
else
    echo "GRUB cgroup parameter already set. Skipping."
fi

# --- Configure Satlab services ---

echo -e "\n${C_YELLOW}Configuring Satlab services...${C_NC}"
ln -sf /etc/satlab/satlab-upstart.service /etc/systemd/system/satlab-upstart.service
ln -sf /etc/satlab/host_server/satlab-host-server.service /etc/systemd/system/satlab-host-server.service
ln -sf /etc/satlab/satlab-apply-staged-updates.service /etc/systemd/system/
ln -sf /etc/satlab/satlab-stage-updates.service /etc/systemd/system/
ln -sf /etc/satlab/satlab-stage-updates.timer /etc/systemd/system/
systemctl daemon-reload
systemctl enable satlab-upstart
systemctl enable satlab-host-server
systemctl enable satlab-apply-staged-updates.service
systemctl enable satlab-stage-updates.service
systemctl enable satlab-stage-updates.timer

# --- Finish the installation ---

echo -e "\n${C_GREEN}Configuration complete! A reboot is required to apply all changes.${C_NC}"
while true; do
    read -p "Do you want to reboot now? (y/n) " yn
    case $yn in
        [Yy]* ) reboot; break;;
        [Nn]* ) echo -e "${C_GREEN}Reboot cancelled.${C_NC}"; break;;
        * ) echo -e "${C_GREEN}Please answer yes or no.${C_NC}";;
    esac
done