#!/bin/bash
#
# Script: check_ip_bans
#
# Usage: check_ip_bans
# - Script does not accept any input
#
# Description: A tool for reviewing active DROP/REJECT rules and comparing that to a list of IPs that should Never be banned.
#
# Author: Keith Muska
#
# Changelog:
# * 07.13.2026: Keith Muska
# - Initial write with assistance of Gemini 3.1 Pro
#

# List of Whitelisted IPs
IP_WHITELIST_URL="https://fs01.us-north-4.mojocloud.com/whitelisted_ips"

# List of Whitelisted ASNs
ASN_WHITELIST_URL="https://fs01.us-north-4.mojocloud.com/whitelisted_asn"

# Ensure we have root privileges to read firewall rules
if [ "$(id -u)" -ne 0 ]; then
    echo "UNKNOWN: This script must be run as root to access firewall rules."
    exit 3
fi

# Determine available Python executable (Prefers 3, falls back to 2/default)
if command -v python3 >/dev/null 2>&1; then
    PY_CMD="python3"
elif command -v python >/dev/null 2>&1; then
    PY_CMD="python"
elif command -v python2 >/dev/null 2>&1; then
    PY_CMD="python2"
else
    echo "UNKNOWN: Required command 'python' is not installed."
    exit 3
fi

# Check for curl
if ! command -v curl >/dev/null 2>&1; then
    echo "UNKNOWN: Required command 'curl' is not installed."
    exit 3
fi

# Check for DNS tools
if ! command -v dig >/dev/null 2>&1 && ! command -v host >/dev/null 2>&1; then
    echo "UNKNOWN: Either 'dig' or 'host' is required for ASN lookups."
    exit 3
fi

# Setup Temporary Files
BANNED_FILE=$(mktemp /tmp/banned_ips.XXXXXX)
WHITELIST_IP_FILE=$(mktemp /tmp/whitelist_ips.XXXXXX)
WHITELIST_ASN_FILE=$(mktemp /tmp/whitelist_asns.XXXXXX)

# Cleanup trap to ensure temporary files are deleted upon exit
trap 'rm -f "$BANNED_FILE" "$WHITELIST_IP_FILE" "$WHITELIST_ASN_FILE"' EXIT

# 1. Fetch whitelisted IPs
if ! curl -s "$IP_WHITELIST_URL" | grep -vE '^\s*#|^\s*$' > "$WHITELIST_IP_FILE"; then
    echo "UNKNOWN: Failed to download IP whitelist from $IP_WHITELIST_URL"
    exit 3
fi

# 2. Fetch whitelisted ASNs
if ! curl -s "$ASN_WHITELIST_URL" | grep -vE '^\s*#|^\s*$' > "$WHITELIST_ASN_FILE"; then
    echo "UNKNOWN: Failed to download ASN whitelist from $ASN_WHITELIST_URL"
    exit 3
fi

# 3. Gather banned IPs from all available firewall managers

# -- Firewalld --
if command -v firewall-cmd >/dev/null 2>&1 && firewall-cmd --state >/dev/null 2>&1; then
    firewall-cmd --zone=drop --list-sources 2>/dev/null | tr ' ' '\n' | grep -v '^ipset:' >> "$BANNED_FILE"
    
    DROP_IPSETS=$(firewall-cmd --zone=drop --list-sources 2>/dev/null | tr ' ' '\n' | grep '^ipset:' | cut -d':' -f2)
    for set in $DROP_IPSETS; do
        if command -v ipset >/dev/null 2>&1; then
            ipset list "$set" 2>/dev/null | grep -E '^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+(/[0-9]+)?$' >> "$BANNED_FILE"
        fi
    done
    
    firewall-cmd --list-all-zones 2>/dev/null | grep -i 'rule' | grep -iE 'drop|reject' | grep -oE 'address="[0-9\./]+"' | cut -d'"' -f2 >> "$BANNED_FILE"
fi

# -- UFW --
if command -v ufw >/dev/null 2>&1 && ufw status | grep -q "Status: active"; then
    ufw status 2>/dev/null | grep -iE 'deny|reject' | grep -oE '[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+(/[0-9]+)?' >> "$BANNED_FILE"
fi

# -- iptables --
if command -v iptables-save >/dev/null 2>&1; then
    iptables-save 2>/dev/null | grep -iE -- '-j (DROP|REJECT)' | grep -oE '(-s|-d|--source|--destination)[ =][0-9\./]+' | sed -E 's/.*[ =]//' >> "$BANNED_FILE"
    
    SETS=$(iptables-save 2>/dev/null | grep -iE -- '-j (DROP|REJECT)' | grep -oE '(--match-set|--set) [^ ]+' | awk '{print $2}')
    if [ -n "$SETS" ] && command -v ipset >/dev/null 2>&1; then
        for s in $SETS; do
            ipset list "$s" 2>/dev/null | grep -E '^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+(/[0-9]+)?$' >> "$BANNED_FILE"
        done
    fi
