#!/bin/bash

#############################################################################
# Firewall Operations - Platform-specific IP blocking and management
# Provides variables and functions for adding/removing IPs across all firewalls
# Must be sourced AFTER lib/system-detect.sh has set SYS_* variables
#############################################################################

# Source guard
if [ -n "${_FIREWALL_OPERATIONS_LOADED:-}" ]; then
    return 0
fi
readonly _FIREWALL_OPERATIONS_LOADED=1

#############################################################################
# CSF FIREWALL OPERATIONS
#############################################################################

derive_csf_operations() {
    export SYS_CSF_ALLOW="/etc/csf/csf.allow"
    export SYS_CSF_DENY="/etc/csf/csf.deny"
    export SYS_CSF_WHITELIST="/etc/csf/csf.whitelist"
    export SYS_CSF_REGEX="/etc/csf/csf.regex"
    export SYS_CSF_IGNOREAUTO="/etc/csf/csf.ignoreauto"
    export SYS_CSF_IGNORE="/etc/csf/csf.ignore"
    export SYS_CSF_LOG="/var/log/lfd.log"
    export SYS_CSF_QUEUE="/var/spool/csf"

    # CSF command paths
    export SYS_CSF_BIN="/usr/local/csf/bin"
    export SYS_CSF_CMD="/usr/sbin/csf"
    export SYS_CSF_IP_CMD="/usr/local/csf/bin/csftest.pl"

    # CSF IP blocking command format
    export SYS_CSF_BAN_CMD="csf -d"      # csf -d IP
    export SYS_CSF_UNBAN_CMD="csf -ar"   # csf -ar IP
    export SYS_CSF_ALLOW_CMD="csf -a"    # csf -a IP
}

#############################################################################
# FIREWALLD OPERATIONS
#############################################################################

derive_firewalld_operations() {
    export SYS_FIREWALLD_CONFIG="/etc/firewalld"
    export SYS_FIREWALLD_ZONES="/etc/firewalld/zones"
    export SYS_FIREWALLD_IPSETS="/etc/firewalld/ipsets"
    export SYS_FIREWALLD_SERVICES="/etc/firewalld/services"
    export SYS_FIREWALLD_LOG="/var/log/firewalld"
    export SYS_FIREWALLD_DB="/var/lib/firewalld"

    # firewalld command format
    export SYS_FIREWALLD_BAN_CMD="firewall-cmd --permanent --add-rich-rule='rule family=\"ipv4\" source address=\"IP\" reject'"
    export SYS_FIREWALLD_UNBAN_CMD="firewall-cmd --permanent --remove-rich-rule='rule family=\"ipv4\" source address=\"IP\" reject'"
    export SYS_FIREWALLD_ALLOW_CMD="firewall-cmd --permanent --add-source=IP/32"
    export SYS_FIREWALLD_RELOAD="firewall-cmd --reload"

    # firewalld ipset for mass blocking
    export SYS_FIREWALLD_IPSET_NAME="blocked_ips"
    export SYS_FIREWALLD_IPSET_FILE="/etc/firewalld/ipsets/$SYS_FIREWALLD_IPSET_NAME.xml"
}

#############################################################################
# IPTABLES OPERATIONS
#############################################################################

