#!/bin/bash

################################################################################
# Malware Scanner
################################################################################
# Purpose: Comprehensive malware scanning with multiple engines
# Supports: ImunifyAV, ClamAV, Maldet (LMD)
# Scan scope: Single domain, user account, or entire server
################################################################################

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
source "$SCRIPT_DIR/lib/common-functions.sh" 2>/dev/null || true
source "$SCRIPT_DIR/lib/system-detect.sh" 2>/dev/null || true
source "$SCRIPT_DIR/lib/user-manager.sh" 2>/dev/null || true

# Arrays for docroots and scanners
declare -a docroot_array
declare -a sanitized_docroot
declare -a remove_docroot
declare -a available_scanners

# Scanner detection
detect_scanners() {
    available_scanners=()

    if command -v imunify-antivirus &>/dev/null; then
        available_scanners+=("imunify")
    fi

    if command -v clamscan &>/dev/null; then
        available_scanners+=("clamav")
    fi

    if command -v maldet &>/dev/null; then
        available_scanners+=("maldet")
    fi

    if [ ${#available_scanners[@]} -eq 0 ]; then
        echo -e "${RED}No malware scanners detected!${NC}"
        echo ""
        echo "Available scanners to install:"
        echo "  • ImunifyAV - Commercial, real-time protection"
        echo "  • ClamAV - Open source antivirus"
        echo "  • Maldet (LMD) - Linux Malware Detect"
        echo ""
        return 1
    fi

    return 0
}

# Detect control panel and gather docroots
detect_control_panel() {
    docroot_array=()

    # Detect cPanel
    if [ -f "/etc/userdatadomains" ]; then
        CONTROL_PANEL="cpanel"
        export PATH=/usr/local/cpanel/3rdparty/bin/:$PATH

        while IFS= read -r docroot; do
            [ -n "$docroot" ] && docroot_array+=("$docroot")
        done < <(cut -d= -f5 /etc/userdatadomains | sed 's/==/=/g' | sort -u)

    # Detect Plesk
    elif [ -f "/usr/local/psa/version" ]; then
        CONTROL_PANEL="plesk"

        while IFS= read -r domain; do
            docroot=$(plesk bin site -i "$domain" 2>/dev/null | grep "WWW-Root" | awk '{print $2}')
            [ -n "$docroot" ] && docroot_array+=("$docroot")
        done < <(plesk bin site --list 2>/dev/null)

    # Detect Interworx
    elif [ -d "/usr/local/interworx/" ]; then
        CONTROL_PANEL="interworx"

        while IFS= read -r docroot; do
            [ -n "$docroot" ] && docroot_array+=("$docroot")
        done < <(grep -rh "DocumentRoot" /etc/httpd/conf* 2>/dev/null | grep -Ev '^\s*#|/var/www/html($|$)' | sed 's/DocumentRoot//g' | tr -d " " | sort -u)

    else
        CONTROL_PANEL="none"
        echo -e "${YELLOW}No control panel detected${NC}"
        echo "Manual path selection required"
        return 1
    fi

    # Remove subdirectory docroots (avoid scanning same files twice)
    sanitize_docroots

    return 0
}

# Remove subdirectory docroots from array
sanitize_docroots() {
    remove_docroot=()

    for search_value in "${docroot_array[@]}"; do
        # Count how many paths contain this value
        count=$(printf '%s\n' "${docroot_array[@]}" | grep -c "$search_value" || true)

        if [ "$count" -gt 1 ]; then
            # Find subdirectories and mark for removal
            while IFS= read -r subdir; do
                if [ "$subdir" != "$search_value" ]; then
                    remove_docroot+=("$subdir")
                fi
            done < <(printf '%s\n' "${docroot_array[@]}" | grep "$search_value")
        fi
    done

    # Build sanitized array
    sanitized_docroot=()
    for docroot in "${docroot_array[@]}"; do
        # Check if this docroot is in remove list
        skip=0
        for remove in "${remove_docroot[@]}"; do
            if [ "$docroot" = "$remove" ]; then
                skip=1
                break
            fi
        done

        if [ $skip -eq 0 ]; then
            sanitized_docroot+=("$docroot")
        fi
    done
}

# Get docroots for specific user
get_user_docroots() {
    local username="$1"
    local user_docroots=()

    if [ "$CONTROL_PANEL" = "cpanel" ]; then
        while IFS= read -r docroot; do
            [ -n "$docroot" ] && user_docroots+=("$docroot")
        done < <(grep "^${username}:" /etc/userdatadomains | cut -d= -f5 | sed 's/==/=/g' | sort -u)
    else
        echo -e "${RED}User-specific scanning only supported on cPanel${NC}"
        return 1
    fi

    echo "${user_docroots[@]}"
}

# Get docroot for specific domain
get_domain_docroot() {
    local domain="$1"
    local domain_docroot=""

    if [ "$CONTROL_PANEL" = "cpanel" ]; then
        domain_docroot=$(grep "^${domain}:" /etc/userdatadomains | cut -d= -f5 | sed 's/==/=/g')
    elif [ "$CONTROL_PANEL" = "plesk" ]; then
        domain_docroot=$(plesk bin site -i "$domain" 2>/dev/null | grep "WWW-Root" | awk '{print $2}')
    else
        echo -e "${RED}Domain lookup only supported on cPanel/Plesk${NC}"
        return 1
    fi

    echo "$domain_docroot"
}

# Memory check before scanning
check_memory() {
    local total_mem=$(free -m | awk '/^Mem:/{print $2}')
    local avail_mem=$(free -m | awk '/^Mem:/{print $7}')

    local min_total=2048  # 2GB
    local min_avail=512   # 512MB

    if [ "$total_mem" -lt "$min_total" ] || [ "$avail_mem" -lt "$min_avail" ]; then
        echo -e "${YELLOW}WARNING: Low memory detected${NC}"
        echo "Total: ${total_mem}MB | Available: ${avail_mem}MB"
        echo ""
        echo "Running a full scan may cause high load or OOM conditions."
        echo ""
        read -p "Continue anyway? (yes/no): " confirm

        if [ "$confirm" != "yes" ]; then
            echo "Scan cancelled"
            return 1
        fi
    fi

    return 0
}

# ImunifyAV scanner
scan_imunify() {
    local scan_paths=("$@")

    if ! command -v imunify-antivirus &>/dev/null; then
        echo -e "${RED}ImunifyAV not installed${NC}"
        return 1
    fi

    echo -e "${CYAN}Starting ImunifyAV scan...${NC}"
    echo ""

    # Update signatures
    echo "→ Updating signatures..."
    imunify-antivirus update 2>/dev/null

    # Queue scan paths
    for path in "${scan_paths[@]}"; do
        if [ -d "$path" ]; then
            echo "→ Queuing: $path"
            imunify-antivirus malware on-demand queue put "$path" 2>/dev/null
        fi
    done

    echo ""
    echo -e "${GREEN}✓ Scan queued${NC}"
    echo ""
    echo "Monitor progress:"
    echo "  imunify-antivirus malware on-demand list"
}

# ClamAV scanner
scan_clamav() {
    local scan_paths=("$@")

    if ! command -v clamscan &>/dev/null; then
        echo -e "${RED}ClamAV not installed${NC}"
        return 1
    fi

    # Create log directory
    local log_dir="$SCRIPT_DIR/logs/malware-scans"
    mkdir -p "$log_dir"

    local log_file="$log_dir/clamav_$(date +%Y%m%d_%H%M%S).log"

    echo -e "${CYAN}Starting ClamAV scan...${NC}"
    echo ""

    # Update signatures
    if command -v freshclam &>/dev/null; then
        echo "→ Updating signatures..."
        freshclam 2>/dev/null || true
    fi

    echo "→ Scanning paths..."
    echo ""

    # Log scan details
    {
        echo "ClamAV Malware Scan"
        echo "Date: $(date)"
        echo "Paths:"
        printf '%s\n' "${scan_paths[@]}"
        echo ""
        echo "Results:"
        echo "========================================"
    } > "$log_file"

    # Run scan
    clamscan --infected --recursive "${scan_paths[@]}" >> "$log_file" 2>&1 &
    local scan_pid=$!

    echo "Scan running in background (PID: $scan_pid)"
    echo "Log file: $log_file"
    echo ""
    echo "Monitor with: tail -f $log_file"

    # Store scan info in reference DB
    store_reference "malware_scan_clamav_latest" "$log_file"
}

# Maldet scanner
scan_maldet() {
    local scan_paths=("$@")

    if ! command -v maldet &>/dev/null; then
        echo -e "${RED}Maldet not installed${NC}"
        return 1
    fi

    # Create temp file with paths
    local path_file="/tmp/maldet_paths_$$.txt"
    printf '%s\n' "${scan_paths[@]}" > "$path_file"

    echo -e "${CYAN}Starting Maldet scan...${NC}"
    echo ""

    # Update signatures
    echo "→ Updating signatures..."
    maldet -u 2>/dev/null || true

    echo "→ Starting scan..."
    maldet -b -f "$path_file"

    rm -f "$path_file"

    echo ""
    echo "View results:"
    echo "  maldet -l"
}

# Main scan menu
show_scan_menu() {
    while true; do
        print_banner "Malware Scanner"

        echo "Available Scanners:"
        for scanner in "${available_scanners[@]}"; do
            echo "  • ${scanner^}"
        done
        echo ""

        echo "Scan Scope:"
        echo "  1. Scan entire server"
        echo "  2. Scan specific user"
        echo "  3. Scan specific domain"
        echo "  4. Scan custom path"
        echo ""
        echo "  5. View scan results"
        echo "  6. Scanner settings"
        echo ""
        echo "  0. Back to main menu"
        echo ""

        read -p "Select option: " choice

        case $choice in
            1) scan_entire_server ;;
            2) scan_user_account ;;
            3) scan_domain ;;
            4) scan_custom_path ;;
            5) view_scan_results ;;
            6) scanner_settings ;;
            0) return 0 ;;
            *) echo -e "${RED}Invalid option${NC}"; sleep 1 ;;
        esac
    done
}

