#!/usr/bin/env python3
"""
aria_geoblock.py — Geoblocking via ipset + iptables

Liest eine .ini-Datei mit einer Laenderliste, laedt die zugehoerigen
IP-Bereiche (CIDR) kostenlos von ipdeny.com (kein API-Key noetig) und
blockt sie am Linux-Kernel-Paketfilter (iptables + ipset).

WARUM ipset+iptables statt "Verbindung live pruefen":
  - Live pro Connection ein GeoIP-Lookup zu machen ist langsam, braucht
    eine (meist kostenpflichtige/registrierungspflichtige) GeoIP-DB und
    blockt den Traffic erst NACH dem TCP-Handshake.
  - ipset haelt die komplette IP-Range-Liste jedes Landes im Kernel vor;
    iptables matched dagegen in O(1)/sehr schnell und droppt Pakete
    schon auf Netzwerk-Ebene, bevor sie den Server ueberhaupt erreichen.
  - Das ist der Standard-Ansatz den auch fail2ban/CSF-Firewalls nutzen.

VORAUSSETZUNGEN (auf dem Ziel-Host, NICHT im ARIA-Container):
  - root-Rechte (iptables/ipset aendern den Kernel-Netfilter)
  - Pakete installiert: iptables, ipset
      Debian/Ubuntu: apt install iptables ipset
  - Internetzugang zum Laden der Zonefiles von ipdeny.com

BENUTZUNG:
  sudo python3 aria_geoblock.py --config geoblock.ini --apply
      -> laedt IP-Ranges, befuellt das ipset, setzt die iptables-Regel

  sudo python3 aria_geoblock.py --config geoblock.ini --apply --dry-run
      -> zeigt nur an was gemacht wuerde, aendert nichts

  sudo python3 aria_geoblock.py --config geoblock.ini --remove
      -> entfernt die iptables-Regel und loescht das ipset wieder

  sudo python3 aria_geoblock.py --config geoblock.ini --status
      -> zeigt aktuellen Stand (Set vorhanden? wie viele Eintraege? Regel aktiv?)

.ini-Format siehe mitgelieferte aria_geoblock.ini (Beispiel).

LOGGING:
  Jeder Lauf schreibt nach stderr (systemd-Journal). Zusaetzlich kann in
  der .ini unter "log_file" ein Pfad angegeben werden — dort wird bei
  jedem Lauf eine Zeile angehaengt, inkl. einer maschinenlesbaren
  "RESULT ..."-Zeile (status=OK/ERROR, Anzahl Laender/IP-Bereiche). Genau
  diese Zeile wertet das mitgelieferte checkmk-Plugin aus.

AUTOMATISCH AKTUELL HALTEN (Systemstart + taeglich):
  Siehe systemd/aria-geoblock.service + aria-geoblock.timer sowie
  install_aria_geoblock.sh im gleichen Repo — das richtet Timer fuer
  "beim Booten" + "einmal taeglich" automatisch ein.
  --apply ist idempotent: es flusht das bestehende Set und befuellt es neu,
  legt aber KEINE doppelte iptables-Regel an (wird vorher geprueft).

PORT-EINSCHRAENKUNG (optional):
  Per Default blockt das Script ALLE Ports/Protokolle fuer die gelisteten
  Laender (klassisches Geoblocking). Ueber "ports" in der .ini kann man das
  auf bestimmte Anwendungen/Ports eingrenzen — z.B. nur den Testserver
  (8899) blocken, waehrend SSH (22) fuer die Laender offen bleibt. Praktisch
  zum gefahrlosen Testen, damit man sich nicht selbst aussperrt. Siehe
  Kommentare in aria_geoblock.ini.
"""

import argparse
import configparser
import datetime
import ipaddress
import subprocess
import sys
import time
import urllib.request
import urllib.error

ZONE_URL_TMPL = "https://www.ipdeny.com/ipblocks/data/countries/{cc}.zone"