derive_iptables_operations() {
    export SYS_IPTABLES_CONFIG="/etc/sysconfig/iptables"
    export SYS_IPTABLES_RULES_DIR="/etc/iptables"
    export SYS_IPTABLES_STATE_DIR="/proc/net"
    export SYS_IPTABLES_LOG="/var/log/messages"

    # iptables command format
    export SYS_IPTABLES_BAN_CMD="iptables -I INPUT -s IP -j DROP"
    export SYS_IPTABLES_UNBAN_CMD="iptables -D INPUT -s IP -j DROP"
    export SYS_IPTABLES_ALLOW_CMD="iptables -I INPUT -s IP -j ACCEPT"
    export SYS_IPTABLES_SAVE="iptables-save > /etc/iptables/rules.v4"

    # iptables ipset for mass blocking
    export SYS_IPTABLES_IPSET_NAME="blocked_ips"
    export SYS_IPTABLES_IPSET_LIST="ipset list $SYS_IPTABLES_IPSET_NAME"
    export SYS_IPTABLES_IPSET_CREATE="ipset create $SYS_IPTABLES_IPSET_NAME hash:ip"
    export SYS_IPTABLES_IPSET_ADD="ipset add $SYS_IPTABLES_IPSET_NAME IP"
    export SYS_IPTABLES_IPSET_DEL="ipset del $SYS_IPTABLES_IPSET_NAME IP"
    export SYS_IPTABLES_IPSET_FLUSH="ipset flush $SYS_IPTABLES_IPSET_NAME"
}

#############################################################################
# UFW (Ubuntu Firewall) OPERATIONS
#############################################################################

derive_ufw_operations() {
    export SYS_UFW_CONFIG="/etc/ufw"
    export SYS_UFW_BEFORE_RULES="/etc/ufw/before.rules"
    export SYS_UFW_AFTER_RULES="/etc/ufw/after.rules"
    export SYS_UFW_RULES_DIR="/etc/ufw/user.d"
    export SYS_UFW_LOG="/var/log/ufw.log"
    export SYS_UFW_DB="/etc/ufw/user_rules"

    # UFW command format
    export SYS_UFW_BAN_CMD="ufw deny from IP"
    export SYS_UFW_UNBAN_CMD="ufw delete deny from IP"
    export SYS_UFW_ALLOW_CMD="ufw allow from IP"
    export SYS_UFW_RELOAD="ufw reload"

    # UFW ipset for mass blocking (using before.rules)
    export SYS_UFW_IPSET_NAME="blocked_ips"
    export SYS_UFW_BEFORE_RULES_CUSTOM="/etc/ufw/before.rules.d/10-blocked-ips"
}

#############################################################################
# IMUNIFY FIREWALL OPERATIONS
#############################################################################

derive_imunify_operations() {
    export SYS_IMUNIFY_CONFIG="/etc/sysconfig/imunify360"
    export SYS_IMUNIFY_CLI="/usr/bin/imunify360-agent"
    export SYS_IMUNIFY_LOG="/var/log/imunify360"
    export SYS_IMUNIFY_LOG_MAIN="/var/log/imunify360/imunify360.log"
    export SYS_IMUNIFY_DB="/var/lib/imunify360"
    export SYS_IMUNIFY_BLOCKLIST="/var/lib/imunify360/blocklist"
    export SYS_IMUNIFY_WHITELIST="/var/lib/imunify360/whitelist"

    # Imunify command format (via CLI)
    export SYS_IMUNIFY_BAN_CMD="imunify360-agent blacklist add --ip IP"
    export SYS_IMUNIFY_UNBAN_CMD="imunify360-agent blacklist remove --ip IP"
    export SYS_IMUNIFY_ALLOW_CMD="imunify360-agent whitelist add --ip IP"
    export SYS_IMUNIFY_LIST_BLOCKED="imunify360-agent blacklist list"
    export SYS_IMUNIFY_LIST_ALLOWED="imunify360-agent whitelist list"
}

#############################################################################
# PLESK FIREWALL OPERATIONS
#############################################################################

derive_plesk_firewall_operations() {
    export SYS_PLESK_FW_CONFIG="/etc/sysconfig/plesk-firewall"
    export SYS_PLESK_FW_RULES="/etc/sysconfig/plesk-firewall.rules"
    export SYS_PLESK_FW_LOG="/var/log/plesk-firewall.log"
    export SYS_PLESK_FW_WHITELIST="/etc/sysconfig/plesk-firewall.whitelist"
    export SYS_PLESK_FW_BLACKLIST="/etc/sysconfig/plesk-firewall.blacklist"

    # Plesk firewall command (via plesk CLI)
    export SYS_PLESK_FW_CMD="/usr/local/psa/bin/firewall"
}

