#!/bin/bash

# Variables to store the checklist status
CHECKLIST=()

# Global variables
MIN_FREE_MEMORY=500   # MB
MIN_CPU_CORES=2       # Cores
MIN_CPU_FREE=0.25     # 25%
MIN_HOME_SPACE=500    # MB
MIN_OPT_SPACE=2000    # GB

# Function to display usage
usage() {
    echo
    echo "Usage: $0 <username> <CHAI= URL> [--update]"
    echo "  <username>: The username to be used by CHAI."
    echo "  <CHAI URL>: To check if CHAI is reachable."
    echo "  --update: Optional flag to allow script to do modifications."
    echo
    exit 1
}

# Parse command-line arguments
if [ "$#" -lt 2 ]; then
    usage
fi

ssh_user="$1"
sudoers_entry="$ssh_user ALL=(ALL) NOPASSWD:/opt/.ch-tools/*/*/*, /opt/.ch-tools/*/*, /bin/mkdir, /bin/echo, /bin/chmod 755 /opt/.ch-tools/*, /bin/chmod 755 /home/$ssh_user/chcmd, /bin/chmod -R 755 /opt/.ch-tools, /bin/chmod -R 755 /home/$ssh_user/.ch-tools, /bin/chown $ssh_user\: /opt/.ch-tools/*, /bin/chown -R $ssh_user\: /home/$ssh_user/chcmd, /bin/chown -R $ssh_user\: /opt/.ch-tools, /bin/chown -R $ssh_user\: /home/$ssh_user/.ch-tools"
url="$2"

# Check if --update is present
if [ "$#" -ge 3 ]; then
    if [ "$3" == "--update" ]; then
        update_flag=1
    else
        echo "Error: Unknown option: $3. Did you mean '--update'?"
        update_flag=0
    fi
else
    update_flag=0  # Set to 0 if --update is not provided
fi

# Check if a username was provided
if [ -z "$ssh_user" ]; then
    echo "Error: No username provided."
    usage
fi

# Check if the user exists
if ! id "$ssh_user" &>/dev/null; then
    echo "Error: User $ssh_user does not exist."
    CHECKLIST+=("User existence check: Failed")
    echo "Checklist for requirements:"
    for item in "${CHECKLIST[@]}"; do
        echo "- $item"
    done
    exit 1
fi

# Function to check resources
check_resources() {
    echo
    echo "============= Checking resource requirements ============="

    # Check free memory using basic method
    free_memory=$(grep MemAvailable /proc/meminfo | awk '{print $2}')
    free_memory=$((free_memory / 1024)) # Convert to MB
    if [ $? -ne 0 ]; then
        echo "Failed to check free memory."
        CHECKLIST+=("Free memory check: Failed")
        return 1
    fi

    if [ $free_memory -ge $MIN_FREE_MEMORY ]; then
        echo "Free memory: Passed ($free_memory MB available)"
        CHECKLIST+=("Free memory: Passed")
    else
        echo "Free memory: Failed ($free_memory MB available). Minimum required is $MIN_FREE_MEMORY MB."
        CHECKLIST+=("Free memory: Failed ($free_memory MB available)")
    fi

    # Check CPU cores
    cpu_cores=$(nproc)

    if [ $? -ne 0 ]; then
        echo "Failed to check CPU cores."
        CHECKLIST+=("CPU cores check: Failed")
        return 1
    fi

    if [ $cpu_cores -ge $MIN_CPU_CORES ]; then
        echo "CPU: Passed ($cpu_cores cores)"
        CHECKLIST+=("CPU: Passed")
    else
        echo "CPU: Failed ($cpu_cores cores). Minimum required is $MIN_CPU_CORES cores."
        CHECKLIST+=("CPU: Failed ($cpu_cores cores)")
    fi

    # Check /home/ space
    home_free=$(df -m /home | awk 'NR==2 {print $4}')

    if [ $? -ne 0 ]; then
        echo "Failed to check /home/ space."
        CHECKLIST+=("/home/ space check: Failed")
        return 1
    fi

    if [ $home_free -ge $MIN_HOME_SPACE ]; then
        echo "/home/ space: Passed ($home_free MB available)"
        CHECKLIST+=("/home/ space: Passed")
    else
        echo "/home/ space: Failed ($home_free MB available). Minimum required is $MIN_HOME_SPACE MB."
        CHECKLIST+=("/home/ space: Failed ($home_free MB available)")
    fi

    # Check /opt/ space
    opt_free=$(df -m /opt | awk 'NR==2 {print $4}')

    if [ $? -ne 0 ]; then
        echo "Failed to check /opt/ space."
        CHECKLIST+=("/opt/ space check: Failed")
        return 1
    fi

    if [ $opt_free -ge $MIN_OPT_SPACE ]; then
        echo "/opt/ space: Passed ($opt_free MB available)"
        CHECKLIST+=("/opt/ space: Passed")
    else
        echo "/opt/ space: Failed ($opt_free MB available). Minimum required is $MIN_OPT_SPACE MB."
        CHECKLIST+=("/opt/ space: Failed ($opt_free MB available)")
    fi
}

