Enhance UX and add error handling to the system setup script.

setup.sh

@@ -8,9 +8,16 @@ timezone="America/New_York"
 sshd_config="/etc/ssh/sshd_config"
 docker_installed=false
 
-# Function to prompt for hostname
+# Function to prompt for hostname with validation
 prompt_for_hostname() {
-    read -p "Please enter the desired hostname: " hostname
+    while true; do
+        read -p "Please enter the desired hostname: " hostname
+        if [[ -z "$hostname" ]]; then
+            echo "Hostname cannot be empty. Please try again."
+        else
+            break
+        fi
+    done
 }
 
 # Function to add hostname and IP to /etc/hosts
@@ -18,7 +25,7 @@ update_hosts_file() {
     local ip_address
     ip_address=$(hostname -I | awk '{print $1}')
     echo "Updating /etc/hosts with IP $ip_address and hostname $hostname..."
-    
+
     # Check if the entry already exists
     if grep -q "$ip_address" /etc/hosts; then
         echo "Entry for $ip_address already exists in /etc/hosts."
@@ -29,22 +36,33 @@ update_hosts_file() {
     fi
 }
 
-# Function to prompt for Docker installation
+# Function to prompt for Docker installation with validation
 prompt_for_docker_install() {
-    read -p "Would you like to install Docker CE? (y/n): " install_docker
+    while true; do
-    if [[ "$install_docker" =~ ^[Yy]$ ]]; then
+        read -p "Would you like to install Docker CE? (y/n): " install_docker
-        install_docker
+        case "$install_docker" in
-    else
+            [Yy]*) install_docker ;;
-        echo "Skipping Docker installation."
+            [Nn]*) echo "Skipping Docker installation."; break ;;
-    fi
+            *) echo "Invalid input. Please enter 'y' or 'n'." ;;
+        esac
+    done
 }
 
 # Function to install Docker
 install_docker() {
     echo "Installing Docker CE..."
     curl -fsSL https://get.docker.com -o get-docker.sh
+    if [ $? -ne 0 ]; then
+        echo "Error downloading Docker installation script."
+        exit 1
+    fi
     sh get-docker.sh
 
+    if [ $? -ne 0 ]; then
+        echo "Error installing Docker."
+        exit 1
+    fi
+
     # Start and enable Docker service
     systemctl start docker
     systemctl enable docker
@@ -52,12 +70,19 @@ install_docker() {
 
     usermod -aG docker "$username"
     echo "User $username added to the docker group."
-
 }
 
-# Function to create a limited user account
+# Function to create a limited user account with validation
 create_user_account() {
-    read -p "Please enter the username for the new user account: " username
+    while true; do
+        read -p "Please enter the username for the new user account: " username
+        if [[ -z "$username" ]]; then
+            echo "Username cannot be empty. Please try again."
+        else
+            break
+        fi
+    done
+
     read -sp "Please enter the password for the new user account: " password
     echo
 
@@ -71,100 +96,116 @@ create_user_account() {
     fi
 }
 
-# Check if the script is run as root
+# Function to check if running as root
-if [ "$(id -u)" -ne 0 ]; then
+check_root() {
-    echo "This script must be run as root. Please use sudo or switch to root."
+    if [ "$(id -u)" -ne 0 ]; then
-    exit 1
+        echo "This script must be run as root. Please use sudo or switch to root."
-fi
+        exit 1
+    fi
+}
 
-# Update the package manager
+# Function to update package manager cache
-echo "Updating package manager..."
+update_package_manager() {
-if command -v dnf &> /dev/null; then
+    echo "Updating package manager..."
-    dnf -y makecache
+    if command -v dnf &> /dev/null; then
-elif command -v yum &> /dev/null; then
+        dnf -y makecache
-    yum -y makecache
+    elif command -v yum &> /dev/null; then
-else
+        yum -y makecache
-    echo "Neither dnf nor yum found. This script only works on RHEL-based distributions."
+    else
-    exit 1
+        echo "Neither dnf nor yum found. This script only works on RHEL-based distributions."
-fi
+        exit 1
+    fi
+}
 
-# Upgrade installed packages
+# Function to upgrade installed packages
-echo "Upgrading installed packages..."
+upgrade_packages() {
-if command -v dnf &> /dev/null; then
+    echo "Upgrading installed packages..."
-    dnf -y upgrade
+    if command -v dnf &> /dev/null; then
-elif command -v yum &> /dev/null; then
+        dnf -y upgrade
-    yum -y update
+    elif command -v yum &> /dev/null; then
-fi
+        yum -y update
+    fi
+}
 
-# Change the timezone
+# Function to change the timezone
-echo "Setting timezone to $timezone..."
+set_timezone() {
-timedatectl set-timezone "$timezone"
+    echo "Setting timezone to $timezone..."
+    timedatectl set-timezone "$timezone"
+}
 
-# Change the hostname
+# Function to change the hostname
-prompt_for_hostname
+set_hostname() {
-echo "Setting hostname to $hostname..."
+    prompt_for_hostname
-hostnamectl set-hostname "$hostname"
+    echo "Setting hostname to $hostname..."
+    hostnamectl set-hostname "$hostname"
+}
 
-# Update /etc/hosts
+# Function to modify sshd_config
+modify_sshd_config() {
+    echo "Modifying SSH configuration..."
+    if [ -f "$sshd_config" ]; then
+        # Set PermitRootLogin to no
+        sed -i 's/^PermitRootLogin .*/PermitRootLogin no/' "$sshd_config" || echo "PermitRootLogin no" >> "$sshd_config"
+
+        # Set PasswordAuthentication to no
+        if grep -q '^#PasswordAuthentication' "$sshd_config"; then
+            sed -i 's/^#PasswordAuthentication.*/PasswordAuthentication no/' "$sshd_config"
+        else
+            sed -i 's/^PasswordAuthentication .*/PasswordAuthentication no/' "$sshd_config" || echo "PasswordAuthentication no" >> "$sshd_config"
+        fi
+
+        # Ensure AddressFamily inet is set
+        if grep -q '^#AddressFamily' "$sshd_config"; then
+            sed -i 's/^#AddressFamily.*/AddressFamily inet/' "$sshd_config"
+        else
+            sed -i 's/^AddressFamily .*/AddressFamily inet/' "$sshd_config" || echo "AddressFamily inet" >> "$sshd_config"
+        fi
+    else
+        echo "sshd_config file not found. Exiting."
+        exit 1
+    fi
+}
+
+# Function to restart SSH service
+restart_ssh() {
+    echo "Restarting SSH service..."
+    systemctl restart sshd
+}
+
+# Function to wait for SSH key copy
+wait_for_key_copy() {
+    echo "Waiting for your public key to be copied..."
+    while true; do
+        if [ -s "/home/$username/.ssh/authorized_keys" ]; then
+            echo "Public key has been successfully copied."
+            break
+        fi
+        sleep 2
+    done
+}
+
+# Function to clean up unused packages
+cleanup() {
+    echo "Cleaning up..."
+    if command -v dnf &> /dev/null; then
+        dnf -y autoremove
+    elif command -v yum &> /dev/null; then
+        yum -y autoremove
+    fi
+}
+
+# Main execution
+check_root
+update_package_manager
+upgrade_packages
+set_timezone
+set_hostname
 update_hosts_file
-
-# Create a limited user account
 create_user_account
-
+modify_sshd_config
-# Modify sshd_config
+restart_ssh
-echo "Modifying SSH configuration..."
+wait_for_key_copy
-if [ -f "$sshd_config" ]; then
-    # Set PermitRootLogin to no
-    sed -i 's/^PermitRootLogin .*/PermitRootLogin no/' "$sshd_config" || echo "PermitRootLogin no" >> "$sshd_config"
-
-    # Set PasswordAuthentication to no
-    if grep -q '^#PasswordAuthentication' "$sshd_config"; then
-        sed -i 's/^#PasswordAuthentication.*/PasswordAuthentication no/' "$sshd_config"
-    else
-        sed -i 's/^PasswordAuthentication .*/PasswordAuthentication no/' "$sshd_config" || echo "PasswordAuthentication no" >> "$sshd_config"
-    fi
-
-    # Ensure AddressFamily inet is set
-    if grep -q '^#AddressFamily' "$sshd_config"; then
-        sed -i 's/^#AddressFamily.*/AddressFamily inet/' "$sshd_config"
-    else
-        sed -i 's/^AddressFamily .*/AddressFamily inet/' "$sshd_config" || echo "AddressFamily inet" >> "$sshd_config"
-    fi
-else
-    echo "sshd_config file not found. Exiting."
-    exit 1
-fi
-
-# Get the IP address of the machine
-ip_address=$(hostname -I | awk '{print $1}')
-
-# Prompt the user to copy their public key
-echo "Please copy your public SSH key to the server using the following command:"
-echo "ssh-copy-id $username@$ip_address"
-
-# Wait until the authorized_keys file is no longer empty
-echo "Waiting for your public key to be copied..."
-while true; do
-    if [ -s "/home/$username/.ssh/authorized_keys" ]; then
-        echo "Public key has been successfully copied."
-        break
-    fi
-    sleep 2
-done
-
-# Restart SSH service to apply changes
-echo "Restarting SSH service..."
-systemctl restart sshd
-
-# Prompt for Docker installation
 prompt_for_docker_install
-
+cleanup
-# Clean up
-echo "Cleaning up..."
-if command -v dnf &> /dev/null; then
-    dnf -y autoremove
-elif command -v yum &> /dev/null; then
-    yum -y autoremove
-fi
 
 echo "System update complete! Timezone set to $timezone, hostname set to $hostname, limited user created, sshd_config modified, and Docker installation completed if selected."