# Wird in main() aus der .ini gesetzt (log_file). None = keine Datei-Logs.
LOG_FILE = None


def _timestamp():
    return datetime.datetime.now().astimezone().isoformat(timespec="seconds")


def _write_log_file(line):
    if not LOG_FILE:
        return
    try:
        with open(LOG_FILE, "a", encoding="utf-8") as f:
            f.write(f"{_timestamp()} {line}\n")
    except OSError as e:
        print(f"[geoblock] WARNUNG: konnte Log-Datei nicht schreiben ({LOG_FILE}): {e}",
              file=sys.stderr)


def log(msg):
    line = f"[geoblock] {msg}"
    print(line, file=sys.stderr)
    _write_log_file(line)


def log_result(status, action, countries=None, entries=None, msg=""):
    """Schreibt eine strukturierte, maschinenlesbare Ergebniszeile.

    Wird vom checkmk-Plugin (aria_geoblock_checkmk) ausgewertet, um den
    Zustand des letzten Laufs zu beurteilen (OK/WARN/CRIT/UNKNOWN).
    """
    parts = [f"RESULT ts={_timestamp()}", f"status={status}", f"action={action}"]
    if countries is not None:
        parts.append(f"countries={countries}")
    if entries is not None:
        parts.append(f"entries={entries}")
    safe_msg = msg.replace('"', "'").replace("\n", " ")
    parts.append(f'msg="{safe_msg}"')
    line = "[geoblock] " + " ".join(parts)
    print(line, file=sys.stderr)
    _write_log_file(line)


def run(cmd, dry_run=False, check=True):
    """Fuehrt einen Shell-Befehl (als Liste) aus. Bei dry_run nur anzeigen."""
    printable = " ".join(cmd)
    if dry_run:
        log(f"DRY-RUN: {printable}")
        return None
    result = subprocess.run(cmd, capture_output=True, text=True)
    if check and result.returncode != 0:
        log(f"FEHLER bei: {printable}")
        log(f"stderr: {result.stderr.strip()}")
        raise RuntimeError(f"Befehl fehlgeschlagen: {printable}")
    return result


def load_config(path):
    cfg = configparser.ConfigParser()
    read_ok = cfg.read(path)
    if not read_ok:
        raise FileNotFoundError(f"Konnte .ini nicht lesen: {path}")
    if "geoblock" not in cfg:
        raise ValueError("Abschnitt [geoblock] fehlt in der .ini-Datei")

    sec = cfg["geoblock"]
    countries_raw = sec.get("countries", "")
    countries = [c.strip().lower() for c in countries_raw.split(",") if c.strip()]
    if not countries:
        raise ValueError("Keine Laender in 'countries' angegeben")

    # ports: leer/weggelassen = alle Ports blocken (klassisches Geoblocking).
    # Angegeben (z.B. "80,443,8899") = nur diese Ports blocken.
    ports_raw = sec.get("ports", "").strip()
    ports = [p.strip() for p in ports_raw.split(",") if p.strip()]
    for p in ports:
        _validate_port_spec(p)

    # protocol: welche(s) Protokoll(e) die Port-Einschraenkung betreffen.
    # Nur relevant wenn "ports" gesetzt ist. Default: tcp. Auch "tcp,udp" moeglich.
    protocols_raw = sec.get("protocol", "tcp").strip()
    protocols = [pr.strip().lower() for pr in protocols_raw.split(",") if pr.strip()] or ["tcp"]
    for pr in protocols:
        if pr not in ("tcp", "udp"):
            raise ValueError(f"Ungueltiges Protokoll '{pr}' in 'protocol' (erlaubt: tcp, udp)")

    return {
        "countries": countries,
        "chain": sec.get("chain", "INPUT").strip(),
        "interface": sec.get("interface", "").strip(),
        "log": sec.getboolean("log", fallback=False),
        "log_prefix": sec.get("log_prefix", "GEOBLOCK-DROP:").strip(),
        "ipset_name": sec.get("ipset_name", "geoblock").strip(),
        "log_file": sec.get("log_file", "").strip(),
        "ports": ports,
        "protocols": protocols,
    }


