fix prompts, add in MOTD function

InitMate.sh

@@ -7,11 +7,13 @@ log() {
     echo "$1" | tee -a "$LOG_FILE"
 }
 
-# Check if the script is run as root
+# Function to check if the script is run as root
-if [ "$EUID" -ne 0 ]; then
+check_root() {
+    if [ "$EUID" -ne 0 ]; then
         log "Please run as root"
         exit 1
-fi
+    fi
+}
 
 # Function to detect the OS and package manager
 detect_os() {
@@ -43,6 +45,18 @@ detect_os() {
     esac
 }
 
+# Function to prompt for yes/no input
+prompt_yes_no() {
+    while true; do
+        read -r -p "$1 (y/n): " REPLY
+        case "$REPLY" in
+            [Yy]* ) return 0 ;;
+            [Nn]* ) return 1 ;;
+            * ) log "Please answer yes or no." ;;
+        esac
+    done
+}
+
 # Function to secure SSH
 secure_ssh() {
     log "Securing SSH..."
@@ -75,9 +89,7 @@ setup_firewall() {
         firewall-cmd --permanent --add-port="$SSH_PORT"/tcp
         firewall-cmd --reload
     else
-        log "No firewall detected. Do you want to install UFW? (yes/no)"
+        if prompt_yes_no "No firewall detected. Do you want to install UFW?"; then
-        read -r INSTALL_UFW
-        if [[ $INSTALL_UFW == "yes" ]]; then
             if [[ $PKG_MANAGER == "apt" ]]; then
                 apt update && apt install -y ufw
             elif [[ $PKG_MANAGER == "yum" ]]; then
@@ -90,9 +102,7 @@ setup_firewall() {
         fi
     fi
 
-    log "Do you want to open additional ports? (yes/no)"
+    if prompt_yes_no "Do you want to open additional ports?"; then
-    read -r OPEN_PORTS
-    if [[ $OPEN_PORTS == "yes" ]]; then
         log "Please enter the ports and protocols to open (comma-separated, e.g., 80/tcp,443/tcp,8080/udp):"
         read -r ADDITIONAL_PORTS
         IFS=',' read -r -a PORT_ARRAY <<< "$ADDITIONAL_PORTS"
@@ -105,7 +115,7 @@ setup_firewall() {
             fi
         done
         if command -v firewall-cmd &>/dev/null; then
-            firewall-cmd --reload
+            firewall_cmd --reload
         fi
     fi
 }
@@ -123,9 +133,7 @@ system_update() {
 
 # Function to run additional setup
 run_additional_setup() {
-    log "Do you want to run additional setup scripts? (yes/no)"
+    if prompt_yes_no "Do you want to run additional setup scripts?"; then
-    read -r RUN_ADDITIONAL_SETUP
-    if [[ $RUN_ADDITIONAL_SETUP == "yes" ]]; then
         SCRIPT_DIR=$(dirname "$0")
         if [[ -f "$SCRIPT_DIR/extender.sh" ]]; then
             bash "$SCRIPT_DIR/extender.sh" | tee -a "$LOG_FILE"
@@ -137,9 +145,8 @@ run_additional_setup() {
 
 # Function to add additional users
 add_users() {
-    log "Do you want to add additional users? (yes/no)"
+    if prompt_yes_no "Do you want to add additional users?"; then
-    read -r ADD_USERS
+        while true; do
-    while [[ $ADD_USERS == "yes" ]]; do
             log "Enter the username:"
             read -r USERNAME
             log "Enter the password:"
@@ -147,33 +154,40 @@ add_users() {
             useradd -m "$USERNAME"
             echo "$USERNAME:$PASSWORD" | chpasswd
 
-        log "Do you want to grant sudo access to $USERNAME? (yes/no)"
+            if prompt_yes_no "Do you want to grant sudo access to $USERNAME?"; then
-        read -r GRANT_SUDO
-        if [[ $GRANT_SUDO == "yes" ]]; then
                 usermod -aG sudo "$USERNAME"
                 log "$USERNAME has been granted sudo access."
             fi
 
-        log "Do you want to add a public key for $USERNAME? (yes/no)"
+            if prompt_yes_no "Do you want to add a public key for $USERNAME?"; then
-        read -r ADD_PUBLIC_KEY
-        if [[ $ADD_PUBLIC_KEY == "yes" ]]; then
                 log "Please enter the public key:"
                 read -r PUBLIC_KEY
                 su - "$USERNAME" -c "mkdir -p ~/.ssh && echo '$PUBLIC_KEY' >> ~/.ssh/authorized_keys && chmod 600 ~/.ssh/authorized_keys && chmod 700 ~/.ssh"
                 log "Public key added for $USERNAME."
             fi
 
-        log "Do you want to add another user? (yes/no)"
+            if ! prompt_yes_no "Do you want to add another user?"; then
-        read -r ADD_USERS
+                break
+            fi
         done
+    fi
+}
+
+# Function to set up the MOTD
+setup_motd() {
+    if prompt_yes_no "Do you want to set up the Message of the Day (MOTD)?"; then
+        log "Please enter the MOTD content:"
+        read -r MOTD_CONTENT
+        echo "$MOTD_CONTENT" > /etc/motd
+        log "MOTD has been set."
+    fi
 }
 
 # Main script execution
+check_root
 detect_os
 log "Detected OS: $OS, Package Manager: $PKG_MANAGER"
-log "Is this correct? (yes/no)"
+if ! prompt_yes_no "Is this correct?"; then
-read -r CONFIRM
-if [[ $CONFIRM != "yes" ]]; then
     log "Please enter the correct OS and package manager:"
     read -r OS PKG_MANAGER
     if [[ $PKG_MANAGER != "apt" && $PKG_MANAGER != "yum" ]]; then
@@ -184,21 +198,15 @@ if [[ $CONFIRM != "yes" ]]; then
     fi
 fi
 
-log "Do you want to secure SSH? (yes/no)"
+if prompt_yes_no "Do you want to secure SSH?"; then
-read -r SECURE_SSH
-if [[ $SECURE_SSH == "yes" ]]; then
     secure_ssh
 fi
 
-log "Do you want to set up a firewall? (yes/no)"
+if prompt_yes_no "Do you want to set up a firewall?"; then
-read -r SETUP_FIREWALL
-if [[ $SETUP_FIREWALL == "yes" ]]; then
     setup_firewall
 fi
 
-log "Do you want to perform a system update? (yes/no)"
+if prompt_yes_no "Do you want to perform a system update?"; then
-read -r UPDATE_SYSTEM
-if [[ $UPDATE_SYSTEM == "yes" ]]; then
     system_update
 fi
 
@@ -212,3 +220,6 @@ run_additional_setup
 
 # Add additional users if requested
 add_users
+
+# Set up MOTD if requested
+setup_motd