#!/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:
# - `standalone`: Performs HW specifications checks for a standalone server instance.
# - `cluster`: Performs HW specifications checks for a cluster server instance.
#
# The positive result of the checks do not guarantee that the server is fully compatible with Wallarm Cloud; the checks only
# indicate that the server meets basic requirements.

# Constants
WALLARM_DOMAIN="onprem.wallarm.com"
WALLARM_SUBDOMAINS="configs scripts hibp packages-versions repo registry"
HEALTH_CHECK_PATH="health"
MODE="${1:-standalone}"
TEST_DOCKER_CONTAINER_REGISTRY="registry.onprem.wallarm.com/public/alpine:latest"
SPEEDTEST_URL="https://scripts.onprem.wallarm.com/test_network_speed.tar"
OS_OK=0
if [[ "$MODE" == "standalone" ]]; then
    MIN_RAM_GB=40
    MIN_CPU_CORES=8
    MIN_DISK_ROOT_GB=200
    MIN_DISK_WALLARM_GB=500
else
    MIN_RAM_GB=64
    MIN_CPU_CORES=16
    MIN_DISK_ROOT_GB=200
    MIN_DISK_WALLARM_GB=2300
fi
BECOME_ROOT_METHOD="sudo"
CONNECT_TIMEOUT=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" =~ ^9\. ]] && ok=0
                ;;
        esac
    fi

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

find_become_root_method() {
    if command -v dzdo &>/dev/null; then
        BECOME_ROOT_METHOD="dzdo"
    elif command -v sudo &>/dev/null; then
        BECOME_ROOT_METHOD="sudo"
    else
        log_message "ERROR" "Neither sudo nor dzdo found. Cannot elevate privileges."
        exit 1
    fi
    log_message "INFO" "Using '$BECOME_ROOT_METHOD' to elevate privileges"
}

report_become_root_method() {
    echo "Using ${BECOME_ROOT_METHOD} for privilege escalation"
    echo "Full path: $(command -v $BECOME_ROOT_METHOD)"
    echo "Binary details: $(ls -la $(command -v $BECOME_ROOT_METHOD))"
}

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

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

check_https_reachability() {
    check_status=0
    for host in $WALLARM_SUBDOMAINS; do
        WALLARM_URL="https://$host.$WALLARM_DOMAIN/$HEALTH_CHECK_PATH"
        rc=1
        code=""

        if command -v curl &>/dev/null; then
            code=$(curl -sS -o /dev/null -w "%{http_code}" --connect-timeout $CONNECT_TIMEOUT "$WALLARM_URL" 2>/dev/null)
            rc=$?
        fi

        if [[ $rc -ne 0 || -z "$code" ]] && command -v wget &>/dev/null; then
            code=$(wget --server-response --spider --timeout=$CONNECT_TIMEOUT "$WALLARM_URL" 2>&1 \
                   | awk '/^  HTTP/{print $2; exit}')
            rc=$?
        fi

        # Accept 400, 402, 403, 404 as "reachable"
        if [[ $rc -ne 0 || ! "$code" =~ ^200$ ]]; then
            check_status=1
        fi
    done
    return $check_status
}

