#!/usr/bin/env python3
"""
build-country-cidr.py

Generate per-country CIDR files from RIPE / RIR data, without depending on a
third party like ipdeny. Combines two sources:

  1. RIR delegation statistics ("delegated-*-extended-latest")
     -> country assigned at the ALLOCATION level to the LIR.
  2. (optional) RIPE database inetnum/inet6num dumps
     -> country DECLARED by the resource holder (whois "country:"),
        which catches more-specific assignments the delegation file misses
        (e.g. 77.81.76.0/24, whois=IR but not delegated under IR).

The two sets are unioned and aggregated per address family.

SAFETY / DURABILITY (this is the important part):
  * Every source is downloaded IN FULL to a temporary file before anything is
    parsed. A connection that drops mid-download never contaminates the result.
  * The run is ALL-OR-NOTHING: if any enabled source fails to download or
    fails to parse, NO output files are written and all existing .cidr files
    are left exactly as they were.
  * Output is written atomically (temp file + os.replace), so a file is never
    seen half-written or truncated, even if the process is killed mid-write.
  * A shrink guard refuses to overwrite an existing file if the new result is
    empty or drops below MIN_KEEP_RATIO of the current file's size -- this
    protects you from a source that returns a broken/partial-but-valid file.

Net effect: your current .cidr files are never deleted or truncated on
failure. The worst case of any failure is "files unchanged, exit non-zero".

Stdlib only (urllib + gzip + ipaddress). Python 3.6+.
"""

import gzip
import ipaddress
import os
import sys
import tempfile
import urllib.request

# ----------------------------- configuration ------------------------------- #

COUNTRIES = {"IR", "IQ"}                 # ISO 3166-1 alpha-2, upper-case
OUT_DIR   = "/var/www/download/ipdb"
WORK_DIR  = None                         # temp download dir; None = system temp
WRITE_IPV6 = False

# Source 1: RIR delegation statistics (extended format).
RIRS = {
    "ripencc": "https://ftp.ripe.net/pub/stats/ripencc/delegated-ripencc-extended-latest",
    # "arin":    "https://ftp.arin.net/pub/stats/arin/delegated-arin-extended-latest",
    # "apnic":   "https://ftp.apnic.net/pub/stats/apnic/delegated-apnic-extended-latest",
    # "afrinic": "https://ftp.afrinic.net/pub/stats/afrinic/delegated-afrinic-extended-latest",
    # "lacnic":  "https://ftp.lacnic.net/pub/stats/lacnic/delegated-lacnic-extended-latest",
}
KEEP_STATUS = {"allocated", "assigned"}

# Source 2: RIPE database dump (whois "country:"). Set False for delegation-only.
USE_WHOIS_DB     = True
RIPE_DB_INETNUM  = "https://ftp.ripe.net/ripe/dbase/split/ripe.db.inetnum.gz"
RIPE_DB_INET6NUM = "https://ftp.ripe.net/ripe/dbase/split/ripe.db.inet6num.gz"

TIMEOUT    = 60     # seconds, delegation files (small)
TIMEOUT_DB = 600    # seconds, RIPE DB dump (large)

# Permissions for generated .cidr files (octal). 0o644 = owner rw, group/other r.
FILE_MODE = 0o644

# Safety guards.
MIN_DOWNLOAD_BYTES = 64     # treat a smaller "download" as a failed fetch
MIN_KEEP_RATIO     = 0.5    # refuse to overwrite if new < 50% of existing size;
                            # set to 0 to disable the shrink guard.

# --------------------------------------------------------------------------- #


def download_to_temp(url, timeout, workdir):
    """Download url fully to a temp file. Raise on any error or short read.
    Returns the temp path on success; cleans up on failure."""
    fd, tmp = tempfile.mkstemp(dir=workdir, prefix=".dl-")
    os.close(fd)
    try:
        req = urllib.request.Request(url, headers={"User-Agent": "country-cidr/1.0"})
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            hdr = None
            headers = getattr(resp, "headers", None)
            if headers is not None:
                try:
                    hdr = headers.get("Content-Length")
                except Exception:
                    hdr = None
            expected = int(hdr) if hdr and str(hdr).isdigit() else None
            written = 0
            with open(tmp, "wb") as out:
                while True:
                    chunk = resp.read(1 << 20)
                    if not chunk:
                        break
                    out.write(chunk)
                    written += len(chunk)
        if expected is not None and written != expected:
            raise IOError("incomplete download: got %d of %d bytes"
                          % (written, expected))
        if written < MIN_DOWNLOAD_BYTES:
            raise IOError("download suspiciously small (%d bytes)" % written)
        return tmp
    except BaseException:
        if os.path.exists(tmp):
            os.unlink(tmp)
        raise


