#!/bin/bash
#
# Script to perform system maintenance on a RHEL-based Linux system.
# 
# This script:
#  - Updates package manager and installed packages
#  - Sets the system timezone based on geolocation or user input
#  - Changes the system hostname and updates /etc/hosts
#  - Creates a limited user account and modifies SSH settings
#  - Optionally installs Docker CE
#
# Prerequisites:
#  - The script must be run as root or with sudo privileges
#
# Usage:
#  - Run the script directly on the system or via SSH.
#  - Respond to prompts as required during the script execution.

# Variables
sshd_config="/etc/ssh/sshd_config"
docker_installed=false

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

# 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
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."

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

    # Clean up get-docker.sh
    rm -f get-docker.sh
}

# Function to create a limited user account
create_user_account() {
    read -p "Please enter the username for the new user account: " username
    read -sp "Please enter the password for the new user account: " password
    echo
    read -sp "Please confirm the password: " password_confirm
    echo

    if [ "$password" != "$password_confirm" ]; then
        echo "Passwords do not match. Exiting."
        exit 1
    fi

    # 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
}

# 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

# 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

# Function to set the timezone to UTC in case of an error
set_utc_timezone() {
  echo "Error occurred while determining the timezone. Falling back to UTC."
  timedatectl set-timezone UTC
}

# Function to get the timezone from the ipinfo.io API
get_timezone_from_api() {
  # Fetch geolocation information using the ipinfo.io API
  response=$(curl -s https://ipinfo.io)

  # Check if the curl command succeeded and the response contains the 'timezone' field
  if [ $? -eq 0 ] && echo "$response" | grep -q "timezone"; then
    # Extract the timezone directly from the JSON response
    timezone=$(echo "$response" | jq -r '.timezone')

    echo "Detected timezone: $timezone"
    return 0
  else
    set_utc_timezone
    return 1
  fi
}

# Function to prompt the user to choose a timezone
prompt_for_timezone() {
  echo ""
  echo "Choose a timezone option:"
  echo "1) Use the detected timezone ($1)"
  echo "2) Use UTC"
  echo "3) Enter a custom timezone"
  
  # Read user choice
  read -p "Enter the number corresponding to your choice: " choice
  
  case "$choice" in
    1)
      echo "You chose to use the detected timezone: $1"
      timedatectl set-timezone "$1"
      ;;
    2)
      echo "You chose to use UTC."
      timedatectl set-timezone UTC
      ;;
    3)
      # Ask for a custom timezone
      read -p "Enter your preferred timezone (e.g., Europe/London, America/New_York): " custom_timezone

      if timedatectl list-timezones | grep -q "$custom_timezone"; then
          timedatectl set-timezone "$custom_timezone"
      else
          echo "Invalid timezone. Falling back to UTC."
          timedatectl set-timezone UTC
      fi
      ;;
    *)
      echo "Invalid choice. Falling back to UTC."
      timedatectl set-timezone UTC
      ;;
  esac
}

# Main script execution starts here
echo "Attempting to detect and set the timezone..."

# Try to get the detected timezone from the API
if get_timezone_from_api; then
  # If the timezone was successfully detected, prompt the user for their choice
  prompt_for_timezone "$timezone"
else
  # If no timezone was detected, ask the user to fall back to UTC
  echo "Unable to detect timezone. Falling back to UTC."
  timedatectl set-timezone UTC
fi

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

# Update /etc/hosts
update_hosts_file

# Create a limited user account
create_user_account

# Modify sshd_config
echo "Modifying SSH configuration..."
if [ -f "$sshd_config" ]; then
    # Set PermitRootLogin to no
    if ! grep -q "^PermitRootLogin" "$sshd_config"; then
        echo "PermitRootLogin no" >> "$sshd_config"
    else
        sed -i 's/^PermitRootLogin.*/PermitRootLogin no/' "$sshd_config"
    fi

    # 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

# 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:-UTC}, hostname set to $hostname, limited user created, sshd_config modified, and Docker installation completed if selected."