#!/bin/bash
# load-geo-sets.sh - Populate nftables geolocation sets in batches
# to avoid "No buffer space available" netlink errors.
#
# Usage: load-geo-sets.sh [countrysets_dir]
#
# Run this AFTER loading the main firewall ruleset (firewall.nft),
# which creates the empty au-ipv4 and au-ipv6 sets.

COUNTRYSETS_DIR="${1:-/etc/nftables/firewall/geo-nft/countrysets}"
TABLE_FAMILY="inet"
TABLE_NAME="filter"
BATCH_SIZE=512

# Load elements from a geo-nft define file into an nftables set in batches.
# Args: $1=file  $2=set_name
load_set() {
        local file="$1"
        local set_name="$2"
        local batch=""
        local count=0
        local total=0

        if [ ! -f "$file" ]; then
                echo "Error: $file not found" >&2
                return 1
        fi

        # Extract element lines: skip comments, define line, and braces.
        # Element lines start with whitespace followed by an address.
        while IFS= read -r line; do
                # Strip leading/trailing whitespace and trailing comma
                element="${line#"${line%%[![:space:]]*}"}"
                element="${element%,}"
                element="${element%"${element##*[![:space:]]}"}"

                [ -z "$element" ] && continue

                if [ -z "$batch" ]; then
                        batch="$element"
                else
                        batch="$batch, $element"
                fi
                ((count++))
                ((total++))

                if [ "$count" -ge "$BATCH_SIZE" ]; then
                        echo "add element $TABLE_FAMILY $TABLE_NAME $set_name { $batch }" | nft -f - || {
                                echo "Error: failed adding batch to $set_name at element $total" >&2
                                return 1
                        }
                        batch=""
                        count=0
                fi
        done < <(sed -n '/^[[:space:]][[:space:]]*[0-9a-fA-F]/p' "$file")

        # Load any remaining elements
        if [ -n "$batch" ]; then
                echo "add element $TABLE_FAMILY $TABLE_NAME $set_name { $batch }" | nft -f - || {
                        echo "Error: failed adding final batch to $set_name at element $total" >&2
                        return 1
                }
        fi

        echo "Loaded $total elements into $set_name"
}

load_set "$COUNTRYSETS_DIR/AU.ipv4" "au-ipv4"
load_set "$COUNTRYSETS_DIR/AU.ipv6" "au-ipv6"