def text_lines(path):
    with open(path, "r", encoding="utf-8", errors="replace") as fh:
        for line in fh:
            yield line


def gz_lines(path):
    # Iterating a truncated/corrupt gzip raises -> surfaces as a parse failure.
    with gzip.open(path, "rt", encoding="utf-8", errors="replace") as fh:
        for line in fh:
            yield line


# ---- source 1: delegation statistics -------------------------------------- #

def parse_delegation(lines, wanted, v4, v6):
    for line in lines:
        if not line or line[0] == "#" or "|" not in line:
            continue
        f = line.rstrip("\n").split("|")
        if len(f) < 7:
            continue
        _registry, cc, typ, start, value, _date, status = f[:7]
        if cc not in wanted or status not in KEEP_STATUS:
            continue
        try:
            if typ == "ipv4":
                first = ipaddress.IPv4Address(start)
                last = first + int(value) - 1
                v4[cc].extend(ipaddress.summarize_address_range(first, last))
            elif typ == "ipv6":
                v6[cc].append(ipaddress.ip_network("%s/%s" % (start, value)))
        except (ipaddress.AddressValueError, ValueError) as e:
            print("  skip bad delegation record: %s (%s)"
                  % (line.strip(), e), file=sys.stderr)


# ---- source 2: RIPE DB inetnum / inet6num --------------------------------- #

def parse_ripe_db(lines, wanted, v4, v6):
    r4 = [None]
    n6 = [None]
    ccs = [set()]

    def flush():
        hit = ccs[0] & wanted
        if hit:
            if r4[0]:
                try:
                    nets = list(ipaddress.summarize_address_range(
                        ipaddress.IPv4Address(r4[0][0]),
                        ipaddress.IPv4Address(r4[0][1])))
                    for cc in hit:
                        v4[cc].extend(nets)
                except (ipaddress.AddressValueError, ValueError):
                    pass
            if n6[0]:
                try:
                    net = ipaddress.ip_network(n6[0], strict=False)
                    for cc in hit:
                        v6[cc].append(net)
                except (ipaddress.AddressValueError, ValueError):
                    pass
        r4[0] = None
        n6[0] = None
        ccs[0] = set()

    for raw in lines:
        line = raw.rstrip("\r\n")
        if not line:
            flush()
            continue
        c0 = line[0]
        if c0 in "#%" or c0 in " \t+" or ":" not in line:
            continue
        attr, _, val = line.partition(":")
        attr = attr.strip().lower()
        val = val.strip()
        if attr == "inetnum":
            a = b = None
            if " - " in val:
                a, b = val.split(" - ", 1)
            elif "-" in val:
                a, b = val.split("-", 1)
            if a:
                r4[0] = (a.strip(), b.strip())
        elif attr == "inet6num":
            n6[0] = val
        elif attr == "country":
            if val:
                ccs[0].add(val[:2].upper())
    flush()


# ---- output --------------------------------------------------------------- #

def count_lines(path):
    try:
        with open(path) as fh:
            return sum(1 for _ in fh)
    except OSError:
        return 0


def write_atomic(path, networks):
    d = os.path.dirname(path) or "."
    fd, tmp = tempfile.mkstemp(dir=d, prefix=".tmp-")
    try:
        with os.fdopen(fd, "w") as fh:
            for net in networks:
                fh.write(str(net) + "\n")
        # mkstemp creates the file 0600; set the intended mode BEFORE the
        # rename so the published file is atomically visible with FILE_MODE.
        os.chmod(tmp, FILE_MODE)
        os.replace(tmp, path)          # atomic on same filesystem
    except BaseException:
        if os.path.exists(tmp):
            os.unlink(tmp)
        raise


