File size: 8,090 Bytes
5d97275
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
"""Enforcement with a conscience: every action has a TTL, an audit record and a backend that can be a dry run.

    e = Enforcer(policy, backend="auto")        # auto -> nftables | iptables | windows | dryrun (what the host has)
    e.apply(ip, action, family, score, reason)   # block / rate_limit / challenge; 'monitor' is only logged
    e.expire(now)                                # lifts what has timed out
    e.state()                                    # what is in force now (for the report and the proxy's fast path)

Backends only ever touch a dedicated table / chain / rule group named 'dibaguard', never the host's own rules, and
the layer removes what it added. 'challenge' is a mark the HTTP fast path reads (serve a JS/captcha challenge);
'rate_limit' is a mark for the proxy (token bucket per source) and, with nftables, a limit rule.
"""
from __future__ import annotations

import json
import os
import shutil
import subprocess
import time


class Backend:
    name = "dryrun"

    def __init__(self):
        self.log: list = []

    def run(self, argv: list[str]) -> bool:
        self.log.append(" ".join(argv))
        return True

    def setup(self) -> None:
        pass

    def block(self, ip: str, ttl_s: float = 1800) -> bool:
        return self.run(["dryrun", "block", ip, f"ttl={int(ttl_s)}"])

    def unblock(self, ip: str) -> bool:
        return self.run(["dryrun", "unblock", ip])


class _Shell(Backend):
    def run(self, argv: list[str]) -> bool:
        self.log.append(" ".join(argv))
        try:
            return subprocess.run(argv, capture_output=True, text=True, timeout=10).returncode == 0
        except Exception:  # noqa: BLE001
            return False


class Nftables(_Shell):
    name = "nftables"

    def setup(self) -> None:
        self.run(["nft", "add", "table", "inet", "dibaguard"])
        self.run(["nft", "add", "set", "inet", "dibaguard", "blocked", "{ type ipv4_addr; flags timeout; }"])
        self.run(["nft", "add", "set", "inet", "dibaguard", "blocked6", "{ type ipv6_addr; flags timeout; }"])
        self.run(["nft", "add", "chain", "inet", "dibaguard", "input", "{ type filter hook input priority -10; }"])
        self.run(["nft", "add", "rule", "inet", "dibaguard", "input", "ip", "saddr", "@blocked", "drop"])
        self.run(["nft", "add", "rule", "inet", "dibaguard", "input", "ip6", "saddr", "@blocked6", "drop"])

    def block(self, ip: str, ttl_s: float = 1800) -> bool:
        s = "blocked6" if ":" in ip else "blocked"
        return self.run(["nft", "add", "element", "inet", "dibaguard", s, f"{{ {ip} timeout {int(ttl_s)}s }}"])

    def unblock(self, ip: str) -> bool:
        s = "blocked6" if ":" in ip else "blocked"
        return self.run(["nft", "delete", "element", "inet", "dibaguard", s, f"{{ {ip} }}"])


class Iptables(_Shell):
    name = "iptables"

    def setup(self) -> None:
        self.run(["iptables", "-N", "DIBAGUARD"])
        self.run(["iptables", "-C", "INPUT", "-j", "DIBAGUARD"]) or self.run(["iptables", "-I", "INPUT", "-j", "DIBAGUARD"])

    def block(self, ip: str, ttl_s: float = 1800) -> bool:
        return self.run(["iptables", "-A", "DIBAGUARD", "-s", ip, "-j", "DROP"])

    def unblock(self, ip: str) -> bool:
        return self.run(["iptables", "-D", "DIBAGUARD", "-s", ip, "-j", "DROP"])


class WindowsFirewall(_Shell):
    name = "windows"

    def block(self, ip: str, ttl_s: float = 1800) -> bool:
        return self.run(["netsh", "advfirewall", "firewall", "add", "rule", f"name=dibaguard-{ip}", "dir=in", "action=block", f"remoteip={ip}"])

    def unblock(self, ip: str) -> bool:
        return self.run(["netsh", "advfirewall", "firewall", "delete", "rule", f"name=dibaguard-{ip}"])