get_https_reachability() {
    check_status=0
    for host in $WALLARM_SUBDOMAINS; do
        WALLARM_URL="https://$host.$WALLARM_DOMAIN/$HEALTH_CHECK_PATH"
        log_message "INFO" "Checking HTTPS reachability for $WALLARM_URL"

        rc=1
        code=""

        if command -v curl &>/dev/null; then
            code=$(curl -sS -o /dev/null -w "%{http_code}" --connect-timeout $CONNECT_TIMEOUT "$WALLARM_URL" 2>/dev/null)
            rc=$?
        fi

        if [[ $rc -ne 0 || -z "$code" ]] && command -v wget &>/dev/null; then
            code=$(wget --server-response --spider --timeout=$CONNECT_TIMEOUT "$WALLARM_URL" 2>&1 \
                   | awk '/^  HTTP/{print $2; exit}')
            rc=$?
        fi

        if [[ $rc -ne 0 || ! "$code" =~ ^200$ ]]; then
            log_message "ERROR" "Failed to reach $WALLARM_URL (HTTP $code). Error details:"
            if command -v curl &>/dev/null; then
                curl -v --connect-timeout $CONNECT_TIMEOUT "$WALLARM_URL" || echo "No response from $WALLARM_URL"
            elif command -v wget &>/dev/null; then
                wget --spider --timeout=$CONNECT_TIMEOUT "$WALLARM_URL" || echo "No response from $WALLARM_URL"
            fi
            check_status=1
        else
            log_message "INFO" "Successfully reached $WALLARM_URL (HTTP $code)"
        fi
        echo
        print_separator
    done

    if [[ $check_status -ne 0 ]]; then
        log_message "ERROR" "One or more Wallarm subdomains are not reachable via HTTPS"
    else
        log_message "INFO" "All Wallarm subdomains are reachable via HTTPS"
    fi
}

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_gb() {
    df -BG /var/lib/wallarm-storage 2>/dev/null | awk 'NR==2 {gsub("G", "", $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_docker_pull_status() {
    if command -v docker &>/dev/null; then
        echo "Try to pull test Docker container from Wallarm repo: $TEST_DOCKER_CONTAINER_REGISTRY"
        docker pull "$TEST_DOCKER_CONTAINER_REGISTRY" && echo "Docker pull succeeded"
    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
        ${BECOME_ROOT_METHOD} -n apt-get update -qq && echo "Ubuntu repo reachable"
    elif command -v yum &>/dev/null || command -v dnf &>/dev/null; then
        ${BECOME_ROOT_METHOD} -n yum makecache -q || ${BECOME_ROOT_METHOD} -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 | grep -E '^python([0-9]|$)' | 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
}

get_etc_environments_file() {
    if [[ -f /etc/environment ]]; then
        echo "Contents of /etc/environment:"
        cat /etc/environment
    else
        echo "/etc/environment file not found"
    fi
}

list_docker_config_files() {
    if [ -d /etc/systemd/system/docker.service.d ]; then
        echo "Docker service override files in directory /etc/systemd/system/docker.service.d:"
        find /etc/systemd/system/docker.service.d -type f -exec ls -l {} \;
    else
        echo "No Docker service override files found (directory /etc/systemd/system/docker.service.d does not exist)"
    fi

    if [ -f /etc/systemd/system/docker.service.d/proxy.conf ]; then
        echo "Docker proxy configuration file /etc/systemd/system/docker.service.d/proxy.conf found:"
        cat /etc/systemd/system/docker.service.d/proxy.conf
    else
        echo "No Docker proxy configuration file /etc/systemd/system/docker.service.d/proxy.conf found"
    fi
}

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
        ${BECOME_ROOT_METHOD} dd if=/dev/zero of="$mount_point/test_io.tmp" bs=1M count=1000 oflag=direct seek=$RANDOM 2>&1 | grep -oP '[0-9]+(\.\d+)? [MG]B/s' | awk '{if ($2 == "GB/s") printf("%.0f MB/s\n", $1 * 1024); else print $1, $2}' | tr -d 'MB/s'
        ${BECOME_ROOT_METHOD} rm -f "$mount_point/test_io.tmp"
    else
        echo "$mount_point not mounted or not found"
    fi
}

get_running_processes() {
    ps aux --sort=-%mem
}

get_enabled_services() {
    ${BECOME_ROOT_METHOD} systemctl list-unit-files --type=service --state=enabled --no-pager
}

get_installed_packages() {
    if command -v dpkg &>/dev/null; then
        dpkg -l --no-pager
    elif command -v rpm &>/dev/null; then
        rpm -qa
    else
        echo "No package manager found"
    fi
}

