#!/bin/bash

# Requires: pvesh (Proxmox), dig (dnsutils)

check_deps() {
    missing=()
    command -v pvesh >/dev/null 2>&1 || missing+=("pvesh")
    command -v dig   >/dev/null 2>&1 || missing+=("dnsutils")
    if [ ${#missing[@]} -gt 0 ]; then
        echo "ERROR: missing dependencies: ${missing[*]}"
        exit 1
    fi
}

check_deps

create_ipset() {
    ipset=$1
    comment="auto-generated $(date -Iseconds)"

    if ! pvesh get "/cluster/firewall/ipset/${ipset}" &>/dev/null; then
        echo "creating $ipset"
        pvesh create /cluster/firewall/ipset --name "${ipset}" --comment "$comment"
    else
        pvesh create /cluster/firewall/ipset --name "${ipset}" --rename "${ipset}" --comment "$comment"
    fi
}

manage_ipset() {
    ipset=$1
    shift
    hostnames=("$@")

    declare -A dns_ips
    max_rounds=50

    for hostname in "${hostnames[@]}"; do
        stale=0
        for ((round=1; round<=max_rounds; round++)); do
            dns4=$(dig +short "$hostname" A 2>/dev/null)
            dns6=$(dig +short "$hostname" AAAA 2>/dev/null)
            added=0
            while read -r ip; do
                if [ -n "$ip" ] && [ -z "${dns_ips[$ip]}" ]; then
                    dns_ips[$ip]=1
                    added=1
                fi
            done < <(echo -e "${dns4}\n${dns6}" | grep -v '^;' | grep -v '^$')
            if [ "$added" -eq 1 ]; then
                stale=0
            else
                stale=$((stale + 1))
                [ "$stale" -ge 10 ] && break
            fi
        done
    done

    if [ ${#dns_ips[@]} -eq 0 ]; then
        echo "WARNING: no DNS results for ${hostnames[*]}, skipping"
        return 1
    fi

    declare -A ipset_ips
    if existing=$(pvesh get "/cluster/firewall/ipset/$ipset" --output-format json 2>/dev/null) && [ -n "$existing" ]; then
        ip_list=$(echo "$existing" | grep -o '"cidr"[[:space:]]*:[[:space:]]*"[^"]*"' | sed -E 's/.*"([^"]*)"/\1/' | sed 's|/[0-9]*$||')
        while read -r ip; do
            [ -n "$ip" ] && ipset_ips[$ip]=1
        done < <(echo "$ip_list" | sort -u)
    fi

    for ip in "${!dns_ips[@]}"; do
        if [ -z "${ipset_ips[$ip]}" ]; then
            echo "ADD: $ip"
            pvesh create "/cluster/firewall/ipset/$ipset" --cidr "$ip"
        fi
    done

    for ip in "${!ipset_ips[@]}"; do
        if [ -z "${dns_ips[$ip]}" ]; then
            echo "DEL: $ip"
            pvesh delete "/cluster/firewall/ipset/$ipset/$ip"
        fi
    done
}

if [ $# -lt 2 ]; then
    echo "Usage: $0 <ipset> <hostname> [hostname ...]"
    exit 1
fi

create_ipset "$1"
manage_ipset "$@"