#!/bin/bash

################################################################################
# cPHulk Enablement Script with CSF Whitelist Import
################################################################################
# Purpose: Enable cPHulk brute force protection and import CSF allowed IPs
# Requirements: cPanel server with CSF installed
# Author: Server Toolkit
#
# IMPORTANT NOTES:
# - cPHulk operates SYSTEM-WIDE (not per-user or per-domain)
# - Protects all authentication services: cPanel, WHM, SSH, FTP, Email
# - This script intelligently discovers ALL CSF whitelist files including:
#   * Standard files (/etc/csf/csf.allow, cpanel.allow, etc.)
#   * Include directives (follows recursively)
#   * Hosting panel integrations (LiquidWeb, cPanel, InterWorx, etc.)
#   * Custom locations configured in csf.conf
# - Supports multiple IP formats: simple IPs, s=IP, d=IP, CIDR notation
################################################################################

# Get script directory
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
source "$SCRIPT_DIR/lib/common-functions.sh"
source "$SCRIPT_DIR/lib/system-detect.sh"

# Require root
if [ "$EUID" -ne 0 ]; then
    print_error "This script must be run as root"
    exit 1
fi

print_banner "cPHulk Enablement with CSF Whitelist Import"

# Check if cPanel
if [ "$SYS_CONTROL_PANEL" != "cpanel" ]; then
    print_error "This script is for cPanel servers only"
    print_info "Detected control panel: ${SYS_CONTROL_PANEL:-none}"
    exit 1
fi

# Check if cPHulk exists
if [ ! -x "/usr/local/cpanel/bin/cphulk_pam_ctl" ]; then
    print_error "cPHulk not found. This may not be a standard cPanel installation."
    exit 1
fi

# Check if CSF is installed
if [ ! -x "/usr/sbin/csf" ]; then
    print_warning "CSF not found - will skip CSF whitelist import"
    CSF_AVAILABLE=false
else
    CSF_AVAILABLE=true
fi

print_section "Current cPHulk Status"

# Get current status
CPHULK_STATUS=$(/usr/local/cpanel/bin/cphulk_pam_ctl --status 2>/dev/null)
if echo "$CPHULK_STATUS" | grep -qi "enabled"; then
    print_success "cPHulk is currently ENABLED"
    ALREADY_ENABLED=true
else
    print_info "cPHulk is currently DISABLED"
    ALREADY_ENABLED=false
fi

# Show current whitelist count
CURRENT_WHITELIST=$(/usr/local/cpanel/scripts/cphulkdwhitelist --list 2>/dev/null | grep -v "^$" | wc -l)
print_info "Current cPHulk whitelist entries: $CURRENT_WHITELIST"