fi

# -- nftables --
if command -v nft >/dev/null 2>&1; then
    nft list ruleset 2>/dev/null | awk '/drop|reject/ {for(i=1;i<=NF;i++) if($i ~ /^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+(\/[0-9]+)?$/) print $i}' >> "$BANNED_FILE"
fi

# Clean up banned list
sort -u "$BANNED_FILE" | sed '/^$/d' > "${BANNED_FILE}.tmp" && mv "${BANNED_FILE}.tmp" "$BANNED_FILE"

# 4. Evaluate rules using Py2/Py3 compatible struct bitwise math
$PY_CMD - "$BANNED_FILE" "$WHITELIST_IP_FILE" "$WHITELIST_ASN_FILE" << 'EOF'
import sys
import os
import subprocess
import socket
import struct

banned_file = sys.argv[1]
whitelist_ip_file = sys.argv[2]
whitelist_asns_file = sys.argv[3]

def ip_to_int(ip_str):
    return struct.unpack("!I", socket.inet_aton(ip_str))[0]

def int_to_ip(ip_int):
    return socket.inet_ntoa(struct.pack("!I", ip_int))

def parse_cidr(cidr_str):
    if '/' in cidr_str:
        parts = cidr_str.split('/')
        ip = parts[0]
        prefix = int(parts[1])
    else:
        ip = cidr_str
        prefix = 32
    try:
        ip_int = ip_to_int(ip)
        mask = (0xFFFFFFFF << (32 - prefix)) & 0xFFFFFFFF
        network = ip_int & mask
        broadcast = network | (~mask & 0xFFFFFFFF)
        return (network, broadcast, prefix, cidr_str)
    except Exception:
        return None

# Parse whitelisted IPs/CIDRs
whitelisted_networks = []
try:
    with open(whitelist_ip_file, 'r') as f:
        for line in f:
            line = line.strip()
            if line:
                parsed = parse_cidr(line)
                if parsed:
                    whitelisted_networks.append(parsed)
except Exception:
    pass

# Parse whitelisted ASNs into a set for fast lookup
allowed_asns = set()
try:
    with open(whitelist_asns_file, 'r') as f:
        for line in f:
            line = line.strip()
            if line:
                # Handles multiple ASNs per line just in case they are space-separated
                for asn in line.split():
                    allowed_asns.add(asn)
except Exception:
    pass

# Parse banned IPs/CIDRs
banned_networks = []
try:
    with open(banned_file, 'r') as f:
        for line in f:
            line = line.strip()
            if line:
                parsed = parse_cidr(line)
                if parsed:
                    banned_networks.append(parsed)
except Exception:
    pass

def get_asn(ip_str):
    try:
        # Reverse IPv4 for cymru DNS query
        parts = ip_str.split('.')
        parts.reverse()
        rev = ".".join(parts) + ".origin.asn.cymru.com"
        
        with open(os.devnull, 'w') as devnull:
            try:
                # Use subprocess to query dig or host
                out = subprocess.check_output(["dig", "+short", rev, "TXT"], stderr=devnull).decode("utf-8").strip()
            except OSError:
                out = subprocess.check_output(["host", "-t", "TXT", rev], stderr=devnull).decode("utf-8").strip()
                if "descriptive text" in out:
                    out = out.split("descriptive text ")[1].strip('"')
            
        if out:
            parts = out.strip('"').split("|")
            if len(parts) > 0:
                return parts[0].strip()
    except Exception:
        pass
    return None

critical = False

for b_net in banned_networks:
    b_network, b_broadcast, b_prefix, b_str = b_net
    
    # Check 1: Does it match the File Whitelist (Overlap or Subset)
    is_whitelisted = False
    for w_net in whitelisted_networks:
        w_network, w_broadcast, w_prefix, w_str = w_net
        
        # Overlap Logic: Two CIDR ranges overlap if the highest network address is less than or equal to the lowest broadcast address
        if max(b_network, w_network) <= min(b_broadcast, w_broadcast):
            is_whitelisted = True
            break
            
    if is_whitelisted:
        print("CRITICAL: The following IP should not be blacklisted: " + b_str + " (Matched Whitelist URL)")
        critical = True
        continue
        
    # Check 2: Does it match the ASN Whitelist
    base_ip = int_to_ip(b_network)
    asn_str = get_asn(base_ip)
    
    if asn_str:
        asns = asn_str.split()
        for a in asns:
            if a in allowed_asns:
                print("CRITICAL: The following IP should not be blacklisted: " + b_str + " (Matched Allowed ASN: " + a + ")")
                critical = True
                break

if critical:
    sys.exit(2)
else:
    print("OK: No invalid IP blacklists found")
    sys.exit(0)
EOF

# Bubble up the exit code from the python evaluation block
exit $?