#!/usr/bin/env python3
# =====================================================================
# f2b-webctl – restricted fail2ban control for fail2ban-webui
# Path: /usr/local/sbin/f2b-webctl   (root:root, 0750)
# Config: /etc/f2b-webctl.conf
#
# Runs only as SSH "forced command" of the fail2ban-webui key, e.g.
#   restrict,from="198.51.100.10",command="sudo -n /usr/local/sbin/f2b-webctl" ssh-ed25519 AAAA…
# The requested command comes from $SSH_ORIGINAL_COMMAND
# (when started by hand as root: from the arguments):
#
#   list                 all jails with banned IPs as JSON
#   ban <ip>             ban an IP permanently in the manual jail
#   unban <jail> <ip>    unban an IP in one jail
#   unban-all <ip>       unban an IP in all jails
#   version              agent version
#
# Everything else is rejected. Every action goes to syslog (tag f2b-webctl).
# Errors carry "error_key"/"error_args" so the web UI can translate them.
# =====================================================================
import ipaddress
import json
import os
import re
import shlex
import socket
import subprocess
import sys
import syslog

VERSION = "1.0.0"
CONF_FILE = "/etc/f2b-webctl.conf"
JAIL_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]{0,40}$")
BAN_RE = re.compile(
    r"^(\S+)\s+(\d{4}-\d\d-\d\d \d\d:\d\d:\d\d) \+ (-?\d+) = (\d{4}-\d\d-\d\d \d\d:\d\d:\d\d)")

# Never banned, regardless of the configuration: loopback, private, link-local,
# CGNAT, multicast and reserved ranges
BUILTIN_PROTECTED = (
    "0.0.0.0/8", "10.0.0.0/8", "100.64.0.0/10", "127.0.0.0/8", "169.254.0.0/16",
    "172.16.0.0/12", "192.168.0.0/16", "224.0.0.0/4", "240.0.0.0/4",
    "::/128", "::1/128", "fc00::/7", "fe80::/10", "ff00::/8",
)


# ---------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------
def out(obj, code=0):
    print(json.dumps(obj, ensure_ascii=False))
    sys.exit(code)


def fail(key, text, code=1, **args):
    out({"ok": False, "error_key": key, "error_args": args, "error": text}, code)


def client():
    return os.environ.get("SSH_CLIENT", "local").split(" ")[0]


def log(msg):
    syslog.openlog("f2b-webctl", 0, syslog.LOG_AUTH)
    syslog.syslog(syslog.LOG_NOTICE, f"{msg} (from {client()})")


def load_conf():
    conf = {"MANUAL_JAIL": "manual-block", "PROTECTED": "",
            "PROTECT_OWN_ADDRESSES": "yes", "PROTECT_CLIENT": "yes"}
    try:
        with open(CONF_FILE, encoding="utf-8") as fh:
            for line in fh:
                line = line.strip()
                if line and not line.startswith("#") and "=" in line:
                    key, value = line.split("=", 1)
                    value = " ".join(shlex.split(value, comments=True)) if value.strip() else ""
                    conf[key.strip().upper()] = value
    except FileNotFoundError:
        pass
    return conf


def own_addresses():
    addrs = set()
    try:
        r = subprocess.run(["ip", "-o", "addr", "show"], capture_output=True, text=True, timeout=5)
        addrs.update(m.split("/")[0] for m in re.findall(r"\binet6? (\S+)", r.stdout))
    except (OSError, subprocess.SubprocessError):
        pass
    try:
        for info in socket.getaddrinfo(socket.getfqdn(), None):
            addrs.add(info[4][0])
    except OSError:
        pass
    return addrs


def protected_networks(conf):
    nets = [ipaddress.ip_network(n) for n in BUILTIN_PROTECTED]
    for item in re.split(r"[\s,]+", conf["PROTECTED"]):
        if item:
            try:
                nets.append(ipaddress.ip_network(item, strict=False))
            except ValueError:
                log(f"invalid entry in PROTECTED ignored: {item!r}")
    extra = set()
    if conf["PROTECT_OWN_ADDRESSES"].lower() in ("yes", "true", "1"):
        extra |= own_addresses()
    if conf["PROTECT_CLIENT"].lower() in ("yes", "true", "1") and client() != "local":
        extra.add(client())
    for a in extra:
        try:
            nets.append(ipaddress.ip_network(a.split("%")[0]))
        except ValueError:
            pass
    return nets