if [ "$CSF_AVAILABLE" = true ]; then
    print_section "CSF Whitelist Analysis"

    # Get CSF allowed IPs
    echo "Scanning CSF allow files and following Include directives..."
    echo ""

    # Start with main CSF allow file and follow all includes recursively
    declare -A PROCESSED_FILES  # Track processed files to avoid duplicates
    CSF_FILES=()

    # Function to recursively find included files
    find_csf_files() {
        local current_file="$1"

        # Skip if already processed or doesn't exist
        [ -n "${PROCESSED_FILES[$current_file]}" ] && return
        [ ! -f "$current_file" ] && return

        # Mark as processed
        PROCESSED_FILES["$current_file"]=1
        CSF_FILES+=("$current_file")

        # Look for Include directives
        while IFS= read -r line; do
            if [[ "$line" =~ ^Include[[:space:]]+(.+)$ ]]; then
                local included_file="${BASH_REMATCH[1]}"
                # Trim whitespace
                included_file=$(echo "$included_file" | xargs)
                # Recursively process included file
                find_csf_files "$included_file"
            fi
        done < "$current_file"
    }

    # Method 1: Start with main csf.allow and follow all includes
    if [ -f "/etc/csf/csf.allow" ]; then
        find_csf_files "/etc/csf/csf.allow"
    fi

    # Method 2: Scan for all .allow files in /etc/csf/ (catches files not included)
    while IFS= read -r file; do
        find_csf_files "$file"
    done < <(find /etc/csf/ -maxdepth 1 -type f -name "*.allow" 2>/dev/null)

    # Method 3: Check csf.conf for ALLOWIPS_INCLUDE directive
    # This directive can point to custom include directories
    if [ -f "/etc/csf/csf.conf" ]; then
        ALLOWIPS_INCLUDE=$(grep "^ALLOWIPS_INCLUDE" /etc/csf/csf.conf | cut -d'=' -f2 | tr -d '"' | tr -d ' ')
        if [ -n "$ALLOWIPS_INCLUDE" ] && [ -d "$ALLOWIPS_INCLUDE" ]; then
            echo "Found ALLOWIPS_INCLUDE directory: $ALLOWIPS_INCLUDE"
            while IFS= read -r file; do
                find_csf_files "$file"
            done < <(find "$ALLOWIPS_INCLUDE" -type f 2>/dev/null)
        fi
    fi

    # Method 4: Common custom CSF include locations (hosting panel integrations)
    COMMON_CSF_LOCATIONS=(
        "/usr/local/lp/etc/csf"           # LiquidWeb
        "/usr/local/cpanel/etc/csf"        # cPanel custom
        "/usr/local/interworx/etc/csf"     # InterWorx
        "/var/cpanel/csf"                  # cPanel var
        "/root/.csf"                       # Custom user location
    )

    for location in "${COMMON_CSF_LOCATIONS[@]}"; do
        if [ -d "$location" ]; then
            # Look for all .allow files in these locations
            while IFS= read -r file; do
                find_csf_files "$file"
            done < <(find "$location" -type f -name "*allow*" 2>/dev/null)
        fi
    done

    # Method 5: Parse all Include directives from csf.conf CUSTOM variables
    if [ -f "/etc/csf/csf.conf" ]; then
        # Some systems use CUSTOM1_LOG through CUSTOM9_LOG which may reference allow files
        while IFS= read -r include_path; do
            if [ -f "$include_path" ]; then
                find_csf_files "$include_path"
            fi
        done < <(grep -E "^(CUSTOM|GENERIC).*Include" /etc/csf/csf.conf 2>/dev/null | grep -oE '\/[^"]+' || true)
    fi

    if [ ${#CSF_FILES[@]} -eq 0 ]; then
        print_warning "No CSF .allow files found"
        CSF_AVAILABLE=false
    else
        echo "Found ${#CSF_FILES[@]} CSF allow files (including includes):"
        for file in "${CSF_FILES[@]}"; do
            if [[ "$file" == /etc/csf/* ]]; then
                echo "  • $(basename "$file")"
            else
                echo "  • $file"
            fi
        done
        echo ""
    fi

    # Extract IPs from all CSF files
    CSF_ALLOW_IPS=()
    declare -A IP_SOURCE_COUNT  # Track which files contributed IPs

    if [ ${#CSF_FILES[@]} -gt 0 ]; then
        for csf_file in "${CSF_FILES[@]}"; do
            file_ip_count=0
            file_display_name=""

            if [[ "$csf_file" == /etc/csf/* ]]; then
                file_display_name="$(basename "$csf_file")"
            else
                file_display_name="$csf_file"
            fi

            while IFS= read -r line; do
                # Skip comments and empty lines
                [[ "$line" =~ ^#.*$ ]] && continue
                [[ -z "$line" ]] && continue

                # Skip Include directives (already processed)
                [[ "$line" =~ ^Include ]] && continue

                # Try multiple IP extraction methods:
                ip=""

                # Method 1: Simple IP format (IP at start of line)
                # Example: 192.168.100.1 # comment
                # Also handles CIDR: 10.20.4.0/22 # comment (strip /CIDR, just get IP)
                if [ -z "$ip" ]; then
                    ip=$(echo "$line" | awk '{print $1}' | grep -oE '^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}')
                fi

                # Method 2: CSF source format (s=IP)
                # Example: tcp|in|d=4|s=208.74.123.2 # cPanel Auth Server
                if [ -z "$ip" ]; then
                    ip=$(echo "$line" | grep -oE 's=[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}' | cut -d'=' -f2)
                fi

                # Method 3: CSF destination format (d=IP)
                # Example: tcp|out|d=443|d=45.11.128.61/32
                if [ -z "$ip" ]; then
                    # Extract last d= parameter that contains an IP (not a port)
                    ip=$(echo "$line" | grep -oE 'd=[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}' | tail -1 | cut -d'=' -f2)
                fi

                if [ -n "$ip" ]; then
                    # Check if not already in array (deduplicate)
                    if [[ ! " ${CSF_ALLOW_IPS[@]} " =~ " ${ip} " ]]; then
                        CSF_ALLOW_IPS+=("$ip")
                        file_ip_count=$((file_ip_count + 1))
                    fi
                fi
            done < "$csf_file"

            # Track count per file
            if [ $file_ip_count -gt 0 ]; then
                IP_SOURCE_COUNT["$file_display_name"]=$file_ip_count
            fi
        done

        # Show breakdown by file
        echo "IP breakdown by source file:"
        for file in $(printf '%s\n' "${!IP_SOURCE_COUNT[@]}" | sort); do
            printf "  • %-45s %3d IPs\n" "$file" "${IP_SOURCE_COUNT[$file]}"
        done
        echo ""

        print_success "Found ${#CSF_ALLOW_IPS[@]} total unique IPs across all CSF files"
    fi

    if [ ${#CSF_ALLOW_IPS[@]} -gt 0 ]; then
        echo ""
        echo "Sample IPs to be imported:"
        for i in {0..4}; do
            if [ -n "${CSF_ALLOW_IPS[$i]}" ]; then
                echo "  • ${CSF_ALLOW_IPS[$i]}"
            fi
        done
        if [ ${#CSF_ALLOW_IPS[@]} -gt 5 ]; then
            echo "  ... and $((${#CSF_ALLOW_IPS[@]} - 5)) more"
        fi
    fi
fi

echo ""
print_section "Actions to Perform"

if [ "$ALREADY_ENABLED" = false ]; then
    echo "  1. Enable cPHulk brute force protection"
else
    echo "  1. cPHulk already enabled (skip)"
fi

if [ "$CSF_AVAILABLE" = true ] && [ ${#CSF_ALLOW_IPS[@]} -gt 0 ]; then
    echo "  2. Import ${#CSF_ALLOW_IPS[@]} IPs from CSF to cPHulk whitelist"
else
    echo "  2. No CSF IPs to import (skip)"
fi

echo "  3. Display final configuration"

echo ""
read -p "Proceed with these actions? (yes/no): " confirm

if [ "$confirm" != "yes" ]; then
    print_info "Operation cancelled"
    exit 0
fi

echo ""
print_section "Execution"

# Step 1: Enable cPHulk
if [ "$ALREADY_ENABLED" = false ]; then
    print_info "Enabling cPHulk..."
    if /usr/local/cpanel/bin/cphulk_pam_ctl --enable 2>&1; then
        print_success "cPHulk enabled successfully"
    else
        print_error "Failed to enable cPHulk"
        exit 1
    fi
else
    print_info "cPHulk already enabled, skipping"
fi

# Step 2: Import CSF whitelist
if [ "$CSF_AVAILABLE" = true ] && [ ${#CSF_ALLOW_IPS[@]} -gt 0 ]; then
    print_info "Importing CSF whitelist to cPHulk..."

    IMPORTED=0
    SKIPPED=0
    FAILED=0

    for ip in "${CSF_ALLOW_IPS[@]}"; do
        # Check if already in cPHulk whitelist
        if /usr/local/cpanel/scripts/cphulkdwhitelist --list 2>/dev/null | grep -q "$ip"; then
            SKIPPED=$((SKIPPED + 1))
            echo "  [SKIP] $ip (already whitelisted)"
        else
            # Add to cPHulk whitelist
            if whmapi1 cphulkd_add_whitelist ip="$ip" 2>&1 | grep -q "success.*1"; then
                IMPORTED=$((IMPORTED + 1))
                echo "  [OK] $ip"
            else
                FAILED=$((FAILED + 1))
                echo "  [FAIL] $ip"
            fi
        fi
    done

    echo ""
    print_success "Import complete:"
    echo "  • Imported: $IMPORTED"
    echo "  • Skipped (already whitelisted): $SKIPPED"
    if [ $FAILED -gt 0 ]; then
        print_warning "Failed: $FAILED"
    fi
fi

# Step 3: Display final status
echo ""
print_section "Final Configuration"

# Check status
FINAL_STATUS=$(/usr/local/cpanel/bin/cphulk_pam_ctl --status 2>/dev/null)
if echo "$FINAL_STATUS" | grep -qi "enabled"; then
    print_success "cPHulk Status: ENABLED"
else
    print_error "cPHulk Status: DISABLED (unexpected)"
fi

# Count whitelist
FINAL_WHITELIST=$(/usr/local/cpanel/scripts/cphulkdwhitelist --list 2>/dev/null | grep -v "^$" | wc -l)
print_info "cPHulk whitelist entries: $FINAL_WHITELIST"

echo ""
print_section "Next Steps"

echo "1. Configure cPHulk settings in WHM:"
echo "   WHM → Security Center → cPHulk Brute Force Protection"
echo ""
echo "2. Recommended settings:"
echo "   • Brute Force Protection Period: 5 minutes"
echo "   • Maximum Failures per Account: 5"
echo "   • Maximum Failures per IP: 10"
echo ""
echo "3. Add your own IPs to whitelist:"
echo "   whmapi1 cphulkd_add_whitelist ip=YOUR.IP.ADDRESS"
echo ""
echo "4. View currently blocked IPs:"
echo "   whmapi1 cphulkd_list_blocks"
echo ""
echo "5. Remove a blocked IP:"
echo "   whmapi1 cphulkd_remove_block ip=IP.TO.UNBLOCK"

echo ""
print_success "cPHulk setup complete!"