#!/usr/bin/python3
# Managed by Ansible (role host_firewall), do not edit.
"""Rebuilds the FireHOL level1 set firehol_v4 of table inet host.

A failed download or sanity check exits non-zero before anything changes, so
the previous set stays in place.
"""
import ipaddress
import json
import os
import shutil
import subprocess
import sys
import tempfile
import time
import urllib.request

PERSISTED = "/etc/nftables.d/host-firehol.nft"
URL = "https://raw.githubusercontent.com/firehol/blocklist-ipsets/master/firehol_level1.netset"
USER_AGENT = "ruekov-host-firewall (nft-host-firehol)"

# Lower bound that catches an empty or truncated answer.
MIN_ENTRIES = 1000
# A prefix this wide is a parse error, not a bad network.
MIN_PREFIXLEN = 8
# Private, CGNAT (NetBird), loopback, link-local and multicast/reserved space:
# an entry overlapping one of these never lands in the set, so the own
# infrastructure cannot block itself.
EXCLUDED = [ipaddress.ip_network(n) for n in (
    "0.0.0.0/8", "10.0.0.0/8", "100.64.0.0/10", "127.0.0.0/8",
    "169.254.0.0/16", "172.16.0.0/12", "192.168.0.0/16", "224.0.0.0/3")]


def fail(message):
    print(f"nft-host-firehol: {message}", file=sys.stderr)
    sys.exit(1)


def fetch():
    error = None
    for attempt in range(3):
        try:
            request = urllib.request.Request(URL, headers={"User-Agent": USER_AGENT})
            with urllib.request.urlopen(request, timeout=120) as response:
                return response.read().decode("utf-8")
        except OSError as e:
            error = e
            time.sleep(5 * (attempt + 1))
    fail(f"download of {URL} failed: {error}")


def own_addresses():
    result = subprocess.run(["ip", "-j", "-4", "address", "show"], capture_output=True, text=True)
    if result.returncode != 0:
        fail(f"ip address show failed: {result.stderr.strip()}")
    return [ipaddress.ip_address(a["local"]) for link in json.loads(result.stdout)
            for a in link.get("addr_info", []) if a.get("family") == "inet"]


def main():
    networks = []
    for line in fetch().splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        try:
            networks.append(ipaddress.IPv4Network(line))
        except ValueError:
            fail(f"unparsable line: {line!r}")
    if len(networks) < MIN_ENTRIES:
        fail(f"only {len(networks)} entries")

    own = own_addresses()
    kept, excluded, own_skipped = [], 0, 0
    for network in networks:
        if any(network.overlaps(e) for e in EXCLUDED):
            excluded += 1
        elif any(address in network for address in own):
            own_skipped += 1
        else:
            if network.prefixlen < MIN_PREFIXLEN:
                fail(f"suspiciously wide prefix {network}")
            kept.append(network)
    merged = list(ipaddress.collapse_addresses(kept))
    print(f"firehol_v4: {len(networks)} entries, {excluded} in private or reserved space and "
          f"{own_skipped} holding an own address skipped -> {len(merged)} merged")

    elements = ",\n\t\t".join(str(n) for n in merged) + ","
    persisted = ("# Generated by nft-host-firehol, do not edit.\n"
                 "table inet host {\n"
                 "\tset firehol_v4 { type ipv4_addr; flags interval; auto-merge; elements = {\n"
                 f"\t\t{elements}\n\t}} }}\n}}\n")

    # Checked under another table name: a dry run against the live table
    # name changes the live sets for its duration and unlisted addresses match.
    check_rules = persisted.replace("table inet host {", "table inet host_firehol_check {")
    checked = subprocess.run(["nft", "-c", "-f", "-"], input=check_rules, capture_output=True, text=True)
    if checked.returncode != 0:
        fail(f"nft rejects the set: {checked.stderr.strip()}")

    with tempfile.TemporaryDirectory() as tmp:
        previous = None
        if os.path.exists(PERSISTED):
            previous = os.path.join(tmp, "previous.nft")
            shutil.copyfile(PERSISTED, previous)
        with open(PERSISTED + ".new", "w", encoding="utf-8") as f:
            f.write(persisted)
        os.chmod(PERSISTED + ".new", 0o644)
        os.replace(PERSISTED + ".new", PERSISTED)

        # A reload replaces the whole table, sets included, in one transaction.
        active = subprocess.run(["systemctl", "-q", "is-active", "nft-host.service"]).returncode == 0
        if active:
            reloaded = subprocess.run(["systemctl", "reload", "nft-host.service"], capture_output=True, text=True)
            if reloaded.returncode != 0:
                if previous:
                    shutil.copyfile(previous, PERSISTED)
                fail(f"reloading the rule set failed, kept the previous set: {reloaded.stderr.strip()}")
        print("reloaded the rule set" if active else "persisted only")


try:
    main()
except SystemExit:
    raise
except Exception as e:  # keep the previous set on any surprise
    fail(f"{type(e).__name__}: {e}")