#############################################################################
# GENERIC FIREWALL IP BLOCKING FUNCTIONS
#############################################################################

# Block an IP across the detected firewall
firewall_block_ip() {
    local ip="$1"
    local reason="${2:-Security block}"

    if [ -z "$ip" ]; then
        echo "ERROR: IP address required" >&2
        return 1
    fi

    case "$SYS_FIREWALL" in
        csf)
            csf -d "$ip" 2>/dev/null || {
                echo "ERROR: Failed to block $ip in CSF" >&2
                return 1
            }
            ;;
        firewalld)
            firewall-cmd --permanent --add-rich-rule="rule family=\"ipv4\" source address=\"$ip\" reject" 2>/dev/null || {
                echo "ERROR: Failed to block $ip in firewalld" >&2
                return 1
            }
            firewall-cmd --reload 2>/dev/null
            ;;
        iptables)
            if command -v ipset &>/dev/null; then
                ipset add "$SYS_IPTABLES_IPSET_NAME" "$ip" 2>/dev/null || {
                    # Create set if it doesn't exist
                    ipset create "$SYS_IPTABLES_IPSET_NAME" hash:ip 2>/dev/null
                    ipset add "$SYS_IPTABLES_IPSET_NAME" "$ip" 2>/dev/null
                }
            else
                iptables -I INPUT -s "$ip" -j DROP 2>/dev/null || {
                    echo "ERROR: Failed to block $ip with iptables" >&2
                    return 1
                }
            fi
            ;;
        ufw)
            ufw deny from "$ip" 2>/dev/null || {
                echo "ERROR: Failed to block $ip in UFW" >&2
                return 1
            }
            ;;
        plesk)
            # Plesk firewall (when enabled)
            if [ -x "$SYS_PLESK_FW_CMD" ]; then
                "$SYS_PLESK_FW_CMD" -S add-rule -rule_name "Block_$ip" -rule_enable true \
                    -client_name all -remote_address "$ip" -action drop 2>/dev/null || {
                    echo "ERROR: Failed to block $ip in Plesk firewall" >&2
                    return 1
                }
            fi
            ;;
        *)
            echo "ERROR: No firewall configured for IP blocking" >&2
            return 1
            ;;
    esac

    return 0
}

# Unblock an IP across the detected firewall
firewall_unblock_ip() {
    local ip="$1"

    if [ -z "$ip" ]; then
        echo "ERROR: IP address required" >&2
        return 1
    fi

    case "$SYS_FIREWALL" in
        csf)
            csf -ar "$ip" 2>/dev/null || {
                echo "ERROR: Failed to unblock $ip in CSF" >&2
                return 1
            }
            ;;
        firewalld)
            firewall-cmd --permanent --remove-rich-rule="rule family=\"ipv4\" source address=\"$ip\" reject" 2>/dev/null
            firewall-cmd --reload 2>/dev/null
            ;;
        iptables)
            if command -v ipset &>/dev/null; then
                ipset del "$SYS_IPTABLES_IPSET_NAME" "$ip" 2>/dev/null || true
            else
                iptables -D INPUT -s "$ip" -j DROP 2>/dev/null || true
            fi
            ;;
        ufw)
            ufw delete deny from "$ip" 2>/dev/null || true
            ;;
        plesk)
            if [ -x "$SYS_PLESK_FW_CMD" ]; then
                "$SYS_PLESK_FW_CMD" -S remove-rule -rule_name "Block_$ip" 2>/dev/null || true
            fi
            ;;
    esac

    return 0
}

