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


# Function to display the menu
show_menu() {
    echo "=============================="
    echo "       Iptables Manager       "
    echo "=============================="
    echo "1. Add IP address to a chain"
    echo "2. Remove IP address from a chain"
    echo "3. Check if IP address exists in any chain"
    echo "4. Exit"
    echo "=============================="
}

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

# Function to get available chains
get_chains() {
    echo "Available chains:"
    echo "1. INPUT"
    echo "2. OUTPUT"
    echo "3. FORWARD"
}

# Function to check if an IP address is in a specified chain
is_in_chain() {
    local chain=$1
    local ip_address=$2
    sudo iptables -C "$chain" -s "$ip_address" -j ACCEPT &> /dev/null || sudo iptables -C "$chain" -s "$ip_address" -j DROP &> /dev/null
}

# Function to add IP address to a specified chain
add_to_chain() {
    get_chains
    read -p "Select a chain (1-3): " chain_option
    case $chain_option in
        1) chain="INPUT" ;;
        2) chain="OUTPUT" ;;
        3) chain="FORWARD" ;;
        *) echo "Invalid chain selection."; return ;;
    esac

    read -p "Enter the IP address to add: " ip_address
    if is_in_chain "$chain" "$ip_address"; then
        echo "IP address $ip_address is already in the chain $chain."
    else
        read -p "Enter the target action (ACCEPT or DROP): " action
        if [[ "$action" == "ACCEPT" || "$action" == "DROP" ]]; then
            sudo iptables -A "$chain" -s "$ip_address" -j "$action"
            echo "IP address $ip_address added to chain $chain with action $action."
        else
            echo "Invalid action. Please use ACCEPT or DROP."
        fi
    fi
}

# Function to remove IP address from a specified chain
remove_from_chain() {
    get_chains
    read -p "Select a chain (1-3): " chain_option
    case $chain_option in
        1) chain="INPUT" ;;
        2) chain="OUTPUT" ;;
        3) chain="FORWARD" ;;
        *) echo "Invalid chain selection."; return ;;
    esac

    read -p "Enter the IP address to remove: " ip_address
    if is_in_chain "$chain" "$ip_address"; then
        read -p "Enter the target action (ACCEPT or DROP): " action
        if [[ "$action" == "ACCEPT" || "$action" == "DROP" ]]; then
            sudo iptables -D "$chain" -s "$ip_address" -j "$action"
            echo "IP address $ip_address removed from chain $chain with action $action."
        else
            echo "Invalid action. Please use ACCEPT or DROP."
        fi
    else
        echo "IP address $ip_address is not in the chain $chain."
    fi
}

# Function to check if an IP address exists in any chain
check_ip_address() {
    read -p "Enter the IP address to check: " ip_address
    for chain in INPUT OUTPUT FORWARD; do
        if is_in_chain "$chain" "$ip_address"; then
            echo "IP address $ip_address exists in chain $chain."
            return
        fi
    done
    echo "IP address $ip_address does not exist in any chain."
}

# Main script loop
while true; do
    show_menu
    read -p "Select an option [1-4]: " option

    # Check if input is empty
    if [[ -z "$option" ]]; then
        echo "No option selected. Please try again."
        continue
    fi

    case $option in
        1)
            add_to_chain
            ;;
        2)
            remove_from_chain
            ;;
        3)
            check_ip_address
            ;;
        4)
            echo "Exiting..."
            exit 0
            ;;
        *)
            echo "Invalid option. Please try again."
            ;;
    esac
done