#!/bin/bash

# Variables (configurable)
ALLOWED_COUNTRIES="GB,SK"  # Comma-separated country codes
ALLOWED_IPS="1.2.3.4,5.6.7.0/24"  # Comma-separated IPs or CIDRs
USE_MAXMIND=false  # Set to true to use MaxMind GeoIP database
MAXMIND_ACCOUNT_ID="46544156"
MAXMIND_LICENSE_KEY="API KEY"
INTERFACE=""  # Leave empty to auto-detect WAN interface
WAN_INTERFACE=""
APPLY_TO_ALL_PORTS=true  # Set to false to apply to specific ports only
TCP_PORTS=""  # Comma-separated TCP ports (e.g., "22,80,443")
UDP_PORTS=""  # Comma-separated UDP ports (e.g., "53,123")
EXEMPT_PORTS=""  # Comma-separated ports to exempt if APPLY_TO_ALL_PORTS=true
LOG_PREFIX="GEO_BLOCK: "  # Prefix for blocked traffic logs
USE_CLOUDFLARE=true  # Enable Cloudflare proxy IPs by default

# Function to detect WAN interface if not specified
detect_wan_interface() {
    WAN_INTERFACE=$(ip route | grep default | awk '{print $5}')
    if [ -z "$WAN_INTERFACE" ]; then
        echo "WAN interface could not be detected. Please specify manually."
        exit 1
    fi
    echo "Detected WAN interface: $WAN_INTERFACE"
}