# Function to check if user is in sudoers list
check_sudoers() {
    echo "============= Checking sudoers configuration ============="

    # Check if the user is in the sudo group
    sudo_group=$(id -nG "$ssh_user" | grep -w "sudo")
    if [ $? -ne 0 ]; then
        echo "Failed to check if user is in sudo group."
        echo
        CHECKLIST+=("Sudo group check: Failed")
        return 1
    fi

    if [ -z "$sudo_group" ]; then
        echo "User $ssh_user is not in the sudo group."
        CHECKLIST+=("User sudo group membership: Failed")
    else
        echo "User $ssh_user is in the sudo group."
        CHECKLIST+=("User sudo group membership: Passed")
    fi

    # Check and update the sudoers entry
    if sudo grep -qF "$ssh_user ALL=(ALL) NOPASSWD:" /etc/sudoers; then
        entry_found=$(sudo grep "$ssh_user ALL=(ALL) NOPASSWD:" /etc/sudoers)

        if [[ "$entry_found" == "$sudoers_entry" ]]; then
            echo "Sudoers entry already exists."
            CHECKLIST+=("Sudoers entry: Passed")
        else
            echo "Sudoers entry format differs."
            echo "Expected permissions: $sudoers_entry"
            echo
            echo "But found: $entry_found"
            echo
            CHECKLIST+=("Sudoers entry: Format differs")

            # Update the entry if the update flag is set
            if [ $update_flag -eq 1 ]; then
                echo "Updating sudoers entry..."
                # Remove the old entry
                sudo sed -i.bak "/$ssh_user ALL=(ALL) NOPASSWD:/d" /etc/sudoers
                # Add the new entry
                echo "$sudoers_entry" | sudo EDITOR='tee -a' visudo > /dev/null 2>&1
                if [ $? -eq 0 ]; then
                    echo "Sudoers entry updated successfully."
                    echo
                    echo "Updated suoders entry is as follows : "
                    echo "$sudoers_entry"
                    echo
                    CHECKLIST[-1]="Sudoers entry: Passed"
                else
                    echo "Failed to update sudoers entry."
                    CHECKLIST[-1]="Sudoers entry: Failed"
                fi
            fi
        fi
    else
        echo "Sudoers entry is missing."
        CHECKLIST+=("Sudoers entry: Failed")

        if [ $update_flag -eq 1 ]; then
            echo "Adding sudoers entry..."
            echo "$sudoers_entry" | sudo EDITOR='tee -a' visudo > /dev/null 2>&1
            if [ $? -eq 0 ]; then
                echo "Sudoers entry added successfully."
                CHECKLIST[-1]="Sudoers entry: Passed"
            else
                echo "Failed to add sudoers entry."
                CHECKLIST[-1]="Sudoers entry: Failed"
            fi
        fi
    fi

    # Check and update requiretty setting
    if sudo grep -q "Defaults:$ssh_user !requiretty" /etc/sudoers; then
        echo "Requiretty is already disabled for user $ssh_user."
        CHECKLIST+=("Disable requiretty: Passed")
    else
        echo "Requiretty is not disabled for user $ssh_user."
        CHECKLIST+=("Disable requiretty: Failed")

        if [ $update_flag -eq 1 ]; then
            echo "Disabling requiretty for $ssh_user..."
            echo "Defaults:$ssh_user !requiretty" | sudo EDITOR='tee -a' visudo > /dev/null 2>&1
            if [ $? -eq 0 ]; then
                echo "Disabled requiretty for $ssh_user."
                CHECKLIST[-1]="Disable requiretty: Passed"
            else
                echo "Failed to disable requiretty for $ssh_user."
                CHECKLIST[-1]="Disable requiretty: Failed"
            fi
        fi
    fi
}