# Scan entire server
scan_entire_server() {
    echo ""
    print_header "Full Server Scan"

    if ! detect_control_panel; then
        read -p "Press Enter to continue..."
        return 1
    fi

    echo "Control Panel: ${CONTROL_PANEL^}"
    echo "Docroots found: ${#sanitized_docroot[@]}"
    echo ""

    if ! check_memory; then
        read -p "Press Enter to continue..."
        return 1
    fi

    echo "Select scanner:"
    local i=1
    for scanner in "${available_scanners[@]}"; do
        echo "  $i. ${scanner^}"
        ((i++))
    done
    echo ""

    read -p "Scanner: " scanner_choice

    if [ "$scanner_choice" -lt 1 ] || [ "$scanner_choice" -gt ${#available_scanners[@]} ]; then
        echo -e "${RED}Invalid choice${NC}"
        read -p "Press Enter to continue..."
        return 1
    fi

    local selected_scanner="${available_scanners[$((scanner_choice-1))]}"

    case "$selected_scanner" in
        imunify) scan_imunify "${sanitized_docroot[@]}" ;;
        clamav) scan_clamav "${sanitized_docroot[@]}" ;;
        maldet) scan_maldet "${sanitized_docroot[@]}" ;;
    esac

    echo ""
    read -p "Press Enter to continue..."
}

