#!/bin/bash

# Script to update a RHEL-based Linux system, set the timezone, change the hostname,
# update /etc/hosts, modify sshd_config settings, and optionally install Docker CE

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

# Function to prompt for hostname
prompt_for_hostname() {
    read -p "Please enter the desired hostname: " hostname
}

# Function to prompt for SSH key
prompt_for_ssh_key() {
    read -p "Please enter your public SSH key (or press Enter to skip): " ssh_key
    if [ -z "$ssh_key" ]; then
        echo "No SSH key provided. Exiting."
        exit 1
    fi
}

# Function to prompt for Docker installation
prompt_for_docker_install() {
    read -p "Would you like to install Docker CE? (y/n): " install_docker
    if [[ "$install_docker" =~ ^[Yy]$ ]]; then
        install_docker
    else
        echo "Skipping Docker installation."
    fi
}

# Function to install Docker
install_docker() {
    echo "Installing Docker CE..."
    curl -fsSL https://get.docker.com -o get-docker.sh
    sh get-docker.sh

    # Start and enable Docker service
    systemctl start docker
    systemctl enable docker
    echo "Docker installation complete."
}

# Check if the script is run as 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

# Prompt for hostname
prompt_for_hostname

# Prompt for SSH key
prompt_for_ssh_key

# Update the 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

# Upgrade installed 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

# Change the timezone
echo "Setting timezone to $timezone..."
timedatectl set-timezone "$timezone"

# Change the hostname
echo "Setting hostname to $hostname..."
hostnamectl set-hostname "$hostname"

# Get the IP address of the machine
ip_address=$(hostname -I | awk '{print $1}')

# Add an entry to /etc/hosts
echo "Updating /etc/hosts with IP address $ip_address and hostname $hostname..."
if ! grep -q "$ip_address" /etc/hosts; then
    echo "$ip_address $hostname" >> /etc/hosts
else
    echo "Entry for $ip_address already exists in /etc/hosts."
fi

# 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
    sed -i 's/^PasswordAuthentication .*/PasswordAuthentication no/' "$sshd_config" || echo "PasswordAuthentication no" >> "$sshd_config"

    # Ensure AddressFamily inet is not commented out
    sed -i 's/^#\?\(AddressFamily inet\)/\1/' "$sshd_config"
else
    echo "sshd_config file not found. Exiting."
    exit 1
fi

# 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

# Restart SSH service to apply changes
echo "Restarting SSH service..."
systemctl restart sshd

# Prompt for Docker installation
prompt_for_docker_install

echo "System update complete! Timezone set to $timezone, hostname set to $hostname, /etc/hosts updated, and sshd_config modified."