# Function to validate port numbers
validate_ports() {
    local ports=$1
    for port in ${ports//,/ }; do
        if ! [[ "$port" =~ ^[0-9]+$ ]] || [ "$port" -lt 1 ] || [ "$port" -gt 65535 ]; then
            echo "Invalid port number: $port"
            exit 1
        fi
    done
}

# Function to fetch Cloudflare IPv4 ranges
fetch_cloudflare_ips() {
    if [ "$USE_CLOUDFLARE" = true ]; then
        echo "Fetching Cloudflare IPv4 ranges..."
        local cloudflare_url="https://www.cloudflare.com/ips-v4"
        wget -q -O cloudflare_ips.txt "$cloudflare_url" --timeout=10 --tries=2
        if [ $? -ne 0 ]; then
            echo "Failed to download Cloudflare IPv4 ranges."
            exit 1
        fi
        if [ ! -s cloudflare_ips.txt ]; then
            echo "Downloaded Cloudflare file is empty or invalid."
            exit 1
        fi
        # Append Cloudflare IPs to ALLOWED_IPS
        local cloudflare_ips=$(cat cloudflare_ips.txt | tr '\n' ',' | sed 's/,$//')
        if [ -n "$ALLOWED_IPS" ]; then
            ALLOWED_IPS="$ALLOWED_IPS,$cloudflare_ips"
        else
            ALLOWED_IPS="$cloudflare_ips"
        fi
        rm -f cloudflare_ips.txt
        echo "Cloudflare IPv4 ranges added to allowed IPs."
    fi
}

# Function to install the script to /usr/local/bin
install_script() {
    local script_name=$(basename "$0")
    local target_path="/usr/local/bin/$script_name"
    echo "Checking if script $script_name is already installed at $target_path..."

    # Check if the script is running from the target path and is executable
    if [ "$(realpath "$0")" = "$target_path" ] && [ -x "$target_path" ]; then
        echo "Script is already running from $target_path and executable. Skipping installation and service registration."
        return 1  # Indicate to skip service registration
    fi

    # Check if the source script exists
    if [ ! -f "$0" ]; then
        echo "Error: Cannot locate the script file ($0)."
        exit 1
    fi

    # Install the script
    echo "Installing script $script_name to $target_path..."
    cp "$0" "$target_path"
    if [ $? -ne 0 ]; then
        echo "Failed to copy script to $target_path."
        exit 1
    fi
    chmod +x "$target_path"
    if [ $? -ne 0 ]; then
        echo "Failed to set executable permissions on $target_path."
        exit 1
    fi
    echo "Script installed successfully."
    return 0  # Indicate to proceed with service registration
}

# Function to ensure xt_geoip module and dependencies are installed and loaded
ensure_geoip_module() {
    echo "Checking for xt_geoip module and dependencies..."
    # Check if xtables-addons-common is installed
    if ! dpkg -l | grep -q xtables-addons-common; then
        echo "Installing xtables-addons-common..."
        apt update -y
        apt install -y xtables-addons-common ipcalc conntrack
        if [ $? -ne 0 ]; then
            echo "Failed to install xtables-addons-common."
            exit 1
        fi
    fi

    # Check if xt_geoip kernel module is available
    if ! modprobe xt_geoip 2>/dev/null; then
        echo "xt_geoip module not found. Installing kernel headers and building module..."
        # Install kernel headers and dkms for xtables-addons
        local kernel_version=$(uname -r)
        apt install -y linux-headers-"$kernel_version" xtables-addons-dkms
        if [ $? -ne 0 ]; then
            echo "Failed to install kernel headers or xtables-addons-dkms."
            exit 1
        fi
        # Rebuild xtables-addons modules
        dkms autoinstall
        if ! modprobe xt_geoip 2>/dev/null; then
            echo "Failed to load xt_geoip module after installation."
            exit 1
        fi
    fi

    # Ensure GeoIP database directory exists
    mkdir -p /usr/share/xt_geoip
    if [ ! -d /usr/share/xt_geoip ]; then
        echo "Failed to create /usr/share/xt_geoip directory."
        exit 1
    fi

    echo "xt_geoip module and dependencies are ready."
}

# Function to download and parse RIPE IP ranges
download_ripe_ips() {
    echo "Downloading RIPE IP ranges..."
    local primary_url="https://ftp.ripe.net/pub/stats/ripencc/delegated-ripencc-latest"
    local fallback_url="https://ftp.ripe.net/ripe/stats/delegated-ripencc-latest"
    
    # Try primary URL
    wget -q -O ripe_ips.txt "$primary_url" --timeout=10 --tries=2
    if [ $? -ne 0 ]; then
        echo "Primary RIPE URL failed, trying fallback URL..."
        wget -q -O ripe_ips.txt "$fallback_url" --timeout=10 --tries=2
        if [ $? -ne 0 ]; then
            echo "Failed to download RIPE IP ranges from both primary and fallback URLs."
            exit 1
        fi
    fi
    # Parse the downloaded file
    if [ ! -s ripe_ips.txt ]; then
        echo "Downloaded RIPE file is empty or invalid."
        exit 1
    fi
    grep -E "ripencc|ipv4" ripe_ips.txt | grep -E $(echo $ALLOWED_COUNTRIES | sed 's/,/|/g') | awk -F'|' '{if ($5 <= 4294967296) print $4"/"(32 - log($5)/log(2))}' > allowed_ips.txt
    if [ ! -s allowed_ips.txt ]; then
        echo "No IP ranges found for specified countries: $ALLOWED_COUNTRIES"
        exit 1
    fi
}

# Function to download and parse MaxMind IP ranges
download_maxmind_ips() {
    if [ "$USE_MAXMIND" = true ]; then
        if [ -z "$MAXMIND_LICENSE_KEY" ] || [ -z "$MAXMIND_ACCOUNT_ID" ]; then
            echo "MaxMind account ID and license key must be set."
            exit 1
        fi
        echo "Downloading MaxMind IP ranges..."
        wget -q -O GeoLite2-Country-CSV.zip "https://download.maxmind.com/app/geoip_download?edition_id=GeoLite2-Country-CSV&license_key=$MAXMIND_LICENSE_KEY&suffix=zip"
        if [ $? -ne 0 ]; then
            echo "Failed to download MaxMind GeoIP database."
            exit 1
        fi
        unzip -q GeoLite2-Country-CSV.zip
        mkdir -p /usr/share/xt_geoip
        /usr/libexec/xt_geoip/xt_geoip_build -D /usr/share/xt_geoip GeoLite2-Country-CSV_*/GeoLite2-Country-Blocks-IPv4.csv
    fi
}

# Function to apply iptables rules
apply_iptables_rules() {
    echo "Applying iptables rules..."

    # Flush existing rules
    iptables -F INPUT
    iptables -X GEOIP_CHAIN 2>/dev/null

    # Create new chain for GeoIP filtering
    iptables -N GEOIP_CHAIN

    # Allow established and related connections (preserves existing sessions)
    iptables -A INPUT -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT

    # Allow all traffic on non-WAN interfaces (preserves internal traffic)
    for iface in $(ip link | awk -F: '$0 !~ "lo|${WAN_INTERFACE}" {print $2}' | sed 's/^[ \t]*//'); do
        iptables -A INPUT -i "$iface" -j ACCEPT
    done

    # Allow loopback traffic
    iptables -A INPUT -i lo -j ACCEPT

    # Allow specified IPs
    if [ -n "$ALLOWED_IPS" ]; then
        for ip in ${ALLOWED_IPS//,/ }; do
            iptables -A INPUT -s "$ip" -j ACCEPT
        done
    fi

    # Handle port-specific or all-port rules
    if [ "$APPLY_TO_ALL_PORTS" = true ]; then
        # Apply GeoIP rules to all ports except exempted ones
        if [ -n "$EXEMPT_PORTS" ]; then
            for port in ${EXEMPT_PORTS//,/ }; do
                iptables -A INPUT -p tcp --dport "$port" -j ACCEPT
                iptables -A INPUT -p udp --dport "$port" -j ACCEPT
            done
        fi
        for country in ${ALLOWED_COUNTRIES//,/ }; do
            iptables -A INPUT -m geoip --src-cc "$country" -j ACCEPT
        done
    else
        # Apply GeoIP rules to specific TCP ports
        if [ -n "$TCP_PORTS" ]; then
            validate_ports "$TCP_PORTS"
            for port in ${TCP_PORTS//,/ }; do
                for country in ${ALLOWED_COUNTRIES//,/ }; do
                    iptables -A INPUT -p tcp --dport "$port" -m geoip --src-cc "$country" -j ACCEPT
                done
            done
        fi
        # Apply GeoIP rules to specific UDP ports
        if [ -n "$UDP_PORTS" ]; then
            validate_ports "$UDP_PORTS"
            for port in ${UDP_PORTS//,/ }; do
                for country in ${ALLOWED_COUNTRIES//,/ }; do
                    iptables -A INPUT -p udp --dport "$port" -m geoip --src-cc "$country" -j ACCEPT
                done
            done
        fi
    fi

    # Log and drop remaining inbound traffic on WAN interface
    iptables -A INPUT -i "$WAN_INTERFACE" -j LOG --log-prefix "$LOG_PREFIX"
    iptables -A INPUT -i "$WAN_INTERFACE" -j DROP
}

# Function to check if an IP would be blocked
check_ip() {
    local ip=$1
    if [ -z "$ip" ]; then
        echo "Error: No IP provided for checking."
        exit 1
    fi
    echo "Checking if IP $ip would be blocked..."

    # Check if IP is internal (all interfaces except WAN, including Docker/VPN)
    local all_interfaces=$(ip link | awk -F: '{print $2}' | sed 's/^[ \t]*//')
    for iface in $all_interfaces; do
        if [ "$iface" != "$WAN_INTERFACE" ]; then
            # Check exact IP match on interface
            local iface_ips=$(ip addr show "$iface" 2>/dev/null | grep -w inet | awk '{print $2}' | cut -d'/' -f1)
            for iface_ip in $iface_ips; do
                if [ "$ip" = "$iface_ip" ]; then
                    echo "IP $ip is allowed (exact match on interface: $iface)"
                    return 0
                fi
            done
            # Check if IP is in the interface's subnet
            local subnets=$(ip addr show "$iface" 2>/dev/null | grep -w inet | awk '{print $2}')
            for subnet in $subnets; do
                if [ -n "$subnet" ]; then
                    local network=$(echo "$subnet" | cut -d'/' -f1)
                    local mask=$(echo "$subnet" | cut -d'/' -f2)
                    if [ -n "$network" ] && [ -n "$mask" ] && [ "$mask" -ge 0 ] && [ "$mask" -le 32 ]; then
                        local ip_int=$(echo "$ip" | awk -F. '{printf "%d", ($1*16777216)+($2*65536)+($3*256)+$4}')
                        local net_int=$(echo "$network" | awk -F. '{printf "%d", ($1*16777216)+($2*65536)+($3*256)+$4}')
                        local mask_int=$(( (0xffffffff << (32 - $mask)) & 0xffffffff ))
                        if [ $((ip_int & mask_int)) -eq $((net_int & mask_int)) ]; then
                            echo "IP $ip is allowed (in subnet $subnet on interface: $iface)"
                            return 0
                        fi
                    fi
                fi
            done
        fi
    done

    # Check if IP is part of an established connection
    if command -v conntrack >/dev/null && conntrack -L 2>/dev/null | grep -q "src=$ip.*ESTABLISHED"; then
        echo "IP $ip is allowed (part of an established connection)"
        return 0
    fi

    # Check if IP is in ALLOWED_IPS
    if [ -n "$ALLOWED_IPS" ]; then
        for allowed_ip in ${ALLOWED_IPS//,/ }; do
            if [ "$ip" = "$allowed_ip" ]; then
                echo "IP $ip is allowed (matches allowed IP: $allowed_ip)"
                return 0
            fi
            # Check CIDR ranges in ALLOWED_IPS
            if [[ "$allowed_ip" =~ / ]]; then
                local network=$(echo "$allowed_ip" | cut -d'/' -f1)
                local mask=$(echo "$allowed_ip" | cut -d'/' -f2)
                if [ -n "$network" ] && [ -n "$mask" ] && [ "$mask" -ge 0 ] && [ "$mask" -le 32 ]; then
                    local ip_int=$(echo "$ip" | awk -F. '{printf "%d", ($1*16777216)+($2*65536)+($3*256)+$4}')
                    local net_int=$(echo "$network" | awk -F. '{printf "%d", ($1*16777216)+($2*65536)+($3*256)+$4}')
                    local mask_int=$(( (0xffffffff << (32 - $mask)) & 0xffffffff ))
                    if [ $((ip_int & mask_int)) -eq $((net_int & mask_int)) ]; then
                        echo "IP $ip is allowed (matches allowed IP range: $allowed_ip)"
                        return 0
                    fi
                fi
            fi
        done
    fi

    # Check if IP belongs to allowed countries
    if [ "$USE_MAXMIND" = true ] && [ -d "/usr/share/xt_geoip" ]; then
        if command -v geoip-country >/dev/null; then
            for country in ${ALLOWED_COUNTRIES//,/ }; do
                if geoip-country -d /usr/share/xt_geoip "$ip" 2>/dev/null | grep -q "^$country$"; then
                    echo "IP $ip is allowed (matches country: $country)"
                    return 0
                fi
            done
        else
            echo "Warning: geoip-country not found, falling back to iptables check"
            for country in ${ALLOWED_COUNTRIES//,/ }; do
                if iptables -t mangle -A PREROUTING -s "$ip" -m geoip --src-cc "$country" -j ACCEPT 2>/dev/null; then
                    iptables -t mangle -D PREROUTING -s "$ip" -m geoip --src-cc "$country" -j ACCEPT 2>/dev/null
                    echo "IP $ip is allowed (matches country: $country)"
                    return 0
                fi
            done
        fi
    elif [ -f "allowed_ips.txt" ]; then
        # Fallback to allowed_ips.txt for RIPE data, validate CIDR masks
        head -n 1000 allowed_ips.txt | while IFS= read -r range; do
            if [ -n "$range" ]; then
                local network=$(echo "$range" | cut -d'/' -f1)
                local mask=$(echo "$range" | cut -d'/' -f2)
                if [ -n "$network" ] && [ -n "$mask" ] && [ "$mask" -ge 0 ] && [ "$mask" -le 32 ]; then
                    local ip_int=$(echo "$ip" | awk -F. '{printf "%d", ($1*16777216)+($2*65536)+($3*256)+$4}')
                    local net_int=$(echo "$network" | awk -F. '{printf "%d", ($1*16777216)+($2*65536)+($3*256)+$4}')
                    local mask_int=$(( (0xffffffff << (32 - $mask)) & 0xffffffff ))
                    if [ $((ip_int & mask_int)) -eq $((net_int & mask_int)) ]; then
                        echo "IP $ip is allowed (matches country IP range: $range)"
                        return 0
                    fi
                fi
            fi
        done
    fi

    # If none of the above, IP would be blocked for new inbound connections on WAN
    echo "IP $ip is blocked (new inbound connections on WAN interface)"
    return 1
}

# Function to remove iptables rules
remove_iptables_rules() {
    echo "Removing iptables rules..."
    iptables -F INPUT
    iptables -X GEOIP_CHAIN 2>/dev/null
}

# Function to create systemd service
create_systemd_service() {
    local script_name=$(basename "$0")
    local target_path="/usr/local/bin/$script_name"
    echo "Creating systemd service for $script_name..."
    cat <<EOF > /etc/systemd/system/geoip-firewall.service
[Unit]
Description=GeoIP Firewall
After=network.target

[Service]
ExecStart=$target_path
ExecStop=$target_path --remove
RemainAfterExit=true

[Install]
WantedBy=multi-user.target
EOF
    systemctl daemon-reload
    if [ $? -ne 0 ]; then
        echo "Failed to reload systemd daemon."
        exit 1
    fi
    systemctl enable geoip-firewall.service
    if [ $? -ne 0 ]; then
        echo "Failed to enable geoip-firewall.service."
        exit 1
    fi
    systemctl start geoip-firewall.service
    if [ $? -ne 0 ]; then
        echo "Failed to start geoip-firewall.service."
        exit 1
    fi
    echo "Systemd service created and started successfully."
}

# Function to display usage
usage() {
    echo "Usage: $0 [options]"
    echo "Options:"
    echo "  --check <IP>            Check if an IP is blocked"
    echo "  --remove                Remove iptables rules"
    echo "  --interface <iface>     Specify WAN interface"
    echo "  --countries <codes>     Comma-separated country codes (e.g., GB,DE,FR)"
    echo "  --ips <IPs>             Comma-separated IPs or CIDRs"
    echo "  --tcp-ports <ports>     Comma-separated TCP ports (e.g., 22,80,443)"
    echo "  --udp-ports <ports>     Comma-separated UDP ports (e.g., 53,123)"
    echo "  --exempt-ports <ports>  Comma-separated ports to exempt"
    echo "  --all-ports             Apply filter to all ports (default)"
    echo "  --no-all-ports          Apply filter to specific ports only"
    echo "  --use-maxmind           Enable MaxMind GeoIP database"
    echo "  --no-cloudflare         Disable Cloudflare proxy IPs"
    exit 1
}

# Parse command-line options
while [[ $# -gt 0 ]]; do
    case $1 in
        --check)
            check_ip "$2"
            exit 0
            ;;
        --remove)
            remove_iptables_rules
            exit 0
            ;;
        --interface)
            INTERFACE="$2"
            shift 2
            ;;
        --countries)
            ALLOWED_COUNTRIES="$2"
            shift 2
            ;;
        --ips)
            ALLOWED_IPS="$2"
            shift 2
            ;;
        --tcp-ports)
            TCP_PORTS="$2"
            APPLY_TO_ALL_PORTS=false
            shift 2
            ;;
        --udp-ports)
            UDP_PORTS="$2"
            APPLY_TO_ALL_PORTS=false
            shift 2
            ;;
        --exempt-ports)
            EXEMPT_PORTS="$2"
            shift 2
            ;;
        --all-ports)
            APPLY_TO_ALL_PORTS=true
            shift
            ;;
        --no-all-ports)
            APPLY_TO_ALL_PORTS=false
            shift
            ;;
        --use-maxmind)
            USE_MAXMIND=true
            shift
            ;;
        --no-cloudflare)
            USE_CLOUDFLARE=false
            shift
            ;;
        *)
            usage
            ;;
    esac
done

# Main script logic
if [ "$1" = "--check" ]; then
    check_ip "$2"
elif [ "$1" = "--remove" ]; then
    remove_iptables_rules
else
    if [ -z "$INTERFACE" ]; then
        detect_wan_interface
    else
        WAN_INTERFACE="$INTERFACE"
    fi
    install_script
    if [ $? -eq 0 ]; then
        ensure_geoip_module
        fetch_cloudflare_ips
        download_ripe_ips
        download_maxmind_ips
        apply_iptables_rules
        create_systemd_service
    else
        ensure_geoip_module
        fetch_cloudflare_ips
        download_ripe_ips
        download_maxmind_ips
        apply_iptables_rules
        echo "Skipping systemd service creation as script is running from /usr/local/bin."
    fi
fi