# Scan user account
scan_user_account() {
    echo ""
    print_header "Scan User Account"

    if ! detect_control_panel; then
        read -p "Press Enter to continue..."
        return 1
    fi

    # Use user manager to select user
    select_user_interactive "Select user to scan"

    if [ -z "$SELECTED_USER" ]; then
        echo "No user selected"
        read -p "Press Enter to continue..."
        return 1
    fi

    echo ""
    echo "Getting docroots for: $SELECTED_USER"

    local user_paths=($(get_user_docroots "$SELECTED_USER"))

    if [ ${#user_paths[@]} -eq 0 ]; then
        echo -e "${RED}No docroots found for user${NC}"
        read -p "Press Enter to continue..."
        return 1
    fi

    echo "Paths to scan: ${#user_paths[@]}"
    printf '  %s\n' "${user_paths[@]}"
    echo ""

    echo "Select scanner:"
    local i=1
    for scanner in "${available_scanners[@]}"; do
        echo "  $i. ${scanner^}"
        ((i++))
    done
    echo ""

    read -p "Scanner: " scanner_choice

    if [ "$scanner_choice" -lt 1 ] || [ "$scanner_choice" -gt ${#available_scanners[@]} ]; then
        echo -e "${RED}Invalid choice${NC}"
        read -p "Press Enter to continue..."
        return 1
    fi

    local selected_scanner="${available_scanners[$((scanner_choice-1))]}"

    case "$selected_scanner" in
        imunify) scan_imunify "${user_paths[@]}" ;;
        clamav) scan_clamav "${user_paths[@]}" ;;
        maldet) scan_maldet "${user_paths[@]}" ;;
    esac

    echo ""
    read -p "Press Enter to continue..."
}

# Scan domain
scan_domain() {
    echo ""
    print_header "Scan Domain"

    if ! detect_control_panel; then
        read -p "Press Enter to continue..."
        return 1
    fi

    read -p "Enter domain name: " domain

    if [ -z "$domain" ]; then
        echo "No domain entered"
        read -p "Press Enter to continue..."
        return 1
    fi

    local domain_path=$(get_domain_docroot "$domain")

    if [ -z "$domain_path" ] || [ ! -d "$domain_path" ]; then
        echo -e "${RED}Domain not found or docroot doesn't exist${NC}"
        read -p "Press Enter to continue..."
        return 1
    fi

    echo "Docroot: $domain_path"
    echo ""

    echo "Select scanner:"
    local i=1
    for scanner in "${available_scanners[@]}"; do
        echo "  $i. ${scanner^}"
        ((i++))
    done
    echo ""

    read -p "Scanner: " scanner_choice

    if [ "$scanner_choice" -lt 1 ] || [ "$scanner_choice" -gt ${#available_scanners[@]} ]; then
        echo -e "${RED}Invalid choice${NC}"
        read -p "Press Enter to continue..."
        return 1
    fi

    local selected_scanner="${available_scanners[$((scanner_choice-1))]}"

    case "$selected_scanner" in
        imunify) scan_imunify "$domain_path" ;;
        clamav) scan_clamav "$domain_path" ;;
        maldet) scan_maldet "$domain_path" ;;
    esac

    echo ""
    read -p "Press Enter to continue..."
}