def _validate_port_spec(spec):
    """Erlaubt Einzelport ('443') oder Bereich ('8000:9000'), wie iptables --dports."""
    parts = spec.split(":")
    if len(parts) not in (1, 2):
        raise ValueError(f"Ungueltige Port-Angabe '{spec}' in 'ports' (erlaubt: '443' oder '8000:9000')")
    for part in parts:
        if not part.isdigit() or not (0 <= int(part) <= 65535):
            raise ValueError(f"Ungueltige Port-Angabe '{spec}' in 'ports' (Ports muessen 0-65535 sein)")


def build_rule_variants(cfg):
    """Baut die Liste der (proto, dports)-Kombinationen fuer iptables-Regeln.

    Ohne 'ports' in der .ini: eine Variante ohne Proto-/Port-Einschraenkung
    (= klassisches Geoblocking, alle Ports/Protokolle). Mit 'ports': eine
    Variante pro konfiguriertem Protokoll (multiport braucht -p tcp ODER -p udp,
    nicht beides gleichzeitig).
    """
    if not cfg["ports"]:
        return [{"proto": None, "dports": None}]
    dports = ",".join(cfg["ports"])
    return [{"proto": proto, "dports": dports} for proto in cfg["protocols"]]


def _variant_match_args(variant):
    args = []
    if variant["proto"]:
        args += ["-p", variant["proto"]]
    return args


def _variant_port_args(variant):
    if variant["dports"]:
        return ["-m", "multiport", "--dports", variant["dports"]]
    return []


def _variant_label(variant):
    if not variant["dports"]:
        return "alle Ports"
    return f"{variant['proto']}/{variant['dports']}"


def fetch_country_cidrs(country_code):
    """Laedt die CIDR-Liste eines Landes von ipdeny.com. Gibt Liste von Strings zurueck."""
    url = ZONE_URL_TMPL.format(cc=country_code)
    log(f"Lade IP-Ranges fuer '{country_code}' von {url}")
    try:
        with urllib.request.urlopen(url, timeout=20) as resp:
            text = resp.read().decode("utf-8", errors="ignore")
    except urllib.error.HTTPError as e:
        raise RuntimeError(
            f"Laendercode '{country_code}' unbekannt oder ipdeny nicht erreichbar "
            f"(HTTP {e.code}). Pruefe den ISO-3166-1-alpha-2 Code."
        )
    except urllib.error.URLError as e:
        raise RuntimeError(f"Konnte {url} nicht laden: {e}")

    cidrs = []
    for line in text.splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        try:
            ipaddress.ip_network(line)  # Validierung
        except ValueError:
            log(f"  ignoriere ungueltige Zeile: {line}")
            continue
        cidrs.append(line)
    log(f"  -> {len(cidrs)} IP-Bereiche fuer '{country_code}'")
    return cidrs


def ipset_exists(name):
    result = subprocess.run(["ipset", "list", "-n"], capture_output=True, text=True)
    return name in result.stdout.splitlines()


def iptables_rule_exists(chain, ipset_name, interface, variant=None, dry_run=False):
    variant = variant or {"proto": None, "dports": None}
    check_cmd = ["iptables", "-C", chain]
    if interface:
        check_cmd += ["-i", interface]
    check_cmd += _variant_match_args(variant)
    check_cmd += ["-m", "set", "--match-set", ipset_name, "src"]
    check_cmd += _variant_port_args(variant)
    check_cmd += ["-j", "DROP"]
    if dry_run:
        # Im Dry-Run wissen wir es nicht sicher, nehmen konservativ "existiert nicht" an
        return False
    result = subprocess.run(check_cmd, capture_output=True, text=True)
    return result.returncode == 0