# Check if an IP is currently blocked
firewall_is_blocked() {
    local ip="$1"

    if [ -z "$ip" ]; then
        echo "ERROR: IP address required" >&2
        return 1
    fi

    case "$SYS_FIREWALL" in
        csf)
            grep -q "^$ip" "$SYS_CSF_DENY" 2>/dev/null && return 0 || return 1
            ;;
        firewalld)
            firewall-cmd --list-rich-rules 2>/dev/null | grep -q "source address=\"$ip\"" && return 0 || return 1
            ;;
        iptables)
            if command -v ipset &>/dev/null; then
                ipset test "$SYS_IPTABLES_IPSET_NAME" "$ip" 2>/dev/null && return 0 || return 1
            else
                iptables -C INPUT -s "$ip" -j DROP 2>/dev/null && return 0 || return 1
            fi
            ;;
        ufw)
            ufw status numbered 2>/dev/null | grep -q "Deny.*from $ip" && return 0 || return 1
            ;;
        *)
            return 1
            ;;
    esac
}

# Bulk block multiple IPs (format: one IP per line, or space-separated)
firewall_bulk_block_ips() {
    local ips="$1"
    local blocked_count=0
    local failed_count=0

    case "$SYS_FIREWALL" in
        csf)
            while IFS= read -r ip; do
                [ -z "$ip" ] && continue
                if firewall_block_ip "$ip"; then
                    ((blocked_count++))
                else
                    ((failed_count++))
                fi
            done <<< "$ips"
            ;;
        firewalld)
            # Use richd rules for bulk blocks
            while IFS= read -r ip; do
                [ -z "$ip" ] && continue
                if firewall_block_ip "$ip"; then
                    ((blocked_count++))
                else
                    ((failed_count++))
                fi
            done <<< "$ips"
            firewall-cmd --reload 2>/dev/null
            ;;
        iptables)
            # Use ipset for efficient bulk blocking
            if command -v ipset &>/dev/null; then
                ipset create "$SYS_IPTABLES_IPSET_NAME" hash:ip 2>/dev/null || true
                while IFS= read -r ip; do
                    [ -z "$ip" ] && continue
                    if ipset add "$SYS_IPTABLES_IPSET_NAME" "$ip" 2>/dev/null; then
                        ((blocked_count++))
                    else
                        ((failed_count++))
                    fi
                done <<< "$ips"
                # Add rule if not already present
                iptables -C INPUT -m set --match-set "$SYS_IPTABLES_IPSET_NAME" src -j DROP 2>/dev/null || \
                    iptables -I INPUT -m set --match-set "$SYS_IPTABLES_IPSET_NAME" src -j DROP 2>/dev/null
            else
                while IFS= read -r ip; do
                    [ -z "$ip" ] && continue
                    if firewall_block_ip "$ip"; then
                        ((blocked_count++))
                    else
                        ((failed_count++))
                    fi
                done <<< "$ips"
            fi
            ;;
        ufw)
            while IFS= read -r ip; do
                [ -z "$ip" ] && continue
                if firewall_block_ip "$ip"; then
                    ((blocked_count++))
                else
                    ((failed_count++))
                fi
            done <<< "$ips"
            ;;
    esac

    echo "Blocked: $blocked_count, Failed: $failed_count"
    return 0
}

#############################################################################
# MAIN DERIVATION FUNCTION
#############################################################################

derive_all_firewall_operations() {
    case "$SYS_FIREWALL" in
        csf)
            derive_csf_operations
            ;;
        firewalld)
            derive_firewalld_operations
            ;;
        iptables)
            derive_iptables_operations
            ;;
        ufw)
            derive_ufw_operations
            ;;
        *)
            # Check for Imunify even if other firewall is detected
            if command -v imunify360-agent &>/dev/null; then
                derive_imunify_operations
            fi
            # Check for Plesk firewall on Plesk systems
            if [ "$SYS_CONTROL_PANEL" = "plesk" ] && [ -x "$SYS_PLESK_FW_CMD" ] 2>/dev/null; then
                derive_plesk_firewall_operations
            fi
            ;;
    esac
}

# Export functions
export -f firewall_block_ip
export -f firewall_unblock_ip
export -f firewall_is_blocked
export -f firewall_bulk_block_ips

# Auto-run if sourced with detection complete
if [ -n "${SYS_DETECTION_COMPLETE:-}" ]; then
    derive_all_firewall_operations
fi