# Scan custom path
scan_custom_path() {
    echo ""
    print_header "Scan Custom Path"

    read -p "Enter path to scan: " custom_path

    if [ -z "$custom_path" ] || [ ! -d "$custom_path" ]; then
        echo -e "${RED}Path doesn't exist${NC}"
        read -p "Press Enter to continue..."
        return 1
    fi

    echo "Path: $custom_path"
    echo ""

    echo "Select scanner:"
    local i=1
    for scanner in "${available_scanners[@]}"; do
        echo "  $i. ${scanner^}"
        ((i++))
    done
    echo ""

    read -p "Scanner: " scanner_choice

    if [ "$scanner_choice" -lt 1 ] || [ "$scanner_choice" -gt ${#available_scanners[@]} ]; then
        echo -e "${RED}Invalid choice${NC}"
        read -p "Press Enter to continue..."
        return 1
    fi

    local selected_scanner="${available_scanners[$((scanner_choice-1))]}"

    case "$selected_scanner" in
        imunify) scan_imunify "$custom_path" ;;
        clamav) scan_clamav "$custom_path" ;;
        maldet) scan_maldet "$custom_path" ;;
    esac

    echo ""
    read -p "Press Enter to continue..."
}

# View scan results
view_scan_results() {
    echo ""
    print_header "Scan Results"

    echo "Select scanner to view results:"
    local i=1
    for scanner in "${available_scanners[@]}"; do
        echo "  $i. ${scanner^}"
        ((i++))
    done
    echo ""

    read -p "Scanner: " scanner_choice

    if [ "$scanner_choice" -lt 1 ] || [ "$scanner_choice" -gt ${#available_scanners[@]} ]; then
        echo -e "${RED}Invalid choice${NC}"
        read -p "Press Enter to continue..."
        return 1
    fi

    local selected_scanner="${available_scanners[$((scanner_choice-1))]}"

    echo ""

    case "$selected_scanner" in
        imunify)
            echo "Recent ImunifyAV scans:"
            imunify-antivirus malware on-demand list --since $(date --date="7 days ago" '+%s') 2>/dev/null || echo "No scans found"
            ;;
        clamav)
            echo "Recent ClamAV scans:"
            find "$SCRIPT_DIR/logs/malware-scans" -name "clamav_*.log" -mtime -7 2>/dev/null | sort -r | head -5 || echo "No scans found"
            ;;
        maldet)
            echo "Recent Maldet scans:"
            maldet -l 2>/dev/null || echo "No scans found"
            ;;
    esac

    echo ""
    read -p "Press Enter to continue..."
}

# Scanner settings
scanner_settings() {
    echo ""
    print_header "Scanner Settings"

    echo "Settings (placeholder for future enhancements):"
    echo "  • Auto-quarantine infected files"
    echo "  • Email notifications"
    echo "  • Scheduled scans"
    echo "  • Custom exclusions"
    echo ""
    echo "Coming soon..."

    read -p "Press Enter to continue..."
}

# Main execution
main() {
    if ! detect_scanners; then
        exit 1
    fi

    show_scan_menu
}

# Run if executed directly
if [ "${BASH_SOURCE[0]}" = "${0}" ]; then
    main "$@"
fi