def guarded_write(path, networks, label):
    """Write only if the result passes the safety guards; otherwise keep the
    existing file. Returns (written, message)."""
    if not networks:
        return False, "empty result -- kept existing file"
    existing = count_lines(path)
    if existing and MIN_KEEP_RATIO > 0 and len(networks) < existing * MIN_KEEP_RATIO:
        return False, ("new %s=%d < %.0f%% of existing %d -- kept existing file"
                       % (label, len(networks), MIN_KEEP_RATIO * 100, existing))
    write_atomic(path, networks)
    return True, None


# ---- orchestration -------------------------------------------------------- #

def main():
    os.makedirs(OUT_DIR, exist_ok=True)
    workdir = WORK_DIR or tempfile.gettempdir()
    os.makedirs(workdir, exist_ok=True)

    temps = []
    deleg_files = []
    whois_files = []
    failures = []

    # ---- phase 1: download EVERYTHING first (nothing parsed yet) ---------- #
    try:
        for name, url in RIRS.items():
            print("Downloading delegation stats: %s ..." % name)
            try:
                p = download_to_temp(url, TIMEOUT, workdir)
                temps.append(p)
                deleg_files.append((name, p))
            except Exception as e:
                failures.append(("delegation:" + name, e))

        if USE_WHOIS_DB:
            print("Downloading RIPE DB inetnum (large) ...")
            try:
                p = download_to_temp(RIPE_DB_INETNUM, TIMEOUT_DB, workdir)
                temps.append(p)
                whois_files.append(("inetnum", p))
            except Exception as e:
                failures.append(("whois:inetnum", e))
            if WRITE_IPV6:
                print("Downloading RIPE DB inet6num ...")
                try:
                    p = download_to_temp(RIPE_DB_INET6NUM, TIMEOUT_DB, workdir)
                    temps.append(p)
                    whois_files.append(("inet6num", p))
                except Exception as e:
                    failures.append(("whois:inet6num", e))

        if not deleg_files and not whois_files:
            failures.append(("all-sources", "no source downloaded"))

        if failures:
            for name, err in failures:
                print("  DOWNLOAD FAILED [%s]: %s" % (name, err), file=sys.stderr)
            print("Aborting BEFORE writing. Existing .cidr files left untouched.",
                  file=sys.stderr)
            return 2

        # ---- phase 2: parse from local temp files (gated) ---------------- #
        v4 = {cc: [] for cc in COUNTRIES}
        v6 = {cc: [] for cc in COUNTRIES}
        for name, p in deleg_files:
            parse_delegation(text_lines(p), COUNTRIES, v4, v6)
        deleg4 = {cc: set(ipaddress.collapse_addresses(v4[cc])) for cc in COUNTRIES}
        for kind, p in whois_files:
            parse_ripe_db(gz_lines(p), COUNTRIES, v4, v6)

    except Exception as e:
        # Any parse/decompress error (e.g. truncated gzip) lands here.
        print("  PARSE FAILED: %s" % e, file=sys.stderr)
        print("Aborting BEFORE writing. Existing .cidr files left untouched.",
              file=sys.stderr)
        return 2
    finally:
        for p in temps:
            try:
                os.unlink(p)
            except OSError:
                pass

    # ---- phase 3: sanity-checked atomic writes --------------------------- #
    rc = 0
    for cc in sorted(COUNTRIES):
        merged4 = list(ipaddress.collapse_addresses(v4[cc]))
        new4 = len(set(merged4) - deleg4[cc])
        path4 = os.path.join(OUT_DIR, "%s.cidr" % cc.lower())
        wrote, msg = guarded_write(path4, merged4, "IPv4")
        if wrote:
            print("%s: wrote %d IPv4 CIDRs (delegation %d, +%d net-new from whois)"
                  % (cc, len(merged4), len(deleg4[cc]), new4))
        else:
            rc = 1
            print("%s: NOT written (%s)" % (cc, msg), file=sys.stderr)

        if WRITE_IPV6:
            merged6 = list(ipaddress.collapse_addresses(v6[cc]))
            path6 = os.path.join(OUT_DIR, "%s.v6.cidr" % cc.lower())
            wrote6, msg6 = guarded_write(path6, merged6, "IPv6")
            if wrote6:
                print("%s: wrote %d IPv6 CIDRs" % (cc, len(merged6)))
            elif merged6:
                rc = 1
                print("%s: IPv6 NOT written (%s)" % (cc, msg6), file=sys.stderr)
    return rc


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