def apply_geoblock(cfg, dry_run=False):
    name = cfg["ipset_name"]

    # 1) ipset anlegen falls noetig
    if dry_run or not ipset_exists(name):
        run(["ipset", "create", name, "hash:net", "-exist"], dry_run=dry_run)
    else:
        log(f"ipset '{name}' existiert bereits, wird neu befuellt")

    # 2) alle CIDRs aller konfigurierten Laender sammeln
    all_cidrs = []
    for cc in cfg["countries"]:
        all_cidrs.extend(fetch_country_cidrs(cc))

    if not all_cidrs and not dry_run:
        raise RuntimeError("Keine IP-Bereiche geladen — Abbruch, um kein leeres Set zu aktivieren")

    # 3) Set atomar neu befuellen: temp-Set bauen, dann swap (kein Traffic-Loch)
    tmp_name = f"{name}_tmp"
    run(["ipset", "create", tmp_name, "hash:net", "-exist"], dry_run=dry_run)
    run(["ipset", "flush", tmp_name], dry_run=dry_run)
    for cidr in all_cidrs:
        run(["ipset", "add", tmp_name, cidr, "-exist"], dry_run=dry_run, check=False)
    run(["ipset", "create", name, "hash:net", "-exist"], dry_run=dry_run)
    run(["ipset", "swap", tmp_name, name], dry_run=dry_run)
    run(["ipset", "destroy", tmp_name], dry_run=dry_run, check=False)

    log(f"ipset '{name}' befuellt mit {len(all_cidrs)} Bereichen "
        f"aus Laendern: {', '.join(c.upper() for c in cfg['countries'])}")

    # 4) optionale LOG-Regel(n) + DROP-Regel(n) in iptables, nur wenn noch nicht vorhanden.
    #    Eine Variante pro Protokoll wenn 'ports' gesetzt ist, sonst eine Variante (alle Ports).
    set_match = ["-m", "set", "--match-set", name, "src"]
    iface_opt = ["-i", cfg["interface"]] if cfg["interface"] else []
    variants = build_rule_variants(cfg)

    for variant in variants:
        proto_args = _variant_match_args(variant)
        port_args = _variant_port_args(variant)
        label = _variant_label(variant)

        if cfg["log"]:
            log_check = (["iptables", "-C", cfg["chain"]] + iface_opt + proto_args + set_match +
                         port_args + ["-j", "LOG", "--log-prefix", cfg["log_prefix"] + " "])
            exists = False if dry_run else subprocess.run(log_check, capture_output=True).returncode == 0
            if not exists:
                run(["iptables", "-I", cfg["chain"]] + iface_opt + proto_args + set_match +
                    port_args + ["-j", "LOG", "--log-prefix", cfg["log_prefix"] + " "], dry_run=dry_run)
            else:
                log(f"LOG-Regel ({label}) existiert bereits, ueberspringe")

        if iptables_rule_exists(cfg["chain"], name, cfg["interface"], variant=variant, dry_run=dry_run):
            log(f"DROP-Regel ({label}) existiert bereits, ueberspringe (idempotent)")
        else:
            run(["iptables", "-A", cfg["chain"]] + iface_opt + proto_args + set_match +
                port_args + ["-j", "DROP"], dry_run=dry_run)
            log(f"DROP-Regel in Chain '{cfg['chain']}' aktiv fuer Set '{name}' ({label})")

    log("Fertig. Hinweis: Regeln sind NICHT reboot-persistent — "
        "fuer Persistenz z.B. 'iptables-persistent' bzw. 'netfilter-persistent save' nutzen, "
        "und das ipset per systemd-Unit / rc.local vor dem iptables-Restore neu befuellen.")

    return len(cfg["countries"]), len(all_cidrs)


