#!/bin/bash
#
# Wallarm Cloud server instance light-weight check script.
#
# The script performs a series of checks on the server to ensure it meets the basic requirements for running Wallarm Cloud.
# It checks the operating system, RAM, CPU, disk space, user permissions, network connectivity, and other system parameters.
# The script can operate in two modes:
# - `check`: Performs a quick check and outputs the results in a table format.
# - `report`: Performs a detailed report with additional system information and diagnostics.
#

# Constants
WALLARM_DOMAIN="onprem.wallarm.com"
WALLARM_URL="https://${WALLARM_DOMAIN}"
MODE="${1:-check}"
OS_OK=0
MIN_RAM_GB=64
MIN_CPU_CORES=16
MIN_DISK_ROOT_GB=200
MIN_DISK_WALLARM_TB=2.3

# Logging function
log_message() {
    local severity="$1"
    local message="$2"
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] [$severity] $message"
}

print_separator() {
    echo "==================================================================================================================================="
}

validate_os() {
    local os=Unknown ok=1
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        os="${ID}_${VERSION_ID}"
        case "$ID" in
            ubuntu)
                [[ "$VERSION_ID" == "22.04" ]] && ok=0
                ;;
            rhel)
                [[ "$VERSION_ID" =~ ^8\.|^9\. ]] && ok=0
                ;;
        esac
    fi

    if [[ $ok -ne 0 ]]; then
        log_message "ERROR" "Unsupported OS: $os"
        exit 1
    fi
}

is_not_root() {
    [[ "$(id -u)" -ne 0 ]]
}

can_sudo_passwordless() {
    sudo -n true 2>/dev/null
}

check_https_reachability() {
    curl -sS --connect-timeout 10 "$WALLARM_URL" -o /dev/null || wget --quiet --spider "$WALLARM_URL"
}

get_os_info() {
    . /etc/os-release
    echo "$PRETTY_NAME"
}

get_kernel_version() {
    uname -r
}

get_cpu_arch() {
    uname -m
}

get_cpu_cores() {
    nproc
}

get_total_ram_gb() {
    free --giga | awk '/^Mem:/{print $2}'
}

get_disk_usage() {
    df -h
}

get_free_root_gb() {
    df -BG / | awk 'NR==2 {gsub("G", "", $4); print $4}'
}

get_free_wallarm_tb() {
    df -BT /var/lib/wallarm-storage 2>/dev/null | awk 'NR==2 {gsub("T", "", $4); print $4}'
}

get_docker_status() {
    if command -v docker &>/dev/null; then
        echo "Docker service status: $(systemctl is-active docker || echo inactive)"
        docker ps --format "table {{.ID}}\t{{.Image}}\t{{.Status}}" || echo "No running containers"
    else
        echo "Docker not installed"
    fi
}

