#!/usr/bin/env python3
# Copyright (C) 2026 Paul Klumpp
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License.
"""Keep Proxmox cluster firewall IPSets in sync with DNS.

Hostnames are resolved repeatedly (A and AAAA) until no new IP
addresses show up, so that all IPs behind a round-robin or anycast
DNS record are collected. Afterwards the IPSet is synchronized: IPs
that are gone are deleted, new IPs are added. Only the difference is
written -- no flush-and-refill.

Usage:
    manage_ipset.py                          # sync all IPSets from the config
    manage_ipset.py -c /path/to.conf         # ... from a specific config
    manage_ipset.py <ipset> <host> [host...] # sync a single IPSet
    manage_ipset.py github github.com --round 30

Config file format (default: manage_ipset.conf next to this script):

    [defaults]
    round=50
    stale_after=10
    # dns_server=1.1.1.1        # optional, default: system resolver
    # timeout=15

    golang
      host=go.dev

    github
      host=github.com
      host=github.io round=10 stale_after=3

An IPSet may hold several hosts; their IPs are merged into one set.
Per-host options: round, stale_after, dns_server, timeout.

An IPSet is left untouched if any of its hosts fails to resolve, so a
single broken hostname cannot strip the IPs of its siblings.

Requires: python3 (stdlib only), pvesh (pve-manager), dig (dnsutils).
"""

import argparse
import ipaddress
import json
import shutil
import subprocess
import sys
from datetime import datetime
from pathlib import Path

DEFAULTS = {"round": 50, "stale_after": 10, "timeout": 15}


def check_dependencies():
    for cmd, pkg in (("pvesh", "pve-manager"), ("dig", "dnsutils")):
        if shutil.which(cmd) is None:
            print(f"ERROR: missing dependency: {cmd} ({pkg})")
            sys.exit(1)


def dig(*args, timeout=15):
    """Run dig, return stdout. Raises OSError if dig is missing."""
    result = subprocess.run(["dig", *args],
                            capture_output=True, text=True, check=False,
                            timeout=timeout)
    return result.stdout


def extract_ips(lines):
    """Return the valid IP addresses from dig +short output.

    dig follows CNAME chains and returns the final A/AAAA records, but
    it also echoes the intermediate CNAMEs into the output. Anything
    that is not a valid IP address (CNAMEs, TXT, noise) is dropped."""
    ips = set()
    for line in lines.splitlines():
        line = line.strip()
        if not line or line.startswith(";"):
            continue
        try:
            ips.add(str(ipaddress.ip_address(line)))
        except ValueError:
            continue
    return ips


def pvesh(*args):
    """Run pvesh, return stdout as text. None on error."""
    try:
        result = subprocess.run(["pvesh", *args],
                                capture_output=True, text=True)
    except OSError:
        return None
    if result.returncode != 0:
        return None
    return result.stdout


def pvesh_ok(*args):
    return pvesh(*args) is not None


def resolve_host(host, max_rounds=50, stale_after=10, timeout=15,
                 dns_server=None):
    """Resolve A and AAAA records of host, collecting new IPs until
    no new one appears for stale_after rounds (up to max_rounds).

    dns_server: optional address to pass to dig as @server."""
    collected = set()
    stale = 0
    server_args = [f"@{dns_server}"] if dns_server else []
    for _ in range(max_rounds):
        ips = set()
        for qtype in ("A", "AAAA"):
            try:
                out = dig("+short", *server_args, host, qtype,
                          timeout=timeout)
            except (OSError, subprocess.TimeoutExpired):
                out = ""
            ips |= extract_ips(out)
        new = ips - collected
        if new:
            collected |= new
            stale = 0
        else:
            stale += 1
            if stale >= stale_after:
                break
    return collected


def parse_cidrs(payload):
    """Extract plain IP addresses from a pvesh JSON payload.

    Handles both the nested data format ({"data": [...CIDR...]})
    and a bare list of CIDR strings."""
    if isinstance(payload, dict) and "data" in payload:
        entries = payload["data"]
    elif isinstance(payload, list):
        entries = payload
    else:
        entries = []
    ips = set()
    for entry in entries:
        cidr = entry if isinstance(entry, str) else entry.get("cidr", "")
        ips.add(cidr.split("/")[0])
    return ips


def ensure_ipset(name, comment=None):
    if comment is None:
        comment = f"auto-generated {datetime.now().astimezone().isoformat()}"
    if not pvesh_ok("get", f"/cluster/firewall/ipset/{name}"):
        print(f"creating IPSet {name}")
        pvesh("create", "/cluster/firewall/ipset", "--name", name,
               "--comment", comment)
        return
    # existing: refresh the comment only if it differs
    out = pvesh("get", f"/cluster/firewall/ipset/{name}",
                 "--output-format", "json")
    current_comment = ""
    if out:
        try:
            data = json.loads(out)
            if isinstance(data, dict) and "data" in data:
                current_comment = data["data"].get("comment", "")
            elif isinstance(data, dict):
                current_comment = data.get("comment", "")
        except (json.JSONDecodeError, AttributeError):
            pass
    if current_comment != comment:
        pvesh("set", f"/cluster/firewall/ipset/{name}", "--comment", comment)


