#!/bin/bash

# Function to disable ufw
disable_ufw() {
    echo "[INFO] Disabling ufw firewall..."
    sudo ufw disable
    sudo ufw status
    echo "[OK] ufw has been disabled."
}

# Function to disable firewalld
disable_firewalld() {
    echo "[INFO] Disabling firewalld..."
    sudo systemctl stop firewalld
    sudo systemctl disable firewalld
    sudo systemctl status firewalld
    echo "[OK] firewalld has been disabled."
}

# Function to disable iptables
disable_iptables() {
    echo "[INFO] Flushing and disabling iptables rules..."
    sudo iptables -F
    sudo iptables -X
    sudo iptables -t nat -F
    sudo iptables -t nat -X
    sudo iptables -t mangle -F
    sudo iptables -t mangle -X

    sudo iptables -P INPUT ACCEPT
    sudo iptables -P FORWARD ACCEPT
    sudo iptables -P OUTPUT ACCEPT

    echo "[INFO] iptables rules after flushing:"
    sudo iptables -L -n -v

    # Ensure /etc/iptables directory exists
    if [ ! -d /etc/iptables ]; then
        echo "[INFO] /etc/iptables directory not found. Creating it..."
        sudo mkdir -p /etc/iptables
    fi

    # Persist iptables rules on reboot
    if command -v iptables-save &> /dev/null; then
        echo "[INFO] Saving iptables rules for persistence..."
        if sudo iptables-save > /etc/iptables/rules.v4; then
            echo "[OK] iptables rules saved successfully."
        else
            echo "[ERROR] Failed to save iptables rules. You might need 'iptables-persistent'."
            echo "[INFO] Install it using: sudo apt install iptables-persistent"
        fi
    fi
    echo "[OK] iptables has been flushed and disabled."
}

# Detect the OS
detect_os() {
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        OS=$ID
    else
        echo "[ERROR] Unable to detect OS. /etc/os-release not found."
        exit 1
    fi
}

# Check if a service is running
is_running() {
    local service=$1
    if systemctl is-active --quiet "$service"; then
        echo "running"
    else
        echo "stopped"
    fi
}

# Prompt the user to disable a component
prompt_disable() {
    local component=$1
    read -rp "Do you want to stop and disable $component? (yes/no): " choice
    case "$choice" in
        yes|y|Y) 
            echo "[INFO] Proceeding to disable $component..."
            $2
            ;;
        no|n|N) 
            echo "[INFO] Skipping $component."
            ;;
        *)
            echo "[ERROR] Invalid input. Skipping $component."
            ;;
    esac
}

# Main execution based on OS
main() {
    detect_os
    echo "[INFO] Detected OS: $OS"

    echo "[INFO] Checking for active firewall and iptables services..."

    # Check and prompt for ufw
    if command -v ufw &> /dev/null; then
        if [[ $(sudo ufw status | grep -i "active") ]]; then
            echo "[OK] Detected: ufw firewall is active."
            prompt_disable "ufw firewall" disable_ufw
        else
            echo "[INFO] ufw firewall is not active."
        fi
    fi

    # Check and prompt for firewalld
    if systemctl list-units --type=service | grep -q 'firewalld'; then
        if [ "$(is_running firewalld)" == "running" ]; then
            echo "[OK] Detected: firewalld is active."
            prompt_disable "firewalld" disable_firewalld
        else
            echo "[INFO] firewalld is not running."
        fi
    fi

    # Check and prompt for iptables
    if command -v iptables &> /dev/null; then
        IPTABLES_RULES=$(sudo iptables -L -n -v | grep -v "0 0")
        if [ -n "$IPTABLES_RULES" ]; then
            echo "[OK] Detected: iptables rules are active."
            prompt_disable "iptables" disable_iptables
        else
            echo "[INFO] iptables has no active rules."
        fi
    fi

    echo "[OK] Process completed."
}

# Execute the main function
main