#!/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
    (falls der Host selbst keine direkte Verbindung hat: 'proxy' in der
    .ini setzen -- siehe README)

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 shlex
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,
        "proxy": sec.get("proxy", "").strip(),
    }


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 _build_url_opener(proxy):
    """Baut einen urllib-Opener. Mit 'proxy' gesetzt (z.B. 'http://host:3128' oder
    'http://user:pass@host:3128') wird dieser fuer http UND https genutzt -- fuer
    Server, die selbst keine direkte Internetverbindung haben und nur ueber
    einen Vorgelagerten Proxy raus koennen."""
    if not proxy:
        return urllib.request.build_opener()
    handler = urllib.request.ProxyHandler({"http": proxy, "https": proxy})
    return urllib.request.build_opener(handler)


def fetch_country_cidrs(country_code, proxy=""):
    """Laedt die CIDR-Liste eines Landes von ipdeny.com. Gibt Liste von Strings zurueck.

    'proxy' (optional): HTTP(S)-Proxy-URL, falls der Server selbst keine direkte
    Internetverbindung hat und ipdeny.com nur ueber einen Proxy erreichen kann.
    """
    url = ZONE_URL_TMPL.format(cc=country_code)
    if proxy:
        log(f"Lade IP-Ranges fuer '{country_code}' von {url} (ueber Proxy {proxy})")
    else:
        log(f"Lade IP-Ranges fuer '{country_code}' von {url}")
    opener = _build_url_opener(proxy)
    try:
        with opener.open(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 _list_chain_rules(chain):
    """Gibt die aktuellen Regeln einer iptables-Chain als '-A ...'-Zeilen zurueck
    (Format von 'iptables -S <chain>', eine Zeile pro Regel)."""
    result = subprocess.run(["iptables", "-S", chain], capture_output=True, text=True)
    if result.returncode != 0:
        return []
    return [ln for ln in result.stdout.splitlines() if ln.startswith("-A ")]


def purge_geoblock_rules(cfg, dry_run=False):
    """Loescht ALLE Regeln in der Chain, die auf das geoblock-ipset verweisen
    (--match-set <name> src) -- unabhaengig von Proto/Port/Variante.

    Dadurch werden auch alte Regeln aus frueheren --apply-Laeufen mit anderer
    ini (z.B. andere Ports/Protokolle) sauber entfernt, die build_rule_variants(cfg)
    fuer die AKTUELLE Config gar nicht mehr kennt. Wir rekonstruieren das Loesch-
    Kommando direkt aus 'iptables -S' (per shlex, damit z.B. das --log-prefix mit
    Leerzeichen erhalten bleibt), statt die Args selbst nachzubauen -- so passt das
    -D immer exakt zur real existierenden Regel.
    """
    name = cfg["ipset_name"]
    chain = cfg["chain"]
    needle = f"--match-set {name} src"
    removed = 0
    for line in _list_chain_rules(chain):
        if needle not in line:
            continue
        tokens = shlex.split(line)  # z.B. ["-A", "INPUT", "-m", "set", ...]
        # "-A <chain> <rest>"  ->  "iptables -D <chain> <rest>"
        run(["iptables", "-D"] + tokens[1:], dry_run=dry_run, check=False)
        removed += 1
    if removed:
        log(f"{removed} bestehende geoblock-Regel(n) aus Chain '{chain}' entfernt (Set '{name}')")
    else:
        log(f"Keine bestehenden geoblock-Regeln in Chain '{chain}' gefunden")
    return removed


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, proxy=cfg.get("proxy", "")))

    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) Vor dem Neu-Setzen ALLE bestehenden geoblock-Regeln der Chain entfernen
    #    (auch verwaiste aus frueheren Laeufen mit anderer ini, z.B. andere Ports).
    #    Erst JETZT (nach dem Befuellen des Sets), damit waehrend des ipdeny-Downloads
    #    der bestehende Block aktiv bleibt und kein Traffic-Loch entsteht.
    purge_geoblock_rules(cfg, dry_run=dry_run)

    # 5) LOG-Regel(n) + DROP-Regel(n) frisch setzen.
    #    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"]
    # Loescht ALLE geoblock-Regeln der Chain (auch verwaiste aus frueheren Laeufen
    # mit anderer ini), nicht nur die aktuell konfigurierten Varianten -- so bleiben
    # keine Reste stehen, egal wie sich 'ports'/'protocol' zwischenzeitlich geaendert haben.
    purge_geoblock_rules(cfg, dry_run=dry_run)
    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()