def remove_geoblock(cfg, dry_run=False):
    name = cfg["ipset_name"]
    iface_opt = ["-i", cfg["interface"]] if cfg["interface"] else []
    set_match = ["-m", "set", "--match-set", name, "src"]
    # Entfernt fuer alle aktuell konfigurierten Varianten. Falls die .ini seit dem letzten
    # --apply geaendert wurde (z.B. ports entfernt/hinzugefuegt), koennen dadurch alte
    # Regeln mit anderer Variante stehen bleiben -- im Zweifel per 'iptables -L <chain> -n'
    # pruefen und manuell aufraeumen.
    for variant in build_rule_variants(cfg):
        proto_args = _variant_match_args(variant)
        port_args = _variant_port_args(variant)

        for _ in range(5):
            result = run(["iptables", "-D", cfg["chain"]] + iface_opt + proto_args + set_match +
                          port_args + ["-j", "DROP"], dry_run=dry_run, check=False)
            if dry_run or result is None or result.returncode != 0:
                break

        if cfg["log"]:
            for _ in range(5):
                result = run(["iptables", "-D", cfg["chain"]] + iface_opt + proto_args + set_match +
                              port_args + ["-j", "LOG", "--log-prefix", cfg["log_prefix"] + " "],
                              dry_run=dry_run, check=False)
                if dry_run or result is None or result.returncode != 0:
                    break

    run(["ipset", "destroy", name], dry_run=dry_run, check=False)
    log(f"Geoblock-Regeln entfernt, ipset '{name}' geloescht.")


def show_status(cfg):
    name = cfg["ipset_name"]
    if ipset_exists(name):
        result = subprocess.run(["ipset", "list", name, "-t"], capture_output=True, text=True)
        print(result.stdout)
    else:
        print(f"ipset '{name}' existiert nicht (noch nicht aktiv).")

    variants = build_rule_variants(cfg)
    for variant in variants:
        active = iptables_rule_exists(cfg["chain"], name, cfg["interface"], variant=variant)
        print(f"iptables DROP-Regel in Chain '{cfg['chain']}' ({_variant_label(variant)}) aktiv: {active}")
    print(f"Konfigurierte Laender: {', '.join(c.upper() for c in cfg['countries'])}")
    if cfg["ports"]:
        print(f"Port-Einschraenkung: {', '.join(cfg['ports'])} ({'/'.join(cfg['protocols'])})")
    else:
        print("Port-Einschraenkung: keine (alle Ports/Protokolle werden geblockt)")


def main():
    global LOG_FILE

    parser = argparse.ArgumentParser(description="Geoblocking per Laenderliste (.ini) via ipset+iptables")
    parser.add_argument("--config", required=True, help="Pfad zur .ini-Konfigurationsdatei")
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--apply", action="store_true", help="Laender laden und Block-Regeln setzen/aktualisieren")
    group.add_argument("--remove", action="store_true", help="Block-Regeln und ipset wieder entfernen")
    group.add_argument("--status", action="store_true", help="Aktuellen Zustand anzeigen")
    parser.add_argument("--dry-run", action="store_true", help="Nur anzeigen was passieren wuerde, nichts aendern")
    args = parser.parse_args()

    action = "apply" if args.apply else ("remove" if args.remove else "status")

    try:
        cfg = load_config(args.config)
        LOG_FILE = cfg["log_file"] or None

        if args.apply:
            n_countries, n_entries = apply_geoblock(cfg, dry_run=args.dry_run)
            log_result("OK", action, countries=n_countries, entries=n_entries,
                        msg=f"Geoblock aktualisiert ({n_countries} Laender, {n_entries} IP-Bereiche)")
        elif args.remove:
            remove_geoblock(cfg, dry_run=args.dry_run)
            log_result("OK", action, msg="Geoblock-Regeln entfernt")
        elif args.status:
            show_status(cfg)
    except Exception as e:
        log_result("ERROR", action, msg=str(e))
        log(f"FEHLER: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()