#!/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 nft(*args):
    return subprocess.run(["nft", *args], capture_output=True, text=True)


def write(path, text):
    with open(path, "w", encoding="utf-8") as f:
        f.write(text)


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 {"]
        update = []
        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}} }}")
            update.append(f"flush set inet host {name}\n"
                          f"add element inet host {name} {{\n\t{elements(sets[name])}\n}}")
        persisted.append("}\n")

        persisted_path = os.path.join(tmp, "host-spamhaus.nft")
        update_path = os.path.join(tmp, "update.nft")
        write(persisted_path, "\n".join(persisted))
        write(update_path, "\n".join(update) + "\n")

        # The running table gets the new elements in one transaction; without
        # the sets (boot, or a rule set rendered without them) the persisted
        # file is the only copy.
        loaded = all(nft("list", "set", "inet", "host", n).returncode == 0
                     for n in ("spamhaus_v4", "spamhaus_v6"))
        for path in (persisted_path,) + ((update_path,) if loaded else ()):
            checked = nft("-c", "-f", path)
            if checked.returncode != 0:
                fail(f"nft rejects {os.path.basename(path)}: {checked.stderr.strip()}")
        if loaded:
            applied = nft("-f", update_path)
            if applied.returncode != 0:
                fail(f"loading the sets failed: {applied.stderr.strip()}")
        shutil.copyfile(persisted_path, PERSISTED + ".new")
        os.chmod(PERSISTED + ".new", 0o644)
        os.replace(PERSISTED + ".new", PERSISTED)
        print("loaded into the running table" if loaded 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}")
