#!/usr/bin/env python3
"""
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 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 geoblock.py --config geoblock.ini --apply
      -> laedt IP-Ranges, befuellt das ipset, setzt die iptables-Regel

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

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

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

.ini-Format siehe mitgelieferte 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/geoblock.service + geoblock.timer sowie
  install_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 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 (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()