diff --git a/InitMate.sh b/InitMate.sh
index 01feca1..a897a58 100644
--- a/InitMate.sh
+++ b/InitMate.sh
@@ -7,11 +7,13 @@ log() {
     echo "$1" | tee -a "$LOG_FILE"
 }
 
-# Check if the script is run as root
-if [ "$EUID" -ne 0 ]; then
-    log "Please run as root"
-    exit 1
-fi
+# Function to check if the script is run as root
+check_root() {
+    if [ "$EUID" -ne 0 ]; then
+        log "Please run as root"
+        exit 1
+    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)"
-        read -r INSTALL_UFW
-        if [[ $INSTALL_UFW == "yes" ]]; then
+        if prompt_yes_no "No firewall detected. Do you want to install UFW?"; 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)"
-    read -r OPEN_PORTS
-    if [[ $OPEN_PORTS == "yes" ]]; then
+    if prompt_yes_no "Do you want to open additional ports?"; 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)"
-    read -r RUN_ADDITIONAL_SETUP
-    if [[ $RUN_ADDITIONAL_SETUP == "yes" ]]; then
+    if prompt_yes_no "Do you want to run additional setup scripts?"; then
         SCRIPT_DIR=$(dirname "$0")
         if [[ -f "$SCRIPT_DIR/extender.sh" ]]; then
             bash "$SCRIPT_DIR/extender.sh" | tee -a "$LOG_FILE"
@@ -137,43 +145,49 @@ run_additional_setup() {
 
 # Function to add additional users
 add_users() {
-    log "Do you want to add additional users? (yes/no)"
-    read -r ADD_USERS
-    while [[ $ADD_USERS == "yes" ]]; do
-        log "Enter the username:"
-        read -r USERNAME
-        log "Enter the password:"
-        read -r -s PASSWORD
-        useradd -m "$USERNAME"
-        echo "$USERNAME:$PASSWORD" | chpasswd
+    if prompt_yes_no "Do you want to add additional users?"; then
+        while true; do
+            log "Enter the username:"
+            read -r USERNAME
+            log "Enter the password:"
+            read -r -s PASSWORD
+            useradd -m "$USERNAME"
+            echo "$USERNAME:$PASSWORD" | chpasswd
 
-        log "Do you want to grant sudo access to $USERNAME? (yes/no)"
-        read -r GRANT_SUDO
-        if [[ $GRANT_SUDO == "yes" ]]; then
-            usermod -aG sudo "$USERNAME"
-            log "$USERNAME has been granted sudo access."
-        fi
+            if prompt_yes_no "Do you want to grant sudo access to $USERNAME?"; 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)"
-        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
+            if prompt_yes_no "Do you want to add a public key for $USERNAME?"; 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)"
-        read -r ADD_USERS
-    done
+            if ! prompt_yes_no "Do you want to add another user?"; then
+                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)"
-read -r CONFIRM
-if [[ $CONFIRM != "yes" ]]; then
+if ! prompt_yes_no "Is this correct?"; 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)"
-read -r SECURE_SSH
-if [[ $SECURE_SSH == "yes" ]]; then
+if prompt_yes_no "Do you want to secure SSH?"; then
     secure_ssh
 fi
 
-log "Do you want to set up a firewall? (yes/no)"
-read -r SETUP_FIREWALL
-if [[ $SETUP_FIREWALL == "yes" ]]; then
+if prompt_yes_no "Do you want to set up a firewall?"; then
     setup_firewall
 fi
 
-log "Do you want to perform a system update? (yes/no)"
-read -r UPDATE_SYSTEM
-if [[ $UPDATE_SYSTEM == "yes" ]]; then
+if prompt_yes_no "Do you want to perform a system update?"; 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