def sync_ipset(name, ips):
    payload = pvesh("get", f"/cluster/firewall/ipset/{name}",
                    "--output-format", "json")
    current = set()
    if payload:
        try:
            current = parse_cidrs(json.loads(payload))
        except json.JSONDecodeError:
            current = set()
    existing, new = current - ips, ips - current
    for ip in sorted(new):
        print(f"ADD: {ip}")
        pvesh("create", f"/cluster/firewall/ipset/{name}", "--cidr", ip)
    for ip in sorted(existing):
        print(f"DEL: {ip}")
        pvesh("delete", f"/cluster/firewall/ipset/{name}/{ip}")


def sync_group(name, entries):
    """Sync one IPSet from its host entries.

    Returns True if the IPSet was synced, False if it was skipped
    because one or more hosts did not resolve."""
    hosts = ", ".join(e["host"] for e in entries)
    print(f"== {name}: {hosts}")
    all_ips = set()
    unresolved = []
    for entry in entries:
        found = resolve_host(entry["host"],
                             max_rounds=entry["round"],
                             stale_after=entry["stale_after"],
                             timeout=entry["timeout"],
                             dns_server=entry.get("dns_server") or None)
        if not found:
            unresolved.append(entry["host"])
        all_ips |= found
    if unresolved:
        # fail-safe: never let one broken hostname strip its siblings
        print(f"WARNING: no DNS results for {', '.join(unresolved)}, "
              f"skipping {name}")
        return False
    ensure_ipset(name)
    sync_ipset(name, all_ips)
    return True


def parse_config(text):
    """Parse the config.

    Returns (ipsets, defaults) where ipsets maps ipset name -> list of
    host entries (dicts with host, round, stale_after, dns_server,
    timeout) and defaults is the [defaults] section.

    The [defaults] section is terminated by the first line that is not
    a key=value pair (i.e. an IPSet name)."""
    defaults = dict(DEFAULTS)
    ipsets = {}
    current = None
    section = None
    for raw in text.splitlines():
        line = raw.split("#", 1)[0].strip()
        if not line:
            continue
        if line.startswith("[") and line.endswith("]"):
            section = line[1:-1]
            current = None
            continue
        if section == "defaults":
            if "=" not in line:
                # bare line ends the [defaults] section and starts an IPSet
                section = None
            else:
                key, _, value = line.partition("=")
                key = key.strip()
                value = value.strip()
                if value:
                    if key in ("round", "stale_after", "timeout"):
                        defaults[key] = int(value)
                    else:
                        defaults[key] = value
                continue
        if "=" in line:
            entry = _kv(line)
            if current is None:
                continue
            entry["round"] = int(entry.get("round", defaults["round"]))
            entry["stale_after"] = int(entry.get("stale_after",
                                               defaults["stale_after"]))
            entry["timeout"] = int(entry.get("timeout",
                                             defaults["timeout"]))
            entry.setdefault("dns_server", defaults.get("dns_server", ""))
            current.append(entry)
        else:
            current = ipsets.setdefault(line, [])
    return ipsets, defaults


def _kv(line):
    """'host=go.dev round=10' -> {host: 'go.dev', round: '10'}"""
    entry = {}
    for token in line.split():
        if "=" in token:
            key, value = token.split("=", 1)
            entry[key] = value
    return entry


def default_config_path():
    return Path(__file__).parent / "manage_ipset.conf"


def main(argv=None):
    parser = argparse.ArgumentParser(
        description="Keep Proxmox firewall IPSets in sync with DNS.")
    parser.add_argument("-c", "--config",
                        help=f"config file (default: {default_config_path()})")
    parser.add_argument("ipset", nargs="?", help="IPSet name (single mode)")
    parser.add_argument("hosts", nargs="*", help="hostnames (single mode)")
    parser.add_argument("--round", type=int, default=DEFAULTS["round"],
                        help="max. DNS query rounds per host")
    parser.add_argument("--stale-after", type=int,
                        default=DEFAULTS["stale_after"],
                        help="rounds without a new IP before stopping")
    parser.add_argument("--timeout", type=int, default=DEFAULTS["timeout"],
                        help="dig timeout in seconds")
    parser.add_argument("--dns-server",
                        help="resolver to query (default: system resolver)")
    args = parser.parse_args(argv)
    check_dependencies()

    if args.ipset:
        if not args.hosts:
            parser.error("hostname(s) required when an IPSet is given")
        entries = [{"host": host,
                    "round": args.round,
                    "stale_after": args.stale_after,
                    "timeout": args.timeout,
                    "dns_server": args.dns_server or ""}
                   for host in args.hosts]
        sys.exit(0 if sync_group(args.ipset, entries) else 1)

    path = Path(args.config) if args.config else default_config_path()
    try:
        text = path.read_text()
    except OSError as err:
        print(f"ERROR: cannot read {path}: {err}")
        sys.exit(1)
    ipsets, _defaults = parse_config(text)
    if not ipsets:
        print(f"ERROR: no IPSets found in {path}")
        sys.exit(1)
    failed = [name for name, entries in ipsets.items()
              if entries and not sync_group(name, entries)]
    if failed:
        print(f"FAILED: {', '.join(failed)}")
        sys.exit(1)


if __name__ == "__main__":
    main()