#!/bin/bash

##########################################################################
# Basic OS configuration required for NodeZero execution on              #
# RHEL 10+                                                               #
#                                                                        #
# Copyright 2024 -  Horizon3 AI Inc.                                     #
##########################################################################

SCRIPT_VERSION="000"

OS="UNKNOWN"
VER="UNKNOWN"
BUILD="UNKNOWN"

USERNAME=$(whoami)

home_dir="/home/$USERNAME"

banner_file="banner"
n0_utility="n0-rhel.sh"
tmp_advisory_file="/tmp/update-motd-advisory"
tmp_banner_file="/tmp/${banner_file}"

# H3 CLI download
download_url="https://downloads.horizon3ai.com"
download_cli_url="https://downloads.horizon3ai.com/utilities/cli/h3-cli.zip"

# H3 CLI
h3_cli_zip="h3cli.zip"
h3_cli="h3-cli"

# ANSI colors
if which tput &>/dev/null; then
    BOLD=$(tput bold)
    NORMAL=$(tput sgr0)
    RED=$(tput setaf 1)     # Fail
    GREEN=$(tput setaf 2)
    YELLOW=$(tput setaf 3)  # Warning
    BLUE=$(tput setaf 4)
    MAGENTA=$(tput setaf 5) # Ask for Input
    CYAN=$(tput setaf 6)    # Header/Menu
    WHITE=$(tput setaf 7)   # General
    FANCY=$(echo -e "cuu1\nel" | tput -S)
fi

# Create utility functions
function HeaderMsg {
    echo -ne "\n${BOLD}[${CYAN}+${NORMAL}${BOLD}]${CYAN} $@${NORMAL}\n"
}

function MenuMsg {
    echo -ne "\n${BOLD}${CYAN}${NORMAL}${BOLD}${CYAN} $@${NORMAL}"
}

function InfoMsg {
    echo -ne "${BOLD}[${WHITE}+${NORMAL}${BOLD}]${WHITE} $@${NORMAL}\n"
}

function FailMsg {
    echo -ne "${BOLD}[${RED}!${NORMAL}${BOLD}] ${NORMAL}${RED}FAILED: $@${NORMAL}\n"
}

function WarnMsg {
    echo -ne "${BOLD}[${YELLOW}!${NORMAL}${BOLD}] ${NORMAL}${YELLOW}WARNING: $@${NORMAL}\n"
}

function AskMsg {
    echo -ne "${BOLD}[${MAGENTA}-${NORMAL}${BOLD}] ${NORMAL}${MAGENTA}CONFIRM: $@${NORMAL}"
}

function InputMsg {
    echo -ne "${BOLD}[${MAGENTA}-${NORMAL}${BOLD}] ${NORMAL}${MAGENTA}$@${NORMAL}"
}

if [ $USERNAME == "root" ]; then
    FailMsg "This script must be run as a non-root user with sudo privileges"
    exit 1
fi

if [ ! -z "$SUDO_USER" ]; then
    FailMsg "Do not use 'sudo' to run this script, it will automatically use sudo where needed."
    exit 1
fi


function yes_no_validation() {
    # pass in the string to ask yes/no on
    local question=$@
    AskMsg "$question (y|n)? "
    read input

    lower_input=$(echo "$input" | tr '[:upper:]' '[:lower:]') &>/dev/null

    case $lower_input in
        "") return 0 ;;
        y | yes) return 1 ;;
        n | no) return 2 ;;
        q) exit 0 ;;
        *) FailMsg "Invalid option. Enter q to exit. Reasking question..." ; yes_no_validation $question;;
    esac
}

function set_build() {
    # Detect OS type and version
    if [ -f /etc/os-release ]; then
        . /etc/os-release

        if [[ $NAME == "Red Hat Enterprise Linux" ]]; then
            OS="RHEL"
        else
            OS=$NAME
        fi
        VER=$VERSION_ID

    elif [ -f /etc/redhat-release ]; then
        OS="RHEL"
        VER=$(cat /etc/redhat-release | grep -oE '[0-9]+\.[0-9]+' | cut -d'.' -f1)
    else
        OS="UNKNOWN"
        VER="UNKNOWN"
    fi

    # Set build version
    if [[ "$OS" == "RHEL" ]]; then
        if [ "$VER" -lt 10 ]; then
            FailMsg "Unsupported RHEL version detected: $VER (RHEL 10+ required)"
            local unsupported_os_continue="This is an unsupported OS version. Do you want to continue"
            yes_no_validation $unsupported_os_continue
            rc=$?
            if [ $rc -ne 1 ] ; then
                echo -ne "Thank you, exiting...\n"
                exit 1
            fi
        fi
        BUILD="NodeZero-${SCRIPT_VERSION}-${OS}-${VER}"
    else
        FailMsg "Unsupported OS detected: $OS. This script requires RHEL 10+."
        exit 1
    fi

}

