#!/bin/zsh

# Function to extract the list of jails
extract_jail_list() {
    # Run fail2ban-client status and extract the jail list
    JAIL_LIST=$(fail2ban-client status | awk -F'\t' '/Jail list/ {print $2}' | tr ',' '\n' | tr -d ' ')

    # Check if jail list is empty
    if [[ -z "$JAIL_LIST" ]]; then
        echo "No jails found."
        return 1
    fi

    echo "$JAIL_LIST"
    return 0
}

# Function to extract and print all banned IPs
extract_and_print_all_banned_ips() {
    JAIL_LIST=$(extract_jail_list)
    if [[ $? -ne 0 ]]; then
        return 1
    fi

    ALL_BANNED_IPS=""

    for JAIL in ${(f)JAIL_LIST}; do
        STATUS_OUTPUT=$(fail2ban-client status $JAIL)
        
        # Extract the banned IPs
        BANNED_IP_LIST=$(echo "$STATUS_OUTPUT" | grep -oP '(?<=Banned IP list:\t).*')
        if [[ -n "$BANNED_IP_LIST" ]]; then
            ALL_BANNED_IPS+="$BANNED_IP_LIST "
        fi
    done

    # Print all found IPs
    echo "$ALL_BANNED_IPS" | tr ' ' '\n'
}

# Function to filter the IPs based on the provided argument
filter_ips() {
    SEARCH_IP=$1
    extract_and_print_all_banned_ips | grep "$SEARCH_IP"
}

# Main execution
if [[ $# -gt 0 ]]; then
    filter_ips $1
else
    extract_and_print_all_banned_ips
fi