def pick_backend(name: str = "auto") -> Backend:
    if name == "dryrun":
        return Backend()
    if name in ("auto", "nftables") and shutil.which("nft"):
        return Nftables()
    if name in ("auto", "iptables") and shutil.which("iptables"):
        return Iptables()
    if name in ("auto", "windows") and os.name == "nt" and shutil.which("netsh"):
        return WindowsFirewall()
    return Backend()


class Enforcer:
    def __init__(self, policy, backend: str = "auto", audit_path: str | None = None):
        self.policy = policy
        self.backend = pick_backend(backend if policy.enforcing else "dryrun")
        self.backend.setup()
        self.active: dict = {}            # ip -> {"action", "until", "family", "score", "reason", "since"}
        self.audit_path = audit_path
        self.audit: list = []

    def _record(self, ev: dict) -> None:
        ev = dict(ev, ts=time.time(), mode=self.policy.d["mode"], backend=self.backend.name)
        self.audit.append(ev)
        if len(self.audit) > 100_000:
            del self.audit[:50_000]
        if self.audit_path:
            try:
                with open(self.audit_path, "a", encoding="utf-8") as f:
                    f.write(json.dumps(ev, ensure_ascii=False) + "\n")
            except OSError:
                pass

    def key(self, ip: str, action: str) -> str:
        """what an action is keyed on: an IPv6 block covers the address's /64 (policy ipv6_block_prefix, default 64) -
        one host owns a /64 and rotating inside it is free; challenges and rate limits stay per address"""
        pfx = int(self.policy.d.get("ipv6_block_prefix", 64))
        if action == "block" and ":" in ip and "/" not in ip and pfx < 128:
            import ipaddress
            try:
                a = ipaddress.ip_address(ip)
                if a.version == 6 and a.ipv4_mapped is None:
                    return str(ipaddress.ip_network(f"{ip}/{pfx}", strict=False))
            except ValueError:
                pass
        return ip

    def apply(self, ip: str, action: str, family: str, score: int, reason: str = "", now: float | None = None) -> dict:
        """take (or in report mode, record) the policy's action for a verdict; stronger actions replace weaker ones"""
        now = time.time() if now is None else now
        rank = {"monitor": 0, "rate_limit": 1, "challenge": 2, "block": 3}
        if action == "monitor" or self.policy.allowed(ip):
            ev = {"ip": ip, "action": "monitor", "family": family, "score": score, "reason": reason, "taken": False}
            self._record(ev)
            return ev
        ip = self.key(ip, action)
        cur = self.active.get(ip)
        if cur and rank[cur["action"]] >= rank[action]:
            cur["until"] = max(cur["until"], now + self.policy.ttl(action))
            return dict(cur, ip=ip, taken=False, extended=True)
        taken = False
        if self.policy.enforcing and action == "block":
            taken = self.backend.block(ip, self.policy.ttl("block"))
        self.active[ip] = {"action": action, "until": now + self.policy.ttl(action), "family": family, "score": score, "reason": reason, "since": now}
        ev = {"ip": ip, "action": action, "family": family, "score": score, "reason": reason, "taken": taken, "ttl_s": self.policy.ttl(action)}
        self._record(ev)
        return ev

    def expire(self, now: float | None = None) -> list:
        now = time.time() if now is None else now
        lifted = [ip for ip, st in self.active.items() if st["until"] <= now]
        for ip in lifted:
            st = self.active.pop(ip)
            if self.policy.enforcing and st["action"] == "block":
                self.backend.unblock(ip)
            self._record({"ip": ip, "action": "lift", "family": st["family"], "score": st["score"], "reason": "ttl", "taken": self.policy.enforcing})
        return lifted

    def state(self) -> dict:
        """what the fast path needs: blocked / challenged / rate-limited sources with their deadlines"""
        out = {"block": {}, "challenge": {}, "rate_limit": {}}
        for ip, st in self.active.items():
            out[st["action"]][ip] = st["until"]
        return out

    def lookup(self, ip: str) -> str:
        st = self.active.get(ip) or self.active.get(self.key(ip, "block"))
        return st["action"] if st else "allow"