#!/bin/bash

# Script to update a RHEL-based Linux system, set the timezone, change the hostname,
# add a limited user account, modify sshd_config settings, and optionally install Docker CE

# Variables
timezone="America/New_York"
sshd_config="/etc/ssh/sshd_config"
docker_installed=false

# Function to prompt for hostname with validation
prompt_for_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
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."
        sed -i "s/.*$hostname/$ip_address $hostname/" /etc/hosts
    else
        echo "$ip_address $hostname" >> /etc/hosts
        echo "Added $ip_address $hostname to /etc/hosts."
    fi
}

# Function to prompt for Docker installation with validation
prompt_for_docker_install() {
    while true; do
        read -p "Would you like to install Docker CE? (y/n): " install_docker
        case "$install_docker" in
            [Yy]*) install_docker ;;
            [Nn]*) echo "Skipping Docker installation."; break ;;
            *) 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
    echo "Docker installation complete."

    usermod -aG docker "$username"
    echo "User $username added to the docker group."
}

# Function to create a limited user account with validation
create_user_account() {
    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

    # Create the user and add to the wheel group
    if id "$username" &>/dev/null; then
        echo "User $username already exists."
    else
        useradd -m -G wheel "$username"
        echo "$username:$password" | chpasswd
        echo "User $username created and added to the wheel group."
    fi
}

# Function to check if running as root
check_root() {
    if [ "$(id -u)" -ne 0 ]; then
        echo "This script must be run as root. Please use sudo or switch to root."
        exit 1
    fi
}

# Function to update package manager cache
update_package_manager() {
    echo "Updating package manager..."
    if command -v dnf &> /dev/null; then
        dnf -y makecache
    elif command -v yum &> /dev/null; then
        yum -y makecache
    else
        echo "Neither dnf nor yum found. This script only works on RHEL-based distributions."
        exit 1
    fi
}

# Function to upgrade installed packages
upgrade_packages() {
    echo "Upgrading installed packages..."
    if command -v dnf &> /dev/null; then
        dnf -y upgrade
    elif command -v yum &> /dev/null; then
        yum -y update
    fi
}

# Function to change the timezone
set_timezone() {
    echo "Setting timezone to $timezone..."
    timedatectl set-timezone "$timezone"
}

# Function to change the hostname
set_hostname() {
    prompt_for_hostname
    echo "Setting hostname to $hostname..."
    hostnamectl set-hostname "$hostname"
}

# 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_user_account
modify_sshd_config
restart_ssh
wait_for_key_copy
prompt_for_docker_install
cleanup

echo "System update complete! Timezone set to $timezone, hostname set to $hostname, limited user created, sshd_config modified, and Docker installation completed if selected."