#!/usr/bin/env python3
"""Synchronize published IP Blocklist Manager lists into nftables or ipset.

Design goals:
- Fetch Phase 19 public exports only (no API key).
- Revalidate with ETag when possible.
- Validate every downloaded line before touching the firewall.
- Keep the last-known-good firewall state on 404/503/network/validation errors.
- Apply nftables as one checked batch, or ipset via create/populate/swap.
"""

from __future__ import annotations

import argparse
import hashlib
import ipaddress
import json
import logging
import os
import subprocess
import sys
import tempfile
import urllib.error
import urllib.request
from pathlib import Path

LOG = logging.getLogger("blocklist-sync")
MAX_ENTRIES = 100_000
MAX_BYTES = 8 * 1024 * 1024


def atomic_write(path: Path, data: bytes, mode: int = 0o600) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, tmp = tempfile.mkstemp(prefix=path.name + ".", dir=str(path.parent))
    try:
        os.fchmod(fd, mode)
        with os.fdopen(fd, "wb") as fh:
            fh.write(data)
            fh.flush()
            os.fsync(fh.fileno())
        os.replace(tmp, path)
    finally:
        try:
            os.unlink(tmp)
        except FileNotFoundError:
            pass


def load_json(path: Path, default):
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except FileNotFoundError:
        return default


def load_config(path: Path) -> dict:
    cfg = json.loads(path.read_text(encoding="utf-8"))
    base = str(cfg["base_url"]).rstrip("/")
    if not base.startswith("https://") and not cfg.get("allow_http", False):
        raise ValueError("base_url must use HTTPS unless allow_http=true is explicitly set")
    slug = str(cfg.get("slug", "all"))
    if not slug or any(c not in "abcdefghijklmnopqrstuvwxyz0123456789-" for c in slug):
        raise ValueError("slug contains unsupported characters")
    backend = str(cfg.get("backend", "nftables"))
    if backend not in {"nftables", "ipset", "none"}:
        raise ValueError("backend must be nftables, ipset or none")
    return {
        **cfg,
        "base_url": base,
        "slug": slug,
        "backend": backend,
        "state_dir": str(cfg.get("state_dir", "/var/lib/blocklist-sync")),
        "timeout_seconds": int(cfg.get("timeout_seconds", 20)),
    }


def validate_list(body: bytes, family: int) -> list[str]:
    if len(body) > MAX_BYTES:
        raise ValueError("download exceeds local safety limit")
    text = body.decode("utf-8", errors="strict")
    if "\x00" in text:
        raise ValueError("NUL byte in list")
    entries: list[str] = []
    for number, line in enumerate(text.splitlines(), start=1):
        if not line:
            raise ValueError(f"unexpected blank line at {number}")
        if line != line.strip() or line.startswith("#"):
            raise ValueError(f"unexpected syntax at line {number}")
        net = ipaddress.ip_network(line, strict=True)
        if net.version != family:
            raise ValueError(f"wrong address family at line {number}")
        canonical = str(net)
        # A host without /32 or /128 is accepted by ipaddress. The server may
        # legally publish IP/CIDR lines, so do not require textual equality there.
        entries.append(canonical)
        if len(entries) > MAX_ENTRIES:
            raise ValueError("too many entries")
    # Stable unique list; duplicates should not be published, but reject neither
    # availability nor a valid prior state because of harmless duplication.
    return list(dict.fromkeys(entries))


def fetch_representation(cfg: dict, state_dir: Path, family: int, metadata: dict) -> tuple[list[str], str | None, bool]:
    suffix = "ipv4.txt" if family == 4 else "ipv6.txt"
    url = f"{cfg['base_url']}/blocklists/{cfg['slug']}/{suffix}"
    cache_file = state_dir / f"{cfg['slug']}-{family}.txt"
    headers = {"Accept": "text/plain", "User-Agent": "blocklist-sync/1.0"}
    etag = metadata.get(str(family), {}).get("etag")
    if etag:
        headers["If-None-Match"] = etag
    req = urllib.request.Request(url, method="GET", headers=headers)

    try:
        response = urllib.request.urlopen(req, timeout=cfg["timeout_seconds"])
    except urllib.error.HTTPError as exc:
        if exc.code == 304:
            if not cache_file.exists():
                raise RuntimeError(f"server returned 304 for IPv{family} but no local cache exists")
            body = cache_file.read_bytes()
            return validate_list(body, family), etag, False
        if exc.code in {404, 503}:
            raise RuntimeError(f"upstream returned HTTP {exc.code} for IPv{family}; keeping last-known-good state")
        raise RuntimeError(f"upstream returned HTTP {exc.code} for IPv{family}") from exc
    except (urllib.error.URLError, TimeoutError, OSError) as exc:
        raise RuntimeError(f"network error for IPv{family}; keeping last-known-good state") from exc

    with response:
        if response.getcode() != 200:
            raise RuntimeError(f"unexpected HTTP {response.getcode()} for IPv{family}")
        content_type = response.headers.get("Content-Type", "")
        if not content_type.lower().startswith("text/plain"):
            raise RuntimeError(f"unexpected Content-Type for IPv{family}: {content_type!r}")
        body = response.read(MAX_BYTES + 1)
        entries = validate_list(body, family)
        new_etag = response.headers.get("ETag")
        atomic_write(cache_file, body)
        return entries, new_etag, True