def f2b(*args):
    return subprocess.run(["fail2ban-client", *args], capture_output=True, text=True, timeout=30)


def jails():
    r = f2b("status")
    if r.returncode != 0:
        detail = (r.stderr or r.stdout).strip()[:200]
        fail("agent_f2b_down", f"fail2ban is not reachable: {detail}", 1, detail=detail)
    m = re.search(r"Jail list:\s*(.*)", r.stdout)
    return [j.strip() for j in m.group(1).split(",") if j.strip()] if m else []


def valid_ip(value):
    try:
        return str(ipaddress.ip_address(value))
    except ValueError:
        fail("agent_invalid_ip", f"Invalid IP address: {value!r}", 2, value=value)


def result(r):
    return (r.stdout or r.stderr).strip()[:300]


# ---------------------------------------------------------------------
# Commands
# ---------------------------------------------------------------------
def cmd_list(conf):
    data = {"host": socket.gethostname(), "agent_version": VERSION,
            "manual_jail": conf["MANUAL_JAIL"], "jails": {}}
    with_time = True
    for jail in jails():
        bans = []
        if with_time:
            r = f2b("get", jail, "banip", "--with-time")
            if r.returncode == 0:
                for line in r.stdout.splitlines():
                    m = BAN_RE.match(line.strip())
                    if m:
                        bans.append({"ip": m[1], "start": m[2], "bantime": int(m[3]), "end": m[4]})
            else:
                with_time = False   # older fail2ban without --with-time
        if not with_time:
            r = f2b("get", jail, "banip")
            for ip in re.split(r"[\s,]+", r.stdout.strip()):
                try:
                    ipaddress.ip_address(ip)
                except ValueError:
                    continue
                bans.append({"ip": ip, "start": "", "bantime": None, "end": ""})
        total = re.search(r"Total banned:\s*(\d+)", f2b("status", jail).stdout)
        data["jails"][jail] = {"bans": bans, "total": int(total[1]) if total else None}
    out({"ok": True, **data})


def cmd_ban(conf, ip):
    addr = ipaddress.ip_address(ip)
    if any(addr in net for net in protected_networks(conf)):
        log(f"REJECTED ban {ip} (protected address)")
        fail("agent_protected", f"{ip} is protected (private network, own address or PROTECTED list).", 3, ip=ip)
    jail = conf["MANUAL_JAIL"]
    if jail not in jails():
        fail("agent_no_manual_jail", f"Jail {jail!r} does not exist on this host.", 2, jail=jail)
    r = f2b("set", jail, "banip", ip)
    ok = r.returncode == 0
    log(f"ban {ip} -> {jail}: {'OK' if ok else result(r)}")
    out({"ok": ok, "jail": jail, "ip": ip, "result": result(r)}, 0 if ok else 1)


def cmd_unban(jail, ip):
    if jail not in jails():
        fail("agent_unknown_jail", f"Unknown jail: {jail!r}", 2, jail=jail)
    r = f2b("set", jail, "unbanip", ip)
    ok = r.returncode == 0
    log(f"unban {ip} from {jail}: {'OK' if ok else result(r)}")
    out({"ok": ok, "jail": jail, "ip": ip, "result": result(r)}, 0 if ok else 1)


def cmd_unban_all(ip):
    r = f2b("unban", ip)
    ok = r.returncode == 0
    log(f"unban-all {ip}: {'OK' if ok else result(r)}")
    out({"ok": ok, "jail": "*", "ip": ip, "result": result(r)}, 0 if ok else 1)


def main():
    if os.geteuid() != 0:
        fail("agent_not_root", "f2b-webctl must run as root (check the sudo rule).", 4)
    original = os.environ.get("SSH_ORIGINAL_COMMAND")
    args = original.split() if original is not None else sys.argv[1:]
    conf = load_conf()
    if args == ["list"]:
        cmd_list(conf)
    elif args == ["version"]:
        out({"ok": True, "agent_version": VERSION})
    elif len(args) == 2 and args[0] == "ban":
        cmd_ban(conf, valid_ip(args[1]))
    elif len(args) == 3 and args[0] == "unban" and JAIL_RE.fullmatch(args[1]):
        cmd_unban(args[1], valid_ip(args[2]))
    elif len(args) == 2 and args[0] == "unban-all":
        cmd_unban_all(valid_ip(args[1]))
    else:
        log(f"REJECTED unknown command: {' '.join(args)[:100]!r}")
        fail("agent_not_allowed", "Command not allowed.", 2)


if __name__ == "__main__":
    main()