#!/bin/bash
# Execute the script from the URL and remove it after execution:
# wget -qO- https://git.technozone.com.au/vijay/Scripts/raw/branch/main/manage_iptables | bash && rm -f manage_iptables

# Function to display the menu
show_menu() {
    echo "=============================="
    echo "       Iptables Manager       "
    echo "=============================="
    echo "1. List current iptables rules"
    echo "2. Add IP address to whitelist (ACCEPT)"
    echo "3. Remove IP address from blacklist (DROP)"
    echo "4. Add IP address to blacklist (DROP)"
    echo "5. Remove IP address from whitelist (ACCEPT)"
    echo "6. Exit"
    echo "=============================="
}

# Function to list current iptables rules
list_rules() {
    echo "Current iptables rules:"
    sudo iptables -L -n -v
}

# Function to check if an IP address is in the whitelist
is_in_whitelist() {
    local ip_address=$1
    sudo iptables -C INPUT -s "$ip_address" -j ACCEPT &> /dev/null
}

# Function to check if an IP address is in the blacklist
is_in_blacklist() {
    local ip_address=$1
    sudo iptables -C INPUT -s "$ip_address" -j DROP &> /dev/null
}

# Function to add IP address to whitelist
add_to_whitelist() {
    read -p "Enter the IP address to whitelist: " ip_address
    if is_in_whitelist "$ip_address"; then
        echo "IP address $ip_address is already in the whitelist."
    else
        sudo iptables -A INPUT -s "$ip_address" -j ACCEPT
        echo "IP address $ip_address added to whitelist."
    fi
}

# Function to remove IP address from blacklist
remove_from_blacklist() {
    read -p "Enter the IP address to remove from blacklist: " ip_address
    if is_in_blacklist "$ip_address"; then
        sudo iptables -D INPUT -s "$ip_address" -j DROP
        echo "IP address $ip_address removed from blacklist."
    else
        echo "IP address $ip_address is not in the blacklist."
    fi
}

# Function to add IP address to blacklist
add_to_blacklist() {
    read -p "Enter the IP address to blacklist: " ip_address
    if is_in_blacklist "$ip_address"; then
        echo "IP address $ip_address is already in the blacklist."
    else
        sudo iptables -A INPUT -s "$ip_address" -j DROP
        echo "IP address $ip_address added to blacklist."
    fi
}

# Function to remove IP address from whitelist
remove_from_whitelist() {
    read -p "Enter the IP address to remove from whitelist: " ip_address
    if is_in_whitelist "$ip_address"; then
        sudo iptables -D INPUT -s "$ip_address" -j ACCEPT
        echo "IP address $ip_address removed from whitelist."
    else
        echo "IP address $ip_address is not in the whitelist."
    fi
}

# Main script loop
while true; do
    show_menu
    read -p "Select an option [1-6]: " option
    option=$(echo "$option" | tr -d '[:space:]')  # Trim whitespace

    case $option in
        1)
            list_rules
            ;;
        2)
            add_to_whitelist
            ;;
        3)
            remove_from_blacklist
            ;;
        4)
            add_to_blacklist
            ;;
        5)
            remove_from_whitelist
            ;;
        6)
            echo "Exiting..."
            exit 0
            ;;
        *)
            echo "Invalid option. Please try again."
            ;;
    esac
done