reputation: firehol netset cache with local range index

src/netcheck/reputation.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import ipaddress
+import os
+import time
+from bisect import bisect_right
+from pathlib import Path
+
+import httpx
+
+from netcheck.config import Providers
+
+
+def parse_netset(text: str) -> list[str]:
+    entries: list[str] = []
+    for line in text.splitlines():
+        stripped = line.strip()
+        if not stripped or stripped.startswith("#"):
+            continue
+        entries.append(stripped)
+    return entries
+
+
+class NetsetIndex:
+    def __init__(self) -> None:
+        # Sorted (start, end, name) ranges per address family; a lookup is one
+        # bisect plus a short backward scan, which beats a per-CIDR loop over
+        # the ~1M prefixes a full FireHOL set contains.
+        self._ranges: dict[int, list[tuple[int, int, str]]] = {4: [], 6: []}
+        self._starts: dict[int, list[int]] = {4: [], 6: []}
+        self._dirty = False
+
+    def add(self, name: str, cidrs: list[str]) -> None:
+        for entry in cidrs:
+            try:
+                network = ipaddress.ip_network(entry, strict=False)
+            except ValueError:
+                continue
+            self._ranges[network.version].append(
+                (int(network.network_address), int(network.broadcast_address), name)
+            )
+        self._dirty = True
+
+    def _reindex(self) -> None:
+        for version in (4, 6):
+            self._ranges[version].sort()
+            self._starts[version] = [start for start, _, _ in self._ranges[version]]
+        self._dirty = False
+
+    def hits(self, ip: str) -> list[str]:
+        try:
+            address = ipaddress.ip_address(ip)
+        except ValueError:
+            return []
+        if self._dirty:
+            self._reindex()
+        version = address.version
+        value = int(address)
+        found: list[str] = []
+        position = bisect_right(self._starts[version], value)
+        for start, end, name in reversed(self._ranges[version][:position]):
+            if end >= value and name not in found:
+                found.append(name)
+        return found
+
+
+async def refresh_netsets(
+    client: httpx.AsyncClient,
+    providers: Providers,
+    cache_dir: Path,
+) -> NetsetIndex:
+    directory = Path(cache_dir) / "firehol"
+    directory.mkdir(parents=True, exist_ok=True)
+    index = NetsetIndex()
+    for url in providers.firehol_netsets:
+        name = url.rsplit("/", 1)[-1].removesuffix(".netset")
+        path = directory / f"{name}.netset"
+        fresh = (
+            path.exists()
+            and (time.time() - os.path.getmtime(path)) < providers.firehol_refresh_hours * 3600
+        )
+        if not fresh:
+            try:
+                response = await client.get(url)
+                response.raise_for_status()
+                tmp = path.with_suffix(".netset.tmp")
+                tmp.write_text(response.text, encoding="utf-8")
+                os.replace(tmp, path)
+            except httpx.HTTPError:
+                if not path.exists():
+                    continue
+        index.add(name, parse_netset(path.read_text(encoding="utf-8")))
+    return index

tests/fixtures/api/firehol_sample.netset

@@ -0,0 +1,13 @@
+#
+# firehol_level1
+#
+# Source: example
+# Maintainer: FireHOL Team
+#
+0.0.0.0/8
+10.0.0.0/8
+198.51.100.0/24
+203.0.113.44
+2001:db8::/32
+
+192.0.2.0/25

tests/test_reputation.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+from netcheck.reputation import NetsetIndex, parse_netset
+
+
+def test_parse_netset_strips_comments_and_blank_lines(fixtures_dir):
+    cidrs = parse_netset((fixtures_dir / "api" / "firehol_sample.netset").read_text(encoding="utf-8"))
+    assert cidrs == [
+        "0.0.0.0/8",
+        "10.0.0.0/8",
+        "198.51.100.0/24",
+        "203.0.113.44",
+        "2001:db8::/32",
+        "192.0.2.0/25",
+    ]
+
+
+def test_parse_netset_of_an_empty_file_is_empty():
+    assert parse_netset("# only a comment\n\n") == []
+
+
+def test_index_matches_an_ip_inside_a_listed_prefix():
+    index = NetsetIndex()
+    index.add("firehol_level1", ["198.51.100.0/24"])
+    assert index.hits("198.51.100.7") == ["firehol_level1"]
+
+
+def test_index_does_not_match_an_ip_outside_every_prefix():
+    index = NetsetIndex()
+    index.add("firehol_level1", ["198.51.100.0/24"])
+    assert index.hits("203.0.113.7") == []
+
+
+def test_index_matches_a_bare_host_entry():
+    index = NetsetIndex()
+    index.add("abusers", ["203.0.113.44"])
+    assert index.hits("203.0.113.44") == ["abusers"]
+    assert index.hits("203.0.113.45") == []
+
+
+def test_index_reports_every_list_an_ip_appears_on():
+    index = NetsetIndex()
+    index.add("level1", ["198.51.100.0/24"])
+    index.add("level2", ["198.51.100.0/25"])
+    index.add("clean", ["203.0.113.0/24"])
+    assert sorted(index.hits("198.51.100.7")) == ["level1", "level2"]
+
+
+def test_index_handles_ipv6_separately_from_ipv4():
+    index = NetsetIndex()
+    index.add("v6list", ["2001:db8::/32"])
+    index.add("v4list", ["203.0.113.0/24"])
+    assert index.hits("2001:db8::1") == ["v6list"]
+    assert index.hits("203.0.113.1") == ["v4list"]
+
+
+def test_index_ignores_malformed_entries_instead_of_raising():
+    index = NetsetIndex()
+    index.add("junk", ["not-an-ip", "999.1.1.1/24", "198.51.100.0/24"])
+    assert index.hits("198.51.100.7") == ["junk"]
+
+
+def test_index_lookup_of_a_malformed_ip_is_empty_not_an_error():
+    index = NetsetIndex()
+    index.add("level1", ["198.51.100.0/24"])
+    assert index.hits("not-an-ip") == []
+
+
+def test_index_built_from_the_real_fixture_matches_the_expected_entries(fixtures_dir):
+    index = NetsetIndex()
+    index.add(
+        "firehol_level1",
+        parse_netset((fixtures_dir / "api" / "firehol_sample.netset").read_text(encoding="utf-8")),
+    )
+    assert index.hits("10.1.2.3") == ["firehol_level1"]
+    assert index.hits("192.0.2.10") == ["firehol_level1"]
+    assert index.hits("192.0.2.200") == []