def run(cmd: list[str], *, input_text: str | None = None) -> None:
    LOG.debug("running %s", cmd[0])
    subprocess.run(cmd, input=input_text, text=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)


def nft_apply(cfg: dict, v4: list[str], v6: list[str]) -> None:
    table_family = str(cfg.get("nft_family", "inet"))
    table = str(cfg.get("nft_table", "blocklist"))
    set4 = str(cfg.get("nft_set_v4", "blocklist4"))
    set6 = str(cfg.get("nft_set_v6", "blocklist6"))
    lines = [
        f"flush set {table_family} {table} {set4}",
        f"flush set {table_family} {table} {set6}",
    ]
    for entries, setname in ((v4, set4), (v6, set6)):
        for pos in range(0, len(entries), 1000):
            chunk = ", ".join(entries[pos:pos + 1000])
            if chunk:
                lines.append(f"add element {table_family} {table} {setname} {{ {chunk} }}")
    batch = "\n".join(lines) + "\n"
    fd, path = tempfile.mkstemp(prefix="blocklist-sync-", suffix=".nft", dir="/run" if Path("/run").is_dir() else None)
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as fh:
            fh.write(batch)
        run(["nft", "-c", "-f", path])
        run(["nft", "-f", path])
    finally:
        try:
            os.unlink(path)
        except FileNotFoundError:
            pass


def ipset_populate(live: str, temp: str, family: str, entries: list[str]) -> None:
    # hash:net cannot represent a /0. Refuse instead of silently weakening policy.
    if any(x.endswith("/0") for x in entries):
        raise RuntimeError("ipset hash:net cannot store /0; use nftables for this list")
    subprocess.run(["ipset", "destroy", temp], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    run(["ipset", "-exist", "create", live, "hash:net", "family", family, "maxelem", "131072"])
    run(["ipset", "create", temp, "hash:net", "family", family, "maxelem", "131072"])
    try:
        restore = "".join(f"add {temp} {entry}\n" for entry in entries)
        if restore:
            run(["ipset", "restore"], input_text=restore)
        run(["ipset", "swap", live, temp])
    finally:
        subprocess.run(["ipset", "destroy", temp], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


def ipset_apply(cfg: dict, v4: list[str], v6: list[str]) -> None:
    live4 = str(cfg.get("ipset_v4", "blocklist4"))
    live6 = str(cfg.get("ipset_v6", "blocklist6"))
    ipset_populate(live4, live4 + "_new", "inet", v4)
    ipset_populate(live6, live6 + "_new", "inet6", v6)


def main() -> int:
    p = argparse.ArgumentParser()
    p.add_argument("--config", default="/etc/blocklist-sync/config.json")
    p.add_argument("--verbose", action="store_true")
    args = p.parse_args()
    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO, format="%(name)s: %(levelname)s: %(message)s")

    try:
        cfg = load_config(Path(args.config))
        state_dir = Path(cfg["state_dir"])
        state_dir.mkdir(parents=True, exist_ok=True)
        metadata_file = state_dir / f"{cfg['slug']}-http.json"
        metadata = load_json(metadata_file, {})

        v4, etag4, changed4 = fetch_representation(cfg, state_dir, 4, metadata)
        v6, etag6, changed6 = fetch_representation(cfg, state_dir, 6, metadata)

        if not changed4 and not changed6:
            LOG.info("publication unchanged; firewall not modified")
            return 0

        if cfg["backend"] == "nftables":
            nft_apply(cfg, v4, v6)
        elif cfg["backend"] == "ipset":
            ipset_apply(cfg, v4, v6)

        metadata = {
            "4": {"etag": etag4, "sha256": hashlib.sha256(("\n".join(v4) + ("\n" if v4 else "")).encode()).hexdigest()},
            "6": {"etag": etag6, "sha256": hashlib.sha256(("\n".join(v6) + ("\n" if v6 else "")).encode()).hexdigest()},
        }
        atomic_write(metadata_file, (json.dumps(metadata, indent=2) + "\n").encode())
        LOG.info("applied publication slug=%s ipv4=%d ipv6=%d backend=%s", cfg["slug"], len(v4), len(v6), cfg["backend"])
        return 0
    except Exception as exc:
        LOG.error("synchronization failed: %s", exc)
        return 1


if __name__ == "__main__":
    sys.exit(main())