get_package_repos() {
    if command -v apt &>/dev/null; then
        shopt -s nullglob
        local files=(/etc/apt/sources.list /etc/apt/sources.list.d/*.list)
        grep -h ^deb "${files[@]}" 2>/dev/null || echo "No APT sources found"
    elif command -v yum &>/dev/null || command -v dnf &>/dev/null; then
        yum repolist all 2>/dev/null || dnf repolist all 2>/dev/null || echo "No YUM/DNF repos found"
    fi
}

can_install_packages() {
    if command -v apt-get &>/dev/null; then
        sudo -n apt-get update -qq && echo "Ubuntu repo reachable"
    elif command -v yum &>/dev/null || command -v dnf &>/dev/null; then
        sudo -n yum makecache -q || sudo -n dnf makecache -q && echo "RHEL repo reachable"
    fi
}

get_python_versions() {
    log_message "INFO" "Detected Python binaries and versions:"

    # Use 'command -v' to find full paths of executables
    local found=0
    for bin in $(compgen -c python | sort -u); do
        local path
        path=$(command -v "$bin" 2>/dev/null)
        if [[ -n "$path" && -x "$path" ]]; then
            # Try to get version (suppress stderr)
            local version
            version=$("$path" --version 2>&1)
            printf "  %-20s -> %-50s\n" "$version" "$path"
            found=1
        fi
    done

    [[ $found -eq 0 ]] && echo "  No Python binaries found in PATH"
}


get_proxy_config() {
    env | grep -i proxy || echo "No proxy configured"
}

get_fips_status() {
    [[ -f /proc/sys/crypto/fips_enabled ]] && cat /proc/sys/crypto/fips_enabled || echo "FIPS not available"
}

get_selinux_status() {
    if command -v getenforce &>/dev/null; then
        getenforce
    else
        echo "SELinux tools not available"
    fi
}

check_ca_trust() {

    log_message INFO "Checking CA trust for ${WALLARM_DOMAIN}..."
    echo | openssl s_client -connect ${WALLARM_DOMAIN}:443 </dev/null
    echo ""
    echo | openssl s_client -connect ${WALLARM_DOMAIN}:443 </dev/null 2>/dev/null | grep -q "Verify return code: 0" \
        && echo "Trusted" || echo "Untrusted"
}

test_disk_io() {
    local mount_point="$1"
    local label="$2"

    if mountpoint -q "$mount_point"; then
        log_message "INFO" "Disk I/O Test ($label - $mount_point)"
        # Use a small file and sync to flush
        sudo dd if=/dev/zero of="$mount_point/test_io.tmp" bs=64k count=256 conv=fdatasync oflag=direct 2>&1 | grep -E 'copied|bytes'
        sudo rm -f "$mount_point/test_io.tmp"
    else
        echo "$mount_point not mounted or not found"
    fi
}


run_check() {
    print_separator
    echo "                                                           CHECK RESULTS"
    print_separator

    local WIDTH1=35
    local WIDTH2=40
    local WIDTH3=40
    local WIDTH4=10

    printf "%-${WIDTH1}s %-${WIDTH2}s %-${WIDTH3}s %-${WIDTH4}s\n" "Check" "Actual" "Expected" "Status"
    printf "%-${WIDTH1}s %-${WIDTH2}s %-${WIDTH3}s %-${WIDTH4}s\n" "-----" "------" "--------" "------"

    # OS check
    local os=$(get_os_info)
    local os_expected="Ubuntu 22.04 / RHEL 8/9"
    result="OK"
    [[ "$os" == *"Ubuntu 22.04"* || "$os" == *"Red Hat Enterprise Linux"* ]] || result="FAILED"
    printf "%-${WIDTH1}s %-${WIDTH2}s %-${WIDTH3}s %-${WIDTH4}s\n" "OS Version" "$os" "$os_expected" "$result"

    # RAM check
    local ram=$(get_total_ram_gb)
    local ram_expected=">= ${MIN_RAM_GB} GB"
    result="OK"; (( ram < MIN_RAM_GB )) && result="FAILED"
    printf "%-${WIDTH1}s %-${WIDTH2}s %-${WIDTH3}s %-${WIDTH4}s\n" "Total RAM (GB)" "${ram} GB" "$ram_expected" "$result"

    # CPU check
    local cpu=$(get_cpu_cores)
    local cpu_expected=">= ${MIN_CPU_CORES}"
    result="OK"; (( cpu < MIN_CPU_CORES )) && result="FAILED"
    printf "%-${WIDTH1}s %-${WIDTH2}s %-${WIDTH3}s %-${WIDTH4}s\n" "CPU Cores" "$cpu" "$cpu_expected" "$result"

    # Disk check
    local root_free=$(get_free_root_gb)
    local wallarm_free=$(get_free_wallarm_tb)
    wallarm_free=${wallarm_free:-0}
    local wallarm_free_gb=$(awk "BEGIN { printf \"%d\", $wallarm_free * 1024 }")
    local min_wallarm_gb=$(awk "BEGIN { printf \"%d\", $MIN_DISK_WALLARM_TB * 1024 }")
    local disk_expected="root>=${MIN_DISK_ROOT_GB}GB or wallarm>=${MIN_DISK_WALLARM_TB}TB"
    result="FAILED"

    if (( root_free >= MIN_DISK_ROOT_GB )); then
        if (( root_free >= min_wallarm_gb )); then
            result="OK"
        elif (( wallarm_free_gb >= min_wallarm_gb )); then
            result="OK"
        fi
    fi

    local disk_actual="root=${root_free}GB, wallarm=${wallarm_free}TB"
    printf "%-${WIDTH1}s %-${WIDTH2}s %-${WIDTH3}s %-${WIDTH4}s\n" "Disk Space" "$disk_actual" "$disk_expected" "$result"

    # Root check
    local user_actual=$(whoami)
    local user_expected="Not root"
    is_not_root && result="OK" || result="FAILED"
    printf "%-${WIDTH1}s %-${WIDTH2}s %-${WIDTH3}s %-${WIDTH4}s\n" "User is not root" "$user_actual" "$user_expected" "$result"

    # Sudo check
    local sudo_expected="Passwordless sudo"
    can_sudo_passwordless && result="OK" || result="FAILED"
    printf "%-${WIDTH1}s %-${WIDTH2}s %-${WIDTH3}s %-${WIDTH4}s\n" "Sudo passwordless" "$user_actual" "$sudo_expected" "$result"

    # HTTPS check
    local https_expected="HTTPS reachable (403 is OK)"
    check_https_reachability && result="OK" || result="FAILED"
    printf "%-${WIDTH1}s %-${WIDTH2}s %-${WIDTH3}s %-${WIDTH4}s\n" "HTTPS to Wallarm" "Reachable" "$https_expected" "$result"
}

run_report() {
    run_check

    print_separator
    echo "                                                          DETAILED REPORT"
    print_separator

    log_message "INFO" "OS Info: $(get_os_info)"
    print_separator

    log_message "INFO" "Kernel Version: $(get_kernel_version)"
    print_separator

    log_message "INFO" "CPU Architecture: $(get_cpu_arch)"
    print_separator

    log_message "INFO" "CPU Cores: $(get_cpu_cores)"
    print_separator

    log_message "INFO" "Total RAM (GB): $(get_total_ram_gb)"
    print_separator

    log_message "INFO" "Disk Usage:"
    get_disk_usage
    print_separator

    log_message "INFO" "Disk I/O Performance (/)"
    test_disk_io "/" "root"
    print_separator

    log_message "INFO" "Disk I/O Performance (/var/lib/wallarm-storage)"
    test_disk_io "/var/lib/wallarm-storage" "wallarm-storage"
    print_separator

    log_message "INFO" "Docker Status:"
    get_docker_status
    print_separator

    log_message "INFO" "HTTPS Reachability to ${WALLARM_URL}:"
    check_https_reachability && echo "OK" || echo "Failed"
    print_separator

    log_message "INFO" "Package Repositories:"
    get_package_repos
    print_separator

    log_message "INFO" "Package Install Test:"
    can_install_packages
    print_separator

    log_message "INFO" "Python Versions:"
    get_python_versions
    print_separator

    log_message "INFO" "Proxy Configuration:"
    get_proxy_config
    print_separator

    log_message "INFO" "FIPS Status:"
    get_fips_status
    print_separator

    log_message "INFO" "SELinux Status:"
    get_selinux_status
    print_separator

    log_message "INFO" "CA Trust Check for SSL on domain ${WALLARM_DOMAIN}:"
    check_ca_trust
    print_separator
}

main() {
    validate_os

    if ! is_not_root; then
        log_message "ERROR" "Script must not be run as root"
        exit 1
    fi

    if ! can_sudo_passwordless; then
        log_message "ERROR" "User cannot perform passwordless sudo"
        exit 1
    fi

    case "$MODE" in
        check)
            run_check
            ;;
        report)
            run_report
            ;;
        *)
            log_message "ERROR" "Unsupported mode: $MODE. Use 'check' or 'report'."
            exit 1
            ;;
    esac
}

main