#!/usr/bin/python3
# Managed by Ansible (role host_firewall), do not edit.
"""Rebuilds the Spamhaus sets spamhaus_v4 and spamhaus_v6 of table inet host.

Sources are the Spamhaus DROP and DROPv6 prefix lists and the ASN-DROP list,
whose ASNs are expanded to prefixes with the iptoasn bulk files (one download
per address family). Any failed download or sanity check exits non-zero before
anything changes, so the previous sets stay in place.
"""
import gzip
import http.client
import ipaddress
import json
import os
import shutil
import subprocess
import sys
import tempfile
import time
import urllib.request

PERSISTED = "/etc/nftables.d/host-spamhaus.nft"
SPAMHAUS = "https://www.spamhaus.org/drop/"
IPTOASN = "https://iptoasn.com/data/"
USER_AGENT = "ruekov-host-firewall (nft-host-spamhaus)"

# Lower bounds that catch an empty or truncated answer.
MIN_RECORDS = {"drop_v4": 800, "drop_v6": 30, "asndrop": 200}
MIN_TSV_LINES = {4: 200000, 6: 20000}
MIN_ASN_PREFIXES_V4 = 100
# A prefix this wide is a parse error, not a bad network.
MIN_PREFIXLEN = {4: 8, 6: 16}


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


def fetch(url, dest):
    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, open(dest, "wb") as out:
                shutil.copyfileobj(response, out)
            return
        except (OSError, http.client.HTTPException) as e:
            error = e
            time.sleep(5 * (attempt + 1))
    fail(f"download of {url} failed: {error}")


def spamhaus_records(name, tmp):
    path = os.path.join(tmp, name + ".json")
    fetch(f"{SPAMHAUS}{name}.json", path)
    records, meta = [], None
    with open(path, encoding="utf-8") as f:
        for line in f:
            if not line.strip():
                continue
            record = json.loads(line)
            if record.get("type") == "metadata":
                meta = record
            else:
                records.append(record)
    if meta is None or meta.get("records") != len(records):
        fail(f"{name}: record count does not match the metadata record")
    if len(records) < MIN_RECORDS[name]:
        fail(f"{name}: only {len(records)} records")
    return records


def asn_prefixes(asns, family, tmp):
    path = os.path.join(tmp, f"ip2asn-v{family}.tsv.gz")
    fetch(f"{IPTOASN}ip2asn-v{family}.tsv.gz", path)
    networks, lines = [], 0
    with gzip.open(path, "rt", encoding="utf-8") as f:
        for line in f:
            lines += 1
            fields = line.split("\t", 3)
            if len(fields) < 3 or fields[2] not in asns:
                continue
            networks.extend(ipaddress.summarize_address_range(
                ipaddress.ip_address(fields[0]), ipaddress.ip_address(fields[1])))
    if lines < MIN_TSV_LINES[family]:
        fail(f"ip2asn-v{family}: only {lines} lines")
    return networks


def merge(family, *groups):
    networks, skipped = [], 0
    for group in groups:
        for network in group:
            network = ipaddress.ip_network(network)
            if network.prefixlen < MIN_PREFIXLEN[family]:
                fail(f"suspiciously wide prefix {network}")
            if not network.is_global:
                skipped += 1
                continue
            networks.append(network)
    return list(ipaddress.collapse_addresses(networks)), skipped


def main():
    with tempfile.TemporaryDirectory() as tmp:
        drop4 = [r["cidr"] for r in spamhaus_records("drop_v4", tmp)]
        drop6 = [r["cidr"] for r in spamhaus_records("drop_v6", tmp)]
        asns = {str(r["asn"]) for r in spamhaus_records("asndrop", tmp)}
        asn4 = asn_prefixes(asns, 4, tmp)
        asn6 = asn_prefixes(asns, 6, tmp)
        if len(asn4) < MIN_ASN_PREFIXES_V4:
            fail(f"only {len(asn4)} IPv4 prefixes for {len(asns)} ASNs")
        sets = {}
        for family, name, drop, asn in ((4, "spamhaus_v4", drop4, asn4), (6, "spamhaus_v6", drop6, asn6)):
            sets[name], skipped = merge(family, drop, asn)
            print(f"{name}: {len(drop)} DROP + {len(asn)} ASN-DROP prefixes -> "
                  f"{len(sets[name])} merged, {skipped} non-global skipped")
        print(f"ASN-DROP: {len(asns)} ASNs")

        def elements(networks):
            return ",\n\t\t".join(str(n) for n in networks) + ","

        persisted = ["# Generated by nft-host-spamhaus, do not edit.", "table inet host {"]
        for name, kind in (("spamhaus_v4", "ipv4_addr"), ("spamhaus_v6", "ipv6_addr")):
            persisted.append(f"\tset {name} {{ type {kind}; flags interval; auto-merge; elements = {{\n"
                             f"\t\t{elements(sets[name])}\n\t}} }}")
        persisted.append("}\n")

        candidate = os.path.join(tmp, "host-spamhaus.nft")
        with open(candidate, "w", encoding="utf-8") as f:
            f.write("\n".join(persisted))
        # Checked under another table name: a dry run against the live table
        # name changes the live sets for its duration and unlisted addresses match.
        with open(candidate, encoding="utf-8") as f:
            check_rules = f.read().replace("table inet host {", "table inet host_spamhaus_check {")
        checked = subprocess.run(["nft", "-c", "-f", "-"], input=check_rules, capture_output=True, text=True)
        if checked.returncode != 0:
            fail(f"nft rejects the sets: {checked.stderr.strip()}")

        previous = None
        if os.path.exists(PERSISTED):
            previous = os.path.join(tmp, "previous.nft")
            shutil.copyfile(PERSISTED, previous)
        shutil.copyfile(candidate, PERSISTED + ".new")
        os.chmod(PERSISTED + ".new", 0o644)
        os.replace(PERSISTED + ".new", PERSISTED)

        # A reload replaces the whole table, sets included, in one transaction.
        # Flushing and refilling a running interval set instead lets packets of
        # unlisted addresses match while the transaction commits.
        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 sets: {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 sets on any surprise
    fail(f"{type(e).__name__}: {e}")