get_root_volumes() {
    # Print header
    printf "%-25s %-10s %-30s\n" "DEVICE" "SIZE" "MOUNTPOINT"
    printf "%-25s %-10s %-30s\n" "------" "----" "----------"

    df -h --output=source,size,target \
        | tail -n +2 \
        | grep -E -v 'tmpfs|devtmpfs|overlay|shm|none|efivarfs' \
        | while read -r fs size mount; do
            # Only include mounts under / but not /boot
            if [[ "$mount" == /* && "$mount" != "/boot"* ]]; then
                printf "%-25s %-10s %-30s\n" "$fs" "$size" "$mount"
            fi
        done
}


speedtest_download() {
    local url=$SPEEDTEST_URL
    local timeout_sec=120
    local tmpfile="/tmp/wallarm_net_speedtest"
    local tool=""

    # Check for curl or wget
    if command -v curl >/dev/null 2>&1; then
        tool="curl"
    elif command -v wget >/dev/null 2>&1; then
        tool="wget"
    else
        echo "Error: neither curl nor wget is available."
        return 1
    fi

    echo "Starting network speed test to:"
    echo "  $url"
    echo "The test will run up to ${timeout_sec} seconds, attempting ~1GB download."
    echo

    # Record start time
    local start_ts=$(date +%s)

    if [ "$tool" = "curl" ]; then
        timeout "$timeout_sec" curl -L -o "$tmpfile" --progress-bar "$url"
    else
        timeout "$timeout_sec" wget  "$url" -O "$tmpfile"
    fi
    status=$?

    # Record end time
    local end_ts=$(date +%s)
    local elapsed=$((end_ts - start_ts))
    [ $elapsed -eq 0 ] && elapsed=1   # avoid division by zero

    if [ $status -eq 124 ]; then
        echo "Download timed out after ${timeout_sec} seconds. Test not completed successfully."
        rm -f "$tmpfile"
        return 1
    elif [ $status -ne 0 ]; then
        echo "Download failed (exit code $status)."
        rm -f "$tmpfile"
        return 1
    fi

    # Measure bytes downloaded
    local bytes=$(stat -c%s "$tmpfile" 2>/dev/null || stat -f%z "$tmpfile")
    echo "Downloaded $bytes bytes in $elapsed seconds."
    rm -f "$tmpfile"

    # Calculate throughput
    local bits=$((bytes * 8))
    local bps=$((bits / elapsed))  # bits per second
    local mbps=$(awk "BEGIN {printf \"%.2f\", $bps/1000000}")

    echo "Average download speed: $mbps Mbps"
}

run_check() {
    print_separator
    echo "                                                CHECK RESULTS (mode = $MODE)"
    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 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_gb)
    wallarm_free=${wallarm_free:-0}
    local wallarm_free_gb=$wallarm_free
    local min_wallarm_gb=$MIN_DISK_WALLARM_GB
    local disk_expected="root>=${MIN_DISK_ROOT_GB}GB or wallarm>=${MIN_DISK_WALLARM_GB}GB"
    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}GB"
    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="OK"
    check_https_reachability && result="OK" || result="FAILED"
    printf "%-${WIDTH1}s %-${WIDTH2}s %-${WIDTH3}s %-${WIDTH4}s\n" "HTTPS to Wallarm" "$result" "$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" "Root Volumes:"
    get_root_volumes
    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" "Privilege Escalation Method:"
    report_become_root_method
    print_separator

    log_message "INFO" "Docker Status:"
    get_docker_status
    print_separator

    log_message "INFO" "HTTPS Reachability to Wallarm URLs:"
    get_https_reachability
    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" "Content of /etc/environment:"
    get_etc_environments_file
    print_separator

    log_message "INFO" "Docker Service Override Files:"
    list_docker_config_files
    print_separator

    log_message "INFO" "Docker Pull Test:"
    get_docker_pull_status
    print_separator

    log_message "INFO" "Running Processes (sorted by memory usage):"
    get_running_processes
    print_separator

    log_message "INFO" "Enabled Services:"
    get_enabled_services
    print_separator

    log_message "INFO" "Installed Packages:"
    get_installed_packages
    print_separator

    log_message "INFO" "Network Speed Test:"
    speedtest_download
    print_separator
}

main() {
    validate_os

    if ! command -v curl &>/dev/null && ! command -v wget &>/dev/null; then
        log_message "ERROR" "Neither curl nor wget found. Cannot test HTTPS reachability."
        return 1
    fi

    find_become_root_method

    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

    if [[ "$MODE" != "standalone" && "$MODE" != "cluster" ]]; then
        log_message "ERROR" "Unsupported mode: $MODE. Use 'standalone' or 'cluster'."
        exit 1
    fi

    run_report
}

main