# Function to check if SSHD is running
check_sshd() {
    echo "============= Checking if SSHD is running ============="
    sshd_status=$(ps ax | grep -v grep | grep sshd)
    if [ -z "$sshd_status" ]; then
        echo "SSHD is not running."

        # Start SSHD if --update parameter is passed
        if [ $update_flag -eq 1 ]; then
            echo "Attempting to start SSHD..."
            # Start the SSH service
            sudo systemctl start ssh
            ssh_status=$?
            if [ $ssh_status -eq 0 ]; then
                echo "SSHD started successfully."
                CHECKLIST+=("SSHD start: Passed")
            else
                echo "Failed to start SSHD. "
                CHECKLIST+=("SSHD start: Failed")
            fi
        else
            CHECKLIST+=("SSHD running check: Failed")
        fi
    else
        echo "SSHD is running."
        CHECKLIST+=("SSHD running check: Passed. Try SSH from CHAI server to this server.")
    fi
}

# Function to ensure port 443 is open
check_443_firewall() {
    echo "============= Checking outgoing firewall rules for port 443 ============="
    
    # Check if the outgoing rule for port 443 already exists
    if sudo iptables -C OUTPUT -p tcp --dport 443 -j ACCEPT 2>/dev/null; then
        echo "Outgoing rule for port 443 already exists."
        CHECKLIST+=("Outbound firewall rule for Port 443: Passed")
    else
        echo "Outgoing rule for port 443 is not set."

        # Open port 443 if --update parameter is passed
        if [ $update_flag -eq 1 ]; then
            echo "Opening port 443..."
            sudo mkdir -p /etc/iptables  # Ensure the directory exists
            if sudo iptables -A OUTPUT -p tcp --dport 443 -j ACCEPT; then
                echo "Port 443 opened successfully."
                
                # Save the iptables rules to ensure they persist
                if sudo iptables-save > /etc/iptables/rules.v4; then
                    echo "Firewall rules saved to /etc/iptables/rules.v4."
                    CHECKLIST+=("Outbound firewall rule for Port 443: Passed")
                else
                    echo "Failed to save firewall rules."
                    CHECKLIST+=("Outbound firewall rule for Port 443: Failed")
                fi
            else
                echo "Failed to open port 443."
                CHECKLIST+=("Outbound firewall rule for Port 443: Failed")
            fi
        else
            CHECKLIST+=("Outbound firewall rule for Port 443: Failed")
        fi
    fi

    # Print the current status of port 443
    firewall_status=$(sudo iptables -L OUTPUT -n | grep ":443")
    if [ -z "$firewall_status" ]; then
        echo "Port 443 is not currently open."
    else
        echo "Port 443 is open."
    fi
}

# Function to check if outbound connection works
check_outbound_connection() {
    echo "============= Checking outbound connection to $url on port 443 ============="
    if curl -k --connect-timeout 5 "$url" -o /dev/null; then
        echo "Outbound connection to $url succeeded."
        CHECKLIST+=("Outbound connection to $url: Passed")
    else
        echo "Outbound connection to $url failed."
        CHECKLIST+=("Outbound connection to $url: Failed")
    fi
}

# Main script execution
#---------- Check Resources ----------
check_resources
echo

#---------- Check Sudoers ----------
check_sudoers
echo

#---------- Check SSHD ----------
check_sshd
echo

#---------- Check Firewall & Port 443 enable ----------
check_443_firewall
echo

#---------- Check if CHAI URL is reachable ----------
check_outbound_connection
echo

# Print checklist results
echo "Checklist for requirements:"
for item in "${CHECKLIST[@]}"; do
    echo "- $item"
done
echo