function download() {
    filename=${1}  # n0_utility, banner_file
    destination=${2}

    endpoint_url="${download_url}/utilities/${filename}"

    # Check connection to url
    if ! curl -L -I ${endpoint_url} 2>&1 | grep "HTTP.*200" >/dev/null; then
        FailMsg "Failed to connect to ${endpoint_url}"
        return 1
    fi

    # Make a random temp directory
    tmpdir="$(mktemp -d)"

    # Download to temporary dir
    pushd ${tmpdir} > /dev/null
    rm -f ${filename}.*
    curl -s -o ${filename} ${endpoint_url} &> /dev/null
    curl -s -o ${filename}.sha256.checksum ${endpoint_url}.sha256.checksum &> /dev/null

    # Verify checksum
    if sha256sum -c ${filename}.sha256.checksum &> /dev/null; then
        # Checksum valid
        # Cleanup
        rm -f ${filename}.sha256.checksum

        # Move the file to the final destination
        sudo mv ${filename} ${destination}
        popd > /dev/null
        rm -Rf ${tmpdir}

    else
        # Checksum invalid
        FailMsg "Checksum invalid for ${filename}"
        FailMsg "Removing ${filename}..."

        # Cleanup
        rm -f ${filename}.*
        popd > /dev/null
        rm -Rf ${tmpdir}

        exit 1
    fi
}

function install_RHEL_packages() {
    # RHEL 10+ package installation
    sudo dnf check-update || true

    # Perform a system update
    sudo dnf upgrade -y
    sudo dnf distro-sync -y

    # Install EPEL repository
    sudo dnf install -y https://dl.fedoraproject.org/pub/epel/epel-release-latest-$(rpm -E %rhel).noarch.rpm

    # Install basic packages
    sudo dnf install -y wget curl sudo nano openssh-server nmap tcpdump \
        bind-utils unzip jq net-tools cloud-init telnet traceroute

    # Install Podman (rootless) and docker CLI shim
    sudo dnf install -y podman
    rc=$?
    if [ ${rc} -ne 0 ]; then
        FailMsg "Failed to install podman packages"
        exit 1
    fi

    # Virtualization support
    # virtualbox-guest-additions is unsupported on RHEL
    # ec2-instance-connect is not available on RHEL, cloud-init is used instead
    sudo dnf install -y hyperv-daemons open-vm-tools
}

# Configure firewall
function configure_firewall() {
    sudo systemctl stop firewalld
    sudo systemctl disable firewalld
    sudo systemctl mask firewalld

    # Configure SELinux
    sudo setenforce 0
    sudo sed -i 's/^SELINUX=enforcing$/SELINUX=permissive/' /etc/selinux/config
}

function rhel_special_sauce() {
    banner_location=$1
    advisory_location=$2

    n0_ip_and_runners="nodezero-show-ip-and-runners"
    n0_set_region="nodezero-set-region"
    n0_gateway_connect="nodezero-gateway-connect"

    echo "> Downloading set-nodezero-region script..."
    download ${n0_set_region} /etc/profile.d/20-nodezero-set-region.sh
    rc=$?
    if [ ${rc} -ne 0 ]; then
        FailMsg "Failed to download ${n0_set_region}"
        exit 1
    fi
    sudo chmod +x /etc/profile.d/20-nodezero-set-region.sh

    echo "> Downloading show-ip-and-runners script..."
    download ${n0_ip_and_runners} /etc/profile.d/40-nodezero-show-ip-and-runners.sh
    rc=$?
    if [ ${rc} -ne 0 ]; then
        FailMsg "Failed to download ${n0_ip_and_runners}"
        exit 1
    fi
    sudo chmod +x /etc/profile.d/40-nodezero-show-ip-and-runners.sh

    echo "> Downloading nodezero-gateway-connect script..."
    download ${n0_gateway_connect} /etc/profile.d/60-nodezero-gateway-connect.sh
    rc=$?
    if [ ${rc} -ne 0 ]; then
        FailMsg "Failed to download ${n0_gateway_connect}"
        exit 1
    fi
    sudo chmod +x /etc/profile.d/60-nodezero-gateway-connect.sh

    # Warning message displayed in motd
    echo "> Setting up advisory message and banner..."
    sudo mv ${advisory_location} /usr/lib/motd.d/10-update-motd-advisory
    sudo chmod 644 /usr/lib/motd.d/10-update-motd-advisory

    if [ -f /usr/lib/motd.d/10-update-motd-advisory ]; then
        cat /usr/lib/motd.d/10-update-motd-advisory > /tmp/welcome
        echo '' >> /tmp/welcome
        echo '' >> /tmp/welcome
        sudo cp /tmp/welcome /etc/issue.net
        sudo mv /tmp/welcome /etc/issue
    fi

    sudo mv ${banner_location} /usr/lib/motd.d/11-nodezero-banner
    sudo chmod 644 /usr/lib/motd.d/11-nodezero-banner

    # Setup to issue warning message on initial connection
    # Default /etc/ssh/sshd_config has `PrintMotd no`
    echo '> Enabling MOTD...'
    sudo sed -i "s/PrintMotd.*/PrintMotd yes/" /etc/ssh/sshd_config

}

# Writing build string
set_build
echo
echo "> Build Version: ${BUILD}"
echo "${BUILD}" > /tmp/nodezero-build
sudo bash -c 'mv /tmp/nodezero-build /etc/nodezero-build'

# Main installation process
echo "> Detected OS: $OS $VER"
echo "> Home Directory: $home_dir"

# Disable journald to avoid disk space issues
echo
echo "> Disabling journald..."
sudo systemctl stop systemd-journald &> /dev/null
sudo systemctl disable systemd-journald &> /dev/null

# Configure firewall
echo
echo "> Attempting to disable and mask firewalld..."
configure_firewall

# Make sure the OS is updated to the latest packages
# Perform a system update
echo
echo "> Updating OS packages, This may take a bit..."
install_RHEL_packages

# Download the NodeZero banner
download ${banner_file} ${tmp_banner_file}
rc=$?
if [ ${rc} -ne 0 ]; then
    FailMsg "Failed to download ${banner_file}"
    exit 1
fi

# Warning message displayed in motd
cat > ${tmp_advisory_file} << EOT


WARNING : Unauthorized access to this system is forbidden and will be
prosecuted by law. By accessing this system, you agree that your actions
may be monitored if unauthorized usage is suspected.


EOT

rhel_special_sauce $tmp_banner_file $tmp_advisory_file

# Switch to $home_dir directory
echo "> Switching to ${home_dir} directory to download h3-cli..."
cd $home_dir
rc=$?
if [ $rc -ne 0 ]; then
    FailMsg "Failed to switch to ${home_dir} directory, check if this directory exists or if you have permissions to access it."
    exit 1
fi

# Download or update the h3-cli
echo
echo "> Checking connection to downloads.horizon3ai.com..."
if curl -L -I ${download_cli_url} 2>&1 | grep "HTTP.*200" >/dev/null; then

    # Might need to overwrite existing h3-cli
    if [[ -d ${h3_cli} ]]; then
        echo "> Removing existing h3-cli directory..."
        sudo rm -rf ${h3_cli}
    fi

    echo "> Downloading h3-cli..."
    curl -L ${download_cli_url} -o ${h3_cli_zip}
    rc=$?
    if [ ${rc} -ne 0 ]; then
        FailMsg "Failed to download ${download_cli_url}"
        exit 1
    fi

    echo "> Unzipping h3-cli..."
    unzip ${h3_cli_zip} -d ${h3_cli} &> /dev/null
    rc=$?
    if [ ${rc} -ne 0 ]; then
        FailMsg "Failed to unzip ${h3_cli_zip}"
        exit 1
    fi
    rm -f ${h3_cli_zip}

else
  FailMsg "Failed to connect to ${download_url}"
  FailMsg "Unable to download latest h3-cli"
fi

# Download the n0 utility
echo
echo "> Downloading n0 utility script..."
download ${n0_utility} /tmp/${n0_utility}
rc=$?
if [ ${rc} -ne 0 ]; then
  FailMsg "Unable to download ${n0_utility}"
  # Dont remove existing n0 if download fails
  exit 1
fi
sudo mv /tmp/${n0_utility} /usr/local/bin/n0
sudo chmod a+rx /usr/local/bin/n0

# Podman (rootless) is required to run NodeZero
echo
echo "> Configuring rootless podman..."
sudo loginctl enable-linger $USERNAME
systemctl --user enable --now podman.socket

echo "> Rebooting system to apply changes..."
secs=10  # Set countdown timer to 10 seconds
while [ $secs -gt 0 ]; do
   echo -ne "$secs\n"
   sleep 1
   : $((secs--))
done

sudo reboot