| """High availability: state that survives restarts, moves between hosts, and stays coherent across replicas. |
| |
| Snapshots (restart, host move, active-passive shipping) |
| info = snapshot(layer, "/var/lib/dibaguard/snap") # atomic (temp file + rename), keeps the last K, returns size / timings |
| rep = restore(layer, "/var/lib/dibaguard/snap") # newest valid snapshot; refuses one made for other model files |
| snap = Snapshotter(layer, dir, every_s=300, lock=console.lock); snap.restore_on_start(); snap.maybe() # sidecar loop |
| |
| One .npz per snapshot, always read with allow_pickle=False - no member of a snapshot can execute code on load: |
| manifest uint8 JSON: magic, schema version, created, replica, model fingerprint (sha256 over every model array and |
| info, per component: tier1 / seq / http / dbnet), code-layout id (feature layout, horizons, age edges), |
| parts, stats, and the digest - sha256 (HMAC-SHA256 when DIBAGUARD_SNAPSHOT_KEY is set) over the manifest |
| itself and every other member (name, dtype, shape, bytes) |
| state uint8 JSON: the structure; every numeric table it refers to is a member a0, a1, ... |
| keys uint8 JSON: the shared key table - every per-source key (address, host, (host, port), ...) is stored once |
| and referred to by index from int64 members (never an object array) |
| a<i> numeric arrays (float64 / float32 / int64 / int8): references, per-source tables, GRU states, buffers |
| The capture holds the engine lock only while live structures are copied (~0.7 s for 100k sources); key encoding, |
| compression, hashing and the write run after it (docs/HA.md has the measured numbers). |
| Parts - model-dependent: guard (calibration and rolling references, Mondrian cells, anchor, re-calibration buffer, |
| per-source / per-destination evidence, entity risk, stream counters, GRU states of warm sources, per-host HTTP |
| profiles and references, drift, pressure, global rates, alarm log), layer (per-host HTTP bootstrap / re-profile |
| buffers, payload-provenance ring, playbook counts), db (known clients and query shapes, volume references, DbNet |
| reference, decayed counters, MySQL thread map). Model-independent: enforcement (absolute expiry times), feedback |
| (verdicts, allowlist TTLs), rules, allowlist - the only parts a forced restore across different model files brings back. |
| |
| Replicas (active-active behind a load balancer) |
| shared = FileSharedState("/mnt/shared/dibaguard", replica="lb-a") # or RedisSharedState(url, replica="lb-a") |
| sync = ReplicaSync(layer, shared, lock=console.lock); sync.maybe() # every few seconds |
| affinity("203.0.113.7", ["lb-a", "lb-b", "lb-c"]) # the replica that owns this source's /24 |
| (a) shared enforcement view: each replica publishes its block / challenge / rate-limit decisions (absolute expiry), |
| operator lifts and allowlist entries; each imports the merged view - strongest action, latest expiry, and a lift |
| cancels every decision published before it - so every proxy's fast path enforces what any replica decided. |
| (b) source affinity: a consistent (rendezvous) hash of the client's /24 (/64 for IPv6) picks its replica, so a |
| source's evidence accumulates in one place; see docs/HA.md for the HAProxy / nginx equivalents. |
| (c) evidence federation: replicas exchange their top-K sources by accumulated evidence; a source whose evidence is |
| split across replicas is judged on its total against the same conformal threshold the engine uses. |
| """ |
| from __future__ import annotations |
|
|
| import contextlib |
| import gc |
| import hashlib |
| import heapq |
| import hmac |
| import ipaddress |
| import itertools |
| import json |
| import math |
| import os |
| import re |
| import socket |
| import threading |
| import time |
| import zipfile |
| from collections import Counter, OrderedDict, deque |
|
|
| import numpy as np |
|
|
| from dibaguard.decide import ConformalFPR, RollingReference, event_score |
| from dibaguard.families import severity |
| from dibaguard.flows import entity as ent |
|
|
| SCHEMA = 1 |
| MAGIC = "dibaguard-snapshot" |
| MODEL_PARTS = ("guard", "layer", "db") |
| MODEL_INDEPENDENT = ("enforcement", "feedback", "rules", "allowlist") |
| ALL_PARTS = MODEL_PARTS + MODEL_INDEPENDENT |
| RANK = {"monitor": 0, "rate_limit": 1, "challenge": 2, "block": 3} |
| _SNAP = re.compile(r"^dibaguard-snap-(\d{8}T\d{9})Z-([A-Za-z0-9_.-]+)\.npz$") |
| _TMP = re.compile(r"^\.dibaguard-snap-.*\.tmp$") |
| _NAMED = lambda f: f not in ("", "benign", "attack_generic", None) |
|
|
|
|
| class SnapshotError(Exception): |
| """unreadable, corrupted, tampered or incompatible snapshot - nothing was restored""" |
|
|
|
|
| class FingerprintMismatch(SnapshotError): |
| """the snapshot was made for other model files (or another feature layout)""" |
|
|
|
|
| |
| def _jdefault(o): |
| if isinstance(o, np.generic): |
| return o.item() |
| if isinstance(o, np.ndarray): |
| return o.tolist() |
| if isinstance(o, (set, frozenset)): |
| return sorted(o, key=str) |
| if isinstance(o, (bytes, bytearray)): |
| return o.decode("utf-8", "replace") |
| if isinstance(o, deque): |
| return list(o) |
| return str(o) |
|
|
|
|
| def _canon(o) -> bytes: |
| """canonical JSON (sorted keys, ASCII - lone surrogates and Persian text alike are escaped)""" |
| return json.dumps(o, sort_keys=True, separators=(",", ":"), default=_jdefault, ensure_ascii=True).encode("ascii") |
|
|
|
|
| def _jkey(k): |
| """a dict key -> JSON (tuples tagged so they come back as tuples)""" |
| if isinstance(k, tuple): |
| return {"t": [_jkey(x) for x in k]} |
| if isinstance(k, np.generic): |
| return k.item() |
| if k is None or isinstance(k, (str, int, float, bool)): |
| return k |
| return str(k) |
|
|
|
|
| def _ukey(j): |
| return tuple(_ukey(x) for x in j["t"]) if isinstance(j, dict) else j |
|
|
|
|
| def _safe_name(s) -> str: |
| return re.sub(r"[^A-Za-z0-9_.-]", "_", str(s))[:64] or "replica" |
|
|
|
|
| def default_replica() -> str: |
| return _safe_name(os.environ.get("DIBAGUARD_REPLICA") or socket.gethostname()) |
|
|
|
|
| def _key(key): |
| """snapshot MAC key: argument, else DIBAGUARD_SNAPSHOT_KEY, else None (plain sha256)""" |
| if key is None: |
| key = os.environ.get("DIBAGUARD_SNAPSHOT_KEY") or None |
| if isinstance(key, str): |
| key = key.encode("utf-8") |
| return key or None |
|
|
|
|
| def _lock(lock): |
| return lock if lock is not None else contextlib.nullcontext() |
|
|
|
|
| def _replace(tmp: str, final: str, tries: int = 8) -> None: |
| """os.replace with a few retries: on Windows a reader holding the target open makes it fail for a moment""" |
| for i in range(tries): |
| try: |
| os.replace(tmp, final) |
| return |
| except PermissionError: |
| if i == tries - 1: |
| raise |
| time.sleep(0.01 * (i + 1)) |
|
|
|
|
| def _fsync_dir(d: str) -> None: |
| if os.name == "nt": |
| return |
| try: |
| fd = os.open(d, os.O_RDONLY) |
| try: |
| os.fsync(fd) |
| finally: |
| os.close(fd) |
| except OSError: |
| pass |
|
|
|
|
| def _atomic_bytes(path: str, data: bytes, durable: bool = True) -> None: |
| tmp = f"{path}.{os.getpid()}.{threading.get_ident()}.tmp" |
| try: |
| with open(tmp, "wb") as f: |
| f.write(data) |
| f.flush() |
| if durable: |
| os.fsync(f.fileno()) |
| _replace(tmp, path) |
| except BaseException: |
| with contextlib.suppress(OSError): |
| os.remove(tmp) |
| raise |
|
|
|
|
| |
| def _hash_arrays(h, *arrays) -> None: |
| for a in arrays: |
| a = np.asarray(a) |
| h.update(f"|{a.dtype.str}{a.shape}|".encode()) |
| if a.size: |
| h.update(np.ascontiguousarray(a).reshape(-1).view(np.uint8)) |
|
|
|
|
| _FP_CACHE: "OrderedDict[tuple, tuple]" = OrderedDict() |
|
|
|
|
| def model_fingerprint(layer) -> dict: |
| """sha256 per model component (tier1 trees, sequence net, HTTP model incl. its shipped site profile, DbNet) and |
| an id over all of them. Models are immutable at runtime, so the result is cached per model objects.""" |
| g = getattr(layer, "guard", layer) |
| net = getattr(getattr(layer, "dbg", None), "net", None) |
| objs = (g.tier1, g.seq, g.http, net) |
| ck = tuple(id(o) for o in objs) |
| hit = _FP_CACHE.get(ck) |
| if hit is not None and all(a is b for a, b in zip(hit[0], objs)): |
| return dict(hit[1]) |
| comp: dict = {} |
| t1 = g.tier1 |
| h = hashlib.sha256() |
| _hash_arrays(h, t1.feat, t1.border, t1.tree_tests, t1.leaves, t1.bias) |
| h.update(_canon({"names": t1.names, "classes": t1.classes, "meta": t1.meta})) |
| comp["tier1"] = h.hexdigest() |
| comp["seq"] = None |
| if g.seq is not None: |
| h = hashlib.sha256() |
| for k in sorted(g.seq.w): |
| h.update(k.encode()) |
| _hash_arrays(h, g.seq.w[k]) |
| h.update(_canon(g.seq.info)) |
| comp["seq"] = h.hexdigest() |
| comp["http"] = None |
| if g.http is not None: |
| m, h = g.http, hashlib.sha256() |
| for n in ("emb_q", "scale", "W1", "b1", "Wc", "bc", "W2", "b2"): |
| v = getattr(m, n, None) |
| if v is not None: |
| h.update(n.encode()) |
| _hash_arrays(h, v) |
| h.update(_canon(m.info)) |
| site = getattr(m, "site", None) or {} |
| for k in sorted(site): |
| h.update(str(k).encode()) |
| if isinstance(site[k], np.ndarray): |
| _hash_arrays(h, site[k]) |
| else: |
| h.update(_canon(site[k])) |
| comp["http"] = h.hexdigest() |
| comp["dbnet"] = None |
| if net is not None: |
| h = hashlib.sha256() |
| for n in ("emb", "scale", "w1", "b1", "w2", "b2"): |
| _hash_arrays(h, getattr(net, n)) |
| h.update(_canon({"buckets": net.buckets, "thr": net.thr})) |
| comp["dbnet"] = h.hexdigest() |
| comp["id"] = hashlib.sha256(_canon(comp)).hexdigest() |
| _FP_CACHE[ck] = (objs, dict(comp)) |
| while len(_FP_CACHE) > 4: |
| _FP_CACHE.popitem(last=False) |
| return comp |
|
|
|
|
| def code_layout() -> dict: |
| """what the stored per-source counters and reference columns mean; a different layout cannot be restored""" |
| from dibaguard import engine |
| lay = {"features": ent.feature_names(), "horizons": list(ent.HORIZONS_S), "memory": ent.MEMORY, |
| "age_edges": list(engine.AGE_EDGES), "anom_idx": list(engine.ANOM_IDX), "dst_idx": list(engine.DST_IDX)} |
| lay["id"] = hashlib.sha256(_canon(lay)).hexdigest() |
| return lay |
|
|
|
|
| |
| _STRS = frozenset({str, np.str_}) |
|
|
|
|
| class _W: |
| """numeric members + one shared key table (every per-source key is stored once, as an index)""" |
|
|
| def __init__(self): |
| self.arrays: dict = {} |
| self.keys: list = [] |
| self._s: dict = {} |
| self._i: dict = {} |
| self._kix: dict = {} |
| self._deferred: list = [] |
|
|
| def a(self, x, dtype=None) -> str: |
| x = np.asarray(x) if dtype is None else np.asarray(x, dtype) |
| if x.dtype.hasobject: |
| raise SnapshotError("refusing to write an object array (it would need pickle)") |
| name = f"a{len(self.arrays)}" |
| self.arrays[name] = x |
| return name |
|
|
| def k(self, key) -> int: |
| c = key.__class__ |
| d = self._s if c in _STRS else self._i if c is int else None |
| if d is not None: |
| i = d.get(key) |
| if i is None: |
| i = d[key] = len(self.keys) |
| self.keys.append(str(key) if c is np.str_ else key) |
| return i |
| tag = "b" if isinstance(key, (bool, np.bool_)) else "f" if isinstance(key, (float, np.floating)) else "" |
| i = self._kix.get((tag, key)) |
| if i is None: |
| i = self._kix[(tag, key)] = len(self.keys) |
| self.keys.append(_jkey(key)) |
| return i |
|
|
| def ks(self, keys) -> str: |
| k = self.k |
| return self.a(np.fromiter((k(x) for x in keys), np.int64)) |
|
|
| def reserve(self) -> str: |
| name = f"a{len(self.arrays)}" |
| self.arrays[name] = None |
| return name |
|
|
| def defer(self, arr) -> str: |
| """an object column of an array the engine never mutates (it replaces it): encoded later, outside the lock""" |
| name = self.reserve() |
| self._deferred.append((name, lambda arr=arr: np.fromiter((self.k(x) for x in arr.tolist()), np.int64, len(arr)))) |
| return name |
|
|
| def later(self, fn) -> str: |
| """a member computed outside the lock from data already copied under it""" |
| name = self.reserve() |
| self._deferred.append((name, fn)) |
| return name |
|
|
| def finish(self) -> dict: |
| for name, fn in self._deferred: |
| self.arrays[name] = np.asarray(fn()) |
| if self.arrays[name].dtype.hasobject: |
| raise SnapshotError("refusing to write an object array (it would need pickle)") |
| self._deferred = [] |
| self.arrays["keys"] = np.frombuffer(_canon(self.keys), np.uint8) |
| return self.arrays |
|
|
|
|
| class _R: |
| def __init__(self, arrays: dict, keys: list): |
| self.arrays = arrays |
| self.keys = [_ukey(k) for k in keys] |
|
|
| def a(self, name): |
| if name not in self.arrays: |
| raise SnapshotError(f"snapshot member {name!r} is missing") |
| return self.arrays[name] |
|
|
| def ks(self, name) -> list: |
| K = self.keys |
| return [K[i] for i in self.a(name).tolist()] |
|
|
|
|
| def _ref_out(w: _W, r): |
| if r is None: |
| return None |
| return {"window_s": r.window_s, "keep": r.keep, "max_rows": r.max_rows, "cur_k": r.cur_k, |
| "base": None if r.base is None else w.a(r.base), "done": [w.a(x) for x in r.done], |
| "cur": w.a(np.asarray(r.cur, np.float64)) if r.cur else None} |
|
|
|
|
| def _ref_in(rd: _R, d): |
| if d is None: |
| return None |
| r = RollingReference(float(d["window_s"]), int(d["keep"]), int(d["max_rows"])) |
| r.base = None if d["base"] is None else rd.a(d["base"]) |
| r.done = [rd.a(x) for x in d["done"]] |
| if d["cur"] is not None: |
| c = rd.a(d["cur"]) |
| r.cur = c.tolist() if c.ndim == 1 else list(c) |
| r.cur_k = d["cur_k"] |
| r._merged = None |
| return r |
|
|
|
|
| def _od_out(w: _W, od, width: int, keep=None, last=None) -> dict: |
| """OrderedDict key -> [numbers] (LRU order kept); keep: only these keys; last: only the most recent `last`""" |
| if keep is not None: |
| items = [(k, v) for k, v in od.items() if k in keep] |
| elif last is not None and len(od) > last: |
| items = list(itertools.islice(reversed(od.items()), last))[::-1] |
| else: |
| items = list(od.items()) |
| try: |
| v = np.array([v for _, v in items], np.float64).reshape(len(items), width) |
| except ValueError: |
| v = np.array([(list(v) + [0.0] * width)[:width] for _, v in items], np.float64).reshape(len(items), width) |
| keys = [k for k, _ in items] |
| return {"k": w.later(lambda: np.fromiter((w.k(x) for x in keys), np.int64, len(keys))), "v": w.a(v)} |
|
|
|
|
| def _od_in(rd: _R, d: dict, ints=(), maxlen=None) -> OrderedDict: |
| ks, vs = rd.ks(d["k"]), rd.a(d["v"]).tolist() |
| if maxlen is not None and len(ks) > maxlen: |
| ks, vs = ks[-maxlen:], vs[-maxlen:] |
| od: OrderedDict = OrderedDict() |
| for k, v in zip(ks, vs): |
| for j in ints: |
| v[j] = int(v[j]) |
| od[k] = v |
| return od |
|
|
|
|
| def _ent_out(w: _W, table, keep=None) -> dict: |
| """EntityState records: time, decayed counters, gap EWMA, true age and the remembered (host, port) / source keys in |
| LRU order; the per-host counts are rebuilt from those on restore (they are always Counter(seen.values())). |
| Under the lock only the mutable parts are copied (counters, the seen list); arrays and keys are built later.""" |
| recs = [(k, r.t, r.c[:], r.gap_ewma, r.n, list(r.seen.items())) for k, r in table.items() if keep is None or k in keep] |
| width = len(recs[0][2]) if recs else 0 |
| built: dict = {} |
|
|
| def build() -> dict: |
| if built: |
| return built |
| K = w.k |
| sa, sb, sh, off = [], [], [], [0] |
| a_sa, a_sb, a_sh = sa.append, sb.append, sh.append |
| for rec in recs: |
| for sk, host in rec[5]: |
| if sk.__class__ is tuple and len(sk) == 2: |
| a_sa(K(sk[0])) |
| a_sb(K(sk[1])) |
| else: |
| a_sa(K(sk)) |
| a_sb(-1) |
| a_sh(K(host)) |
| off.append(len(sa)) |
| n = len(recs) |
| built.update(k=np.fromiter((K(r[0]) for r in recs), np.int64, n), t=np.fromiter((r[1] for r in recs), np.float64, n), |
| c=np.fromiter(itertools.chain.from_iterable(r[2] for r in recs), np.float64, n * width).reshape(n, width), |
| gap=np.fromiter((r[3] for r in recs), np.float64, n), n=np.fromiter((r[4] for r in recs), np.int64, n), |
| off=np.asarray(off, np.int64), sa=np.asarray(sa, np.int64), sb=np.asarray(sb, np.int64), sh=np.asarray(sh, np.int64)) |
| return built |
| return {x: w.later(lambda x=x: build()[x]) for x in ("k", "t", "c", "gap", "n", "off", "sa", "sb", "sh")} |
|
|
|
|
| def _ent_in(rd: _R, d: dict, maxlen=None) -> OrderedDict: |
| K = rd.keys |
| keys, ts, C, gaps, ns = (rd.a(d[x]).tolist() for x in ("k", "t", "c", "gap", "n")) |
| off, sa, sb, sh = (rd.a(d[x]).tolist() for x in ("off", "sa", "sb", "sh")) |
| start = max(0, len(keys) - maxlen) if maxlen is not None else 0 |
| table: OrderedDict = OrderedDict() |
| for i in range(start, len(keys)): |
| r = ent._Rec(ts[i], 0) |
| r.c, r.gap_ewma, r.n = C[i], gaps[i], int(ns[i]) |
| seen: OrderedDict = OrderedDict() |
| hosts: dict = {} |
| for j in range(off[i], off[i + 1]): |
| key = K[sa[j]] if sb[j] < 0 else (K[sa[j]], K[sb[j]]) |
| h = K[sh[j]] |
| seen[key] = h |
| hosts[h] = hosts.get(h, 0) + 1 |
| r.seen, r.hosts = seen, hosts |
| table[K[keys[i]]] = r |
| return table |
|
|
|
|
| def _buf_out(w: _W, b, rows=None): |
| """the re-calibration buffer / the anchor: column arrays; object columns (src, dst) as key indices""" |
| if b is None: |
| return None |
| n = len(b["t"]) if "t" in b else None |
| out = [] |
| for k, v in b.items(): |
| if isinstance(v, np.ndarray): |
| if rows is not None and n is not None and n > rows and v.shape[:1] == (n,): |
| v = v[n - rows:] |
| out.append([k, {"obj": w.defer(v)} if v.dtype.hasobject else {"arr": w.a(v)}]) |
| else: |
| out.append([k, {"json": v}]) |
| return out |
|
|
|
|
| def _buf_in(rd: _R, d): |
| if d is None: |
| return None |
| b = {} |
| for k, v in d: |
| if "obj" in v: |
| ks = rd.ks(v["obj"]) |
| b[k] = np.fromiter(ks, dtype=object, count=len(ks)) |
| elif "arr" in v: |
| b[k] = rd.a(v["arr"]) |
| else: |
| b[k] = v["json"] |
| return b |
|
|
|
|
| def _prof_out(w: _W, p: dict) -> list: |
| return [[k, {"arr": w.a(v)} if isinstance(v, np.ndarray) else {"json": v}] for k, v in p.items()] |
|
|
|
|
| def _prof_in(rd: _R, d: list) -> dict: |
| return {k: (rd.a(v["arr"]) if "arr" in v else v["json"]) for k, v in d} |
|
|
|
|
| |
| |
| |
| |
| |
| _GUARD_KNOWN = {"state", "site", "rolling", "anom", "anom_s", "anom_age", "seq_age", "evidence", "ev_ref", "evidence_slow", "ev_slow_ref", |
| "evid_drift", "dst_anom", "dst_anom_s", "dst_evidence", "dst_drift", "dst_rate_hi", "dst_ev_ref", "_anchor", "global_http", |
| "global_scan", "_buffer", "_pending", "last_recal", "roll_groups", "entities", "pressure", "drift", "gru", "site_http", "site_http_by_host", |
| "http_ref_by_host", "log", "counts", "warmup_short", |
| "tier1", "seq", "http", "alpha", "t1_idx", "temp1", "temp2", "temph", "shipped", "mondrian_min", "warmup_s", "recal_s", |
| "recal_window_s", "recal_max_rows", "warm_flows", "warm_logit", "gamma", "beta", "global_floor", "global_scan_floor", "_fam"} |
| _LAYER_KNOWN = {"guard", "policy", "position", "enforcer", "feedback", "rules", "llm_interpret", "posture_findings", "playbook_hits", |
| "http_boot_n", "http_reprofile_s", "_http_last", "_http_boot_by", "_http_recent_by", "_http_last_by", "dbg", "_http_ring"} |
| _DB_KNOWN = {"net", "alpha", "tau", "auth_fail_floor", "learn_new_clients", "known_clients", "known_shapes", "vol_ref", "net_ref", "vol", |
| "fails", "exempt_shapes", "log", "counts"} |
| _CLASSES: dict = {} |
|
|
|
|
| def _allowed_classes() -> dict: |
| """the only classes a snapshot may instantiate on restore - looked up by qualified name in this fixed table, built |
| with cls.__new__ and their attributes set (no constructor, no import by name, nothing else can run)""" |
| if not _CLASSES: |
| from dibaguard import campaign, decide |
| from dibaguard.shield import database |
| for mod, names in ((decide, ("ConformalFPR", "DriftMartingale", "EntityRisk", "SitePressure", "RollingReference", "SourceEvidence", |
| "GlobalRate", "Temperature")), |
| (campaign, ("CampaignTracker", "_Key")), (database, ("_Decay",)), (ent, ("_Rec", "EntityState"))): |
| for n in names: |
| c = getattr(mod, n, None) |
| if isinstance(c, type): |
| _CLASSES[f"{c.__module__}.{c.__qualname__}"] = c |
| return _CLASSES |
|
|
|
|
| class _Skip(Exception): |
| pass |
|
|
|
|
| def _gen_out(w: _W, o, depth: int = 0): |
| """plain data -> tagged JSON: lists as lists; {"T"} tuple, {"S"}/{"FS"} set (sorted), {"D"}/{"OD"} dict pairs (order |
| kept), {"DQ", "m"} deque, {"A"} array member, {"NT"} a numeric table (key -> fixed-length row of numbers, the fast |
| path for per-source counters), {"O", "v"} an instance of an allowlisted class""" |
| if depth > 16: |
| raise _Skip("nested too deep") |
| if o is None or isinstance(o, (bool, int, float, str)): |
| return o.item() if isinstance(o, np.generic) else o |
| if isinstance(o, np.generic): |
| return o.item() |
| if isinstance(o, np.ndarray): |
| return {"OA": [_gen_out(w, x, depth + 1) for x in o.tolist()]} if o.dtype.hasobject else {"A": w.a(np.array(o, copy=True))} |
| if isinstance(o, tuple): |
| return {"T": [_gen_out(w, x, depth + 1) for x in o]} |
| if isinstance(o, list): |
| return [_gen_out(w, x, depth + 1) for x in o] |
| if isinstance(o, (set, frozenset)): |
| return {"FS" if isinstance(o, frozenset) else "S": sorted((_gen_out(w, x, depth + 1) for x in o), key=_canon)} |
| if isinstance(o, deque): |
| return {"DQ": [_gen_out(w, x, depth + 1) for x in o], "m": o.maxlen} |
| if isinstance(o, dict): |
| od = isinstance(o, OrderedDict) |
| vals = list(o.values()) |
| if len(vals) >= 64 and vals[0].__class__ is list: |
| width = len(vals[0]) |
| if all(v.__class__ is list and len(v) == width and all(x.__class__ in (int, float) for x in v) for v in vals): |
| ints = [j for j in range(width) if all(v[j].__class__ is int for v in vals)] |
| return {"NT": _od_out(w, o, width), "ints": ints, "od": od} |
| return {"OD" if od else "D": [[_gen_out(w, k, depth + 1), _gen_out(w, v, depth + 1)] for k, v in o.items()]} |
| qn = f"{type(o).__module__}.{type(o).__qualname__}" |
| if qn in _allowed_classes(): |
| names = list(vars(o)) if hasattr(o, "__dict__") else [s for s in getattr(type(o), "__slots__", ()) if hasattr(o, s)] |
| return {"O": qn, "v": [[n, _gen_out(w, getattr(o, n), depth + 1)] for n in names]} |
| raise _Skip(qn) |
|
|
|
|
| def _gen_in(rd: _R, j): |
| if j is None or isinstance(j, (bool, int, float, str)): |
| return j |
| if isinstance(j, list): |
| return [_gen_in(rd, x) for x in j] |
| if "A" in j: |
| return rd.a(j["A"]) |
| if "OA" in j: |
| xs = [_gen_in(rd, x) for x in j["OA"]] |
| return np.fromiter(xs, dtype=object, count=len(xs)) |
| if "T" in j: |
| return tuple(_gen_in(rd, x) for x in j["T"]) |
| if "S" in j: |
| return {_gen_in(rd, x) for x in j["S"]} |
| if "FS" in j: |
| return frozenset(_gen_in(rd, x) for x in j["FS"]) |
| if "DQ" in j: |
| return deque((_gen_in(rd, x) for x in j["DQ"]), maxlen=j.get("m")) |
| if "NT" in j: |
| od = _od_in(rd, j["NT"], tuple(j.get("ints", ()))) |
| return od if j.get("od") else dict(od) |
| if "OD" in j or "D" in j: |
| pairs = ((_gen_in(rd, k), _gen_in(rd, v)) for k, v in j.get("OD", j.get("D"))) |
| return OrderedDict(pairs) if "OD" in j else dict(pairs) |
| if "O" in j: |
| cls = _allowed_classes().get(j["O"]) |
| if cls is None: |
| raise SnapshotError(f"class {j['O']!r} is not allowed in a snapshot") |
| obj = cls.__new__(cls) |
| for n, v in j["v"]: |
| setattr(obj, n, _gen_in(rd, v)) |
| return obj |
| raise SnapshotError(f"unknown generic tag {sorted(j)[:3]}") |
|
|
|
|
| def _extras_out(w: _W, obj, known: set, stats: dict, where: str, max_entries: int) -> list: |
| """every non-scalar attribute of obj not handled explicitly (scalars are configuration); per-key tables over |
| `max_entries` keep their most recent entries""" |
| out = [] |
| for name, val in list(vars(obj).items()): |
| if name in known or val is None or isinstance(val, (bool, int, float, str, np.generic)) or callable(val): |
| continue |
| if isinstance(val, OrderedDict) and len(val) > max_entries: |
| val = OrderedDict(itertools.islice(val.items(), len(val) - max_entries, None)) |
| try: |
| out.append([name, _gen_out(w, val)]) |
| except _Skip as e: |
| stats.setdefault("not_snapshotted", []).append(f"{where}.{name} ({e})") |
| return out |
|
|
|
|
| def _extras_in(rd: _R, pairs: list, obj) -> list: |
| """only attributes the running code still has are set back""" |
| return [(obj, name, _gen_in(rd, v)) for name, v in pairs if hasattr(obj, name)] |
|
|
|
|
| |
| def _latest_t(g) -> float: |
| ts = [] |
| for od, get in ((g.state.src, lambda r: r.t), (g.entities.s, lambda v: v[1]), (g.evidence.e, lambda v: v[1])): |
| with contextlib.suppress(StopIteration): |
| ts.append(float(get(od[next(reversed(od))]))) |
| return max(ts) if ts else 0.0 |
|
|
|
|
| def _choose(g, L, max_sources: int, max_destinations: int): |
| """None (everything fits) or the set of keys to keep: every source in force, plus at most `max_sources` others - |
| up to half of them the riskiest (decayed risk, fast evidence, slow evidence, warm GRU sources; an eighth each), |
| the rest the most recently seen""" |
| t_now = _latest_t(g) |
|
|
| def decayed(tau): |
| return lambda kv: kv[1][0] * math.exp(-max(0.0, t_now - kv[1][1]) / tau) |
|
|
| def pick(forced: set, budget: int, tops: list, recent: list) -> set: |
| others: set = set() |
| q = max(1, budget // 8) |
| for items, key in tops: |
| others.update(k for k, _ in (heapq.nlargest(q, items, key=key) if key else itertools.islice(items, q)) if k not in forced) |
| for od in recent: |
| for k in reversed(od): |
| if len(others) >= budget: |
| break |
| if k not in forced: |
| others.add(k) |
| return forced | others |
|
|
| keep_src = keep_dst = None |
| if max(len(g.state.src), len(g.evidence.e), len(g.evidence_slow.e), len(g.entities.s), len(g.gru)) > max_sources: |
| forced = set(L.enforcer.active) if L is not None else set() |
| keep_src = pick(forced, max_sources, [(g.entities.s.items(), decayed(g.entities.tau)), (g.evidence.e.items(), decayed(g.evidence.tau)), |
| (g.evidence_slow.e.items(), decayed(g.evidence_slow.tau)), (((k, None) for k in reversed(g.gru)), None)], |
| [g.state.src, g.entities.s, g.evidence.e]) |
| if max(len(g.state.dst), len(g.dst_evidence.e)) > max_destinations: |
| keep_dst = pick(set(), max_destinations, [(g.dst_evidence.e.items(), decayed(g.dst_evidence.tau))], [g.state.dst, g.dst_evidence.e]) |
| return keep_src, keep_dst |
|
|
|
|
| |
| def _guard_out(w: _W, g, L, stats: dict, max_sources, max_destinations, buffer_rows, max_log, max_global) -> dict: |
| if hasattr(g, "_flush_buffer"): |
| g._flush_buffer() |
| keep_src, keep_dst = _choose(g, L, max_sources, max_destinations) |
| d: dict = {"site": [[k, w.a(np.asarray(v, np.float64))] for k, v in g.site.cal.items()], |
| "rolling": [[k, _ref_out(w, r)] for k, r in g.rolling.items()], |
| "anom": _ref_out(w, g.anom), "anom_s": _ref_out(w, g.anom_s), |
| "anom_age": [_ref_out(w, r) for r in g.anom_age] if g.anom_age is not None else None, |
| "seq_age": [_ref_out(w, r) for r in g.seq_age]} |
| for name in ("ev_ref", "ev_slow_ref", "dst_anom", "dst_anom_s", "dst_ev_ref"): |
| d[name] = _ref_out(w, getattr(g, name)) |
| d["roll_groups"] = [[k, _ref_out(w, r)] for k, r in g.roll_groups.items()] |
| d["http_ref_by_host"] = [[_jkey(h), _ref_out(w, r)] for h, r in g.http_ref_by_host.items()] |
| d["scalars"] = {"evid_drift": g.evid_drift, "dst_drift": g.dst_drift, "dst_rate_hi": g.dst_rate_hi, "last_recal": g.last_recal, |
| "warmup_short": getattr(g, "warmup_short", None), "_buf_now": getattr(g, "_buf_now", None)} |
| d["counts"] = dict(g.counts) |
| d["anchor"] = _buf_out(w, g._anchor) |
| d["buffer"] = _buf_out(w, g._buffer, buffer_rows) |
| d["evidence"] = _od_out(w, g.evidence.e, 4, keep_src) |
| d["evidence_slow"] = _od_out(w, g.evidence_slow.e, 4, keep_src) |
| d["dst_evidence"] = _od_out(w, g.dst_evidence.e, 4, keep_dst) |
| d["entities"] = _od_out(w, g.entities.s, 4, keep_src) |
| d["state_src"] = _ent_out(w, g.state.src, keep_src) |
| d["state_dst"] = _ent_out(w, g.state.dst, keep_dst) |
| gru = [(k, v) for k, v in g.gru.items() if keep_src is None or k in keep_src] |
| if g.seq is not None: |
| d["gru"] = {"k": w.ks(k for k, _ in gru), "left": w.a(np.array([int(v[1]) for _, v in gru], np.int64)), |
| "h": w.a(np.stack([np.asarray(v[0]) for _, v in gru]) if gru else np.zeros((0, g.seq.layers, g.seq.d), np.float32))} |
| for name in ("global_http", "global_scan"): |
| gr = getattr(g, name) |
| d[name] = {"c": _od_out(w, gr.c, 3, last=max_global), "ref": _ref_out(w, gr.ref)} |
| p = g.pressure |
| d["pressure"] = {"t": p.t, "w": p.w, "a": p.a, "fam": [[f, v] for f, v in p.fam.items()]} |
| d["drift"] = {"log_m": g.drift.log_m, "n": g.drift.n} |
| profs, ix = [], {} |
|
|
| def pref(prof): |
| if prof is None: |
| return None |
| if g.http is not None and prof is g.http.site: |
| return "shipped" |
| if id(prof) not in ix: |
| ix[id(prof)] = len(profs) |
| profs.append(_prof_out(w, prof)) |
| return ix[id(prof)] |
| d["site_http"] = pref(g.site_http) |
| d["site_http_by_host"] = [[_jkey(h), pref(p_)] for h, p_ in g.site_http_by_host.items()] |
| d["profiles"] = profs |
| d["log"] = g.log[-max_log:] if max_log else [] |
| d["extra"] = _extras_out(w, g, _GUARD_KNOWN, stats, "guard", max_sources) |
| stats.update(sources=len(g.state.src) if keep_src is None else len([k for k in g.state.src if k in keep_src]), |
| sources_total=len(g.state.src), destinations=len(g.state.dst) if keep_dst is None else len(keep_dst & set(g.state.dst)), |
| evidence_sources=len(g.evidence.e) if keep_src is None else len(keep_src & set(g.evidence.e)), |
| warm_sources=len(gru), bounded=keep_src is not None or keep_dst is not None, |
| buffer_rows=int(len(g._buffer["t"])) if buffer_rows is None else int(min(buffer_rows, len(g._buffer["t"]))), |
| log=len(d["log"])) |
| return d |
|
|
|
|
| def _layer_out(w: _W, L, ring: int, stats: dict, max_sources: int) -> dict: |
| r = getattr(L, "_http_ring", None) |
| return {"extra": _extras_out(w, L, _LAYER_KNOWN, stats, "layer", max_sources), |
| "http_boot_by": [[k, list(v)] for k, v in getattr(L, "_http_boot_by", {}).items()], |
| "http_recent_by": [[k, list(v)] for k, v in getattr(L, "_http_recent_by", {}).items()], |
| "http_last_by": [[k, v] for k, v in getattr(L, "_http_last_by", {}).items()], |
| "http_last": getattr(L, "_http_last", 0.0), |
| "http_ring": [list(x) for x in itertools.islice(reversed(r), ring)][::-1] if r is not None else [], |
| "ring_maxlen": r.maxlen if r is not None else None, |
| "playbook_hits": dict(getattr(L, "playbook_hits", {})), "posture_findings": list(getattr(L, "posture_findings", []))} |
|
|
|
|
| def _db_out(w: _W, dbg, max_log: int, max_rows: int, stats: dict) -> dict: |
| from dibaguard.shield import database as dbm |
| return {"extra": _extras_out(w, dbg, _DB_KNOWN, stats, "db", max_rows), |
| "known_clients": [[db, sorted(c, key=str)] for db, c in dbg.known_clients.items()], |
| "known_shapes": [[db, sorted(s, key=str)] for db, s in dbg.known_shapes.items()], |
| "vol_ref": [[_jkey(k), [[_jkey(c), v] for c, v in ref.items()]] for k, ref in dbg.vol_ref.items()], |
| "net_ref": [[db, w.a(np.asarray(v))] for db, v in dbg.net_ref.items()], |
| "vol": [[m, _od_out(w, dec.d, 2, last=max_rows)] for m, dec in dbg.vol.items()], |
| "fails": _od_out(w, dbg.fails.d, 2, last=max_rows), |
| "exempt_shapes": sorted(dbg.exempt_shapes, key=str), "log": dbg.log[-max_log:] if max_log else [], "counts": dict(dbg.counts), |
| "mysql_threads": [[tid, u, c] for tid, (u, c) in list(dbm._MYSQL_THREADS.items())]} |
|
|
|
|
| @contextlib.contextmanager |
| def _no_gc(): |
| """the cyclic GC re-scans a heap of millions of per-source objects every few thousand allocations; a capture or a |
| restore allocates hundreds of thousands, so the collector is paused for their duration (refcounting still frees)""" |
| was = gc.isenabled() |
| gc.disable() |
| try: |
| yield |
| finally: |
| if was: |
| gc.enable() |
|
|
|
|
| def capture(layer, **kw) -> tuple[dict, bytes, dict]: |
| """everything that matters operationally -> (numeric members incl. the key table, JSON structure, stats); see |
| _capture for the bounds. Call it under the engine lock (snapshot() holds the lock only for the first phase).""" |
| with _no_gc(): |
| w, blob, stats = _capture(layer, **kw) |
| return w.finish(), blob, stats |
|
|
|
|
| def _capture(layer, *, parts=ALL_PARTS, max_sources: int = 250_000, max_destinations: int = 100_000, buffer_rows: int | None = None, |
| max_log: int = 20_000, max_audit: int = 2_000, ring: int = 5_000, max_global: int = 100_000) -> tuple[_W, bytes, dict]: |
| """The part that needs the engine lock. Every live structure is serialised here (JSON), copied into new arrays, or |
| copied as plain values (stream-counter records), so later in-place changes cannot leak into the snapshot; what is |
| left for w.finish() (outside the lock) is pure encoding of those copies - key indices, the stream-counter arrays, |
| the object columns of the re-calibration buffer and anchor (arrays the engine replaces, never mutates) - and the |
| key table. Bounds: per-source tables keep the sources in force, the riskiest and the most recently seen up to |
| `max_sources`; the re-calibration buffer keeps its most recent `buffer_rows` (all by default); logs their tail.""" |
| g = getattr(layer, "guard", layer) |
| L = layer if hasattr(layer, "enforcer") else None |
| w, st, stats = _W(), {}, {} |
| if "guard" in parts: |
| st["guard"] = _guard_out(w, g, L, stats, max_sources, max_destinations, buffer_rows, max_log, max_global) |
| if L is not None: |
| if "layer" in parts: |
| st["layer"] = _layer_out(w, L, ring, stats, max_sources) |
| if "db" in parts and getattr(L, "dbg", None) is not None: |
| st["db"] = _db_out(w, L.dbg, max_log, max_global, stats) |
| if "enforcement" in parts: |
| st["enforcement"] = {"active": [[ip, dict(s)] for ip, s in L.enforcer.active.items()], "audit": L.enforcer.audit[-max_audit:] if max_audit else []} |
| stats["enforced"] = len(L.enforcer.active) |
| if "feedback" in parts and getattr(L, "feedback", None) is not None: |
| st["feedback"] = {"records": list(L.feedback.records), "allow": [[e, u] for e, u in L.feedback.allow.items()]} |
| if "rules" in parts: |
| st["rules"] = [r.as_dict() for r in L.rules.rules] |
| if "allowlist" in parts: |
| st["allowlist"] = {"allow": list(L.policy.d.get("allow", [])), "policy": L.policy.d} |
| stats["parts"] = sorted(st) |
| return w, _canon(st), stats |
|
|
|
|
| |
| def _digest(arrays: dict, manifest: dict, key) -> str: |
| h = hmac.new(key, digestmod=hashlib.sha256) if key else hashlib.sha256() |
| h.update(_canon({k: v for k, v in manifest.items() if k != "digest"})) |
| for name in sorted(arrays): |
| a = arrays[name] |
| h.update(f"\x00{name}\x00{a.dtype.str}\x00{a.shape}\x00".encode()) |
| if a.size: |
| h.update(np.ascontiguousarray(a).reshape(-1).view(np.uint8)) |
| return h.hexdigest() |
|
|
|
|
| def _target(path: str, replica: str) -> tuple[str, bool]: |
| """(file to write, directory mode?)""" |
| if path.lower().endswith(".npz") and not os.path.isdir(path): |
| os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True) |
| return path, False |
| os.makedirs(path, exist_ok=True) |
| t = time.time() |
| while True: |
| stamp = time.strftime("%Y%m%dT%H%M%S", time.gmtime(t)) + f"{int((t % 1) * 1000):03d}" |
| f = os.path.join(path, f"dibaguard-snap-{stamp}Z-{replica}.npz") |
| if not os.path.exists(f): |
| return f, True |
| t += 0.001 |
|
|
|
|
| def _prune(d: str, replica: str, keep: int) -> list: |
| """keep the newest `keep` snapshots of this replica; drop temp files of crashed writers (older than an hour)""" |
| removed = [] |
| mine = sorted(n for n in os.listdir(d) if (m := _SNAP.match(n)) and m.group(2) == replica) |
| for n in mine[: max(0, len(mine) - max(1, keep))]: |
| with contextlib.suppress(OSError): |
| os.remove(os.path.join(d, n)) |
| removed.append(n) |
| for n in os.listdir(d): |
| p = os.path.join(d, n) |
| if _TMP.match(n): |
| with contextlib.suppress(OSError): |
| if time.time() - os.path.getmtime(p) > 3600: |
| os.remove(p) |
| return removed |
|
|
|
|
| def _write(final: str, arrays: dict, manifest: dict) -> None: |
| d = os.path.dirname(os.path.abspath(final)) |
| tmp = os.path.join(d, f".{os.path.basename(final)}.{os.getpid()}.{threading.get_ident()}.tmp") |
| try: |
| with open(tmp, "wb") as f: |
| np.savez_compressed(f, manifest=np.frombuffer(_canon(manifest), np.uint8), **arrays) |
| f.flush() |
| os.fsync(f.fileno()) |
| _replace(tmp, final) |
| except BaseException: |
| with contextlib.suppress(OSError): |
| os.remove(tmp) |
| raise |
| _fsync_dir(d) |
|
|
|
|
| def snapshot(layer, path: str, *, keep: int = 5, replica: str | None = None, key=None, lock=None, parts=ALL_PARTS, **bounds) -> dict: |
| """Write one snapshot. `path`: a directory (a timestamped file is written there and only the newest `keep` of this |
| replica are kept) or a file ending in .npz. The capture runs under `lock` (the console / engine lock); hashing, |
| compression and the write run outside it. Returns file, bytes, capture_s (lock held), write_s, stats.""" |
| replica = _safe_name(replica or default_replica()) |
| key = _key(key) |
| t0 = time.perf_counter() |
| with _no_gc(): |
| with _lock(lock): |
| w, blob, stats = _capture(layer, parts=parts, **bounds) |
| t1 = time.perf_counter() |
| arrays = w.finish() |
| stats["keys"] = len(w.keys) |
| fp = model_fingerprint(layer) |
| arrays["state"] = np.frombuffer(blob, np.uint8) |
| manifest = {"magic": MAGIC, "schema": SCHEMA, "created": time.time(), "replica": replica, "host": socket.gethostname(), |
| "fingerprint": fp, "layout": code_layout()["id"], "parts": stats["parts"], "stats": stats, |
| "mac": "hmac-sha256" if key else "sha256", "numpy": np.__version__} |
| manifest["digest"] = _digest(arrays, manifest, key) |
| final, dir_mode = _target(path, replica) |
| _write(final, arrays, manifest) |
| pruned = _prune(os.path.dirname(os.path.abspath(final)), replica, keep) if dir_mode else [] |
| t2 = time.perf_counter() |
| return {"file": final, "bytes": os.path.getsize(final), "capture_s": round(t1 - t0, 4), "write_s": round(t2 - t1, 4), |
| "stats": stats, "pruned": pruned, "fingerprint": fp["id"]} |
|
|
|
|
| def _candidates(path: str, replica: str | None = None) -> list: |
| if os.path.isdir(path): |
| names = [(m.group(1), n) for n in os.listdir(path) if (m := _SNAP.match(n)) and (replica is None or m.group(2) == replica)] |
| return [os.path.join(path, n) for _, n in sorted(names, reverse=True)] |
| return [path] if os.path.exists(path) else [] |
|
|
|
|
| def _load(path: str, key, max_bytes: int = 8 << 30) -> tuple[dict, dict]: |
| """read + verify one snapshot file; never unpickles anything (allow_pickle=False), refuses zip bombs""" |
| try: |
| with zipfile.ZipFile(path) as zf: |
| if sum(i.file_size for i in zf.infolist()) > max_bytes: |
| raise SnapshotError(f"{path}: uncompressed size over {max_bytes} bytes") |
| with np.load(path, allow_pickle=False) as z: |
| names = set(z.files) |
| if not {"manifest", "state"} <= names: |
| raise SnapshotError(f"{path}: not a DibaGuard snapshot (no manifest / state)") |
| mb = z["manifest"] |
| if mb.dtype != np.uint8 or mb.ndim != 1: |
| raise SnapshotError(f"{path}: malformed manifest") |
| man = json.loads(mb.tobytes().decode("ascii")) |
| if not isinstance(man, dict) or man.get("magic") != MAGIC: |
| raise SnapshotError(f"{path}: not a DibaGuard snapshot (bad magic)") |
| if not 1 <= int(man.get("schema", 0)) <= SCHEMA: |
| raise SnapshotError(f"{path}: schema {man.get('schema')} is not supported by this version (1..{SCHEMA})") |
| arrays = {n: z[n] for n in names if n != "manifest"} |
| except SnapshotError: |
| raise |
| except Exception as e: |
| raise SnapshotError(f"{path}: unreadable ({type(e).__name__}: {e})") from None |
| mac = man.get("mac") |
| if key is not None and mac != "hmac-sha256": |
| raise SnapshotError(f"{path}: not signed, but a snapshot key is configured (refusing a possible downgrade)") |
| if key is None and mac == "hmac-sha256": |
| raise SnapshotError(f"{path}: signed snapshot - set the key (DIBAGUARD_SNAPSHOT_KEY) to verify it") |
| if mac not in ("sha256", "hmac-sha256") or not hmac.compare_digest(str(man.get("digest", "")), _digest(arrays, man, key)): |
| raise SnapshotError(f"{path}: integrity check failed (digest mismatch: corrupted or tampered)") |
| if any(arrays[n].dtype != np.uint8 or arrays[n].ndim != 1 for n in ("state", "keys") if n in arrays): |
| raise SnapshotError(f"{path}: malformed state") |
| return man, arrays |
|
|
|
|
| def verify(path: str, key=None) -> dict: |
| """integrity check only (no restore): -> the manifest""" |
| return _load(path, _key(key))[0] |
|
|
|
|
| |
| def _guard_in(rd: _R, d: dict, g) -> list: |
| """decode the guard part -> a list of assignments (applied only after everything decoded)""" |
| ops: list = [] |
|
|
| def put(obj, attr, val): |
| ops.append((obj, attr, val)) |
| put(g.site, "cal", {k: rd.a(v).tolist() for k, v in d["site"]}) |
| rolling = dict(g.rolling) |
| rolling.update({k: _ref_in(rd, v) for k, v in d["rolling"]}) |
| put(g, "rolling", rolling) |
| if g.anom is not None and d.get("anom") is not None: |
| put(g, "anom", _ref_in(rd, d["anom"])) |
| put(g, "anom_s", _ref_in(rd, d["anom_s"])) |
| if g.anom_age is not None and d.get("anom_age") is not None: |
| put(g, "anom_age", [_ref_in(rd, x) for x in d["anom_age"]]) |
| put(g, "seq_age", [_ref_in(rd, x) for x in d["seq_age"]]) |
| for name in ("ev_ref", "ev_slow_ref", "dst_anom", "dst_anom_s", "dst_ev_ref"): |
| put(g, name, _ref_in(rd, d[name])) |
| put(g, "roll_groups", {k: _ref_in(rd, v) for k, v in d["roll_groups"]}) |
| put(g, "http_ref_by_host", {_ukey(h): _ref_in(rd, v) for h, v in d["http_ref_by_host"]}) |
| for k, v in d["scalars"].items(): |
| if v is not None: |
| put(g, k, v) |
| put(g, "_anchor", _buf_in(rd, d["anchor"])) |
| if d.get("buffer") is not None: |
| put(g, "_buffer", _buf_in(rd, d["buffer"])) |
| put(g.evidence, "e", _od_in(rd, d["evidence"], (2,), g.evidence.max)) |
| put(g.evidence_slow, "e", _od_in(rd, d["evidence_slow"], (2,), g.evidence_slow.max)) |
| put(g.dst_evidence, "e", _od_in(rd, d["dst_evidence"], (2,), g.dst_evidence.max)) |
| put(g.entities, "s", _od_in(rd, d["entities"], (2, 3), g.entities.max)) |
| put(g.state, "src", _ent_in(rd, d["state_src"], g.state.max)) |
| put(g.state, "dst", _ent_in(rd, d["state_dst"], g.state.max)) |
| if g.seq is not None and d.get("gru") is not None: |
| ks, left, hs = rd.ks(d["gru"]["k"]), rd.a(d["gru"]["left"]).tolist(), rd.a(d["gru"]["h"]) |
| put(g, "gru", OrderedDict((k, (hs[i], int(left[i]))) for i, k in enumerate(ks))) |
| for name in ("global_http", "global_scan"): |
| gr = getattr(g, name) |
| put(gr, "c", _od_in(rd, d[name]["c"], (), gr.max)) |
| put(gr, "ref", _ref_in(rd, d[name]["ref"])) |
| p = d["pressure"] |
| for k in ("t", "w", "a"): |
| put(g.pressure, k, p[k]) |
| put(g.pressure, "fam", {f: v for f, v in p["fam"]}) |
| put(g.drift, "log_m", float(d["drift"]["log_m"])) |
| put(g.drift, "n", int(d["drift"]["n"])) |
| profs = [_prof_in(rd, x) for x in d["profiles"]] |
|
|
| def prof(ref): |
| if ref == "shipped": |
| return g.http.site if g.http is not None else None |
| return None if ref is None else profs[int(ref)] |
| put(g, "site_http", prof(d["site_http"])) |
| put(g, "site_http_by_host", {_ukey(h): prof(r) for h, r in d["site_http_by_host"]}) |
| put(g, "log", list(d["log"]) + list(g.log)) |
| counts = dict(g.counts) |
| for k, v in d["counts"].items(): |
| counts[k] = counts.get(k, 0) + v |
| put(g, "counts", counts) |
| ops += _extras_in(rd, d.get("extra", []), g) |
| return ops |
|
|
|
|
| def _layer_in(rd: _R, d: dict, L) -> list: |
| ops: list = _extras_in(rd, d.get("extra", []), L) + [(L, "_http_boot_by", {k: list(v) for k, v in d["http_boot_by"]}), |
| (L, "_http_recent_by", {k: list(v) for k, v in d["http_recent_by"]}), |
| (L, "_http_last_by", {k: v for k, v in d["http_last_by"]}), (L, "_http_last", d["http_last"]), |
| (L, "playbook_hits", dict(d["playbook_hits"])), (L, "posture_findings", list(d["posture_findings"]))] |
| r = getattr(L, "_http_ring", None) |
| maxlen = r.maxlen if r is not None else d.get("ring_maxlen") |
| ops.append((L, "_http_ring", deque([tuple(x) for x in d["http_ring"]], maxlen=maxlen))) |
| return ops |
|
|
|
|
| def _db_in(rd: _R, d: dict, L) -> list: |
| from dibaguard.shield import database as dbm |
| dbg = L.dbg |
| ops: list = [(dbg, "known_clients", {db: set(c) for db, c in d["known_clients"]}), |
| (dbg, "known_shapes", {db: set(s) for db, s in d["known_shapes"]}), |
| (dbg, "vol_ref", {_ukey(k): {_ukey(c): v for c, v in ref} for k, ref in d["vol_ref"]}), |
| (dbg, "net_ref", {db: rd.a(n) for db, n in d["net_ref"]}), |
| (dbg.fails, "d", _od_in(rd, d["fails"], (), dbg.fails.max)), |
| (dbg, "exempt_shapes", set(dbg.exempt_shapes) | set(d["exempt_shapes"])), |
| (dbg, "log", list(d["log"]) + list(dbg.log)), |
| (dbg, "counts", dict(dbg.counts, **{k: dbg.counts.get(k, 0) + v for k, v in d["counts"].items()}))] |
| for m, dd in d["vol"]: |
| if m in dbg.vol: |
| ops.append((dbg.vol[m], "d", _od_in(rd, dd, (), dbg.vol[m].max))) |
| threads = OrderedDict((tid, (u, c)) for tid, u, c in d["mysql_threads"]) |
| ops.append((None, "mysql_threads", (dbm._MYSQL_THREADS, threads))) |
| return ops + _extras_in(rd, d.get("extra", []), dbg) |
|
|
|
|
| def _restore_allowlist(L, d: dict, mode: str) -> dict: |
| from dibaguard.shield.policy import Policy |
| if mode == "keep": |
| return {"mode": mode} |
| if mode == "replace" and d.get("policy"): |
| p = Policy(d["policy"]) |
| L.policy = p |
| L.enforcer.policy = p |
| return {"mode": mode, "allow": len(p.d["allow"])} |
| cur, added = list(L.policy.d.get("allow", [])), 0 |
| for a in d.get("allow", []): |
| if a in cur: |
| continue |
| try: |
| ipaddress.ip_network(a, strict=False) |
| except ValueError: |
| continue |
| cur.append(a) |
| added += 1 |
| if added: |
| L.policy.d["allow"] = cur |
| L.policy._allow = [ipaddress.ip_network(a, strict=False) for a in cur] |
| return {"mode": "merge", "added": added} |
|
|
|
|
| def _restore_feedback(L, d: dict) -> dict: |
| fb = getattr(L, "feedback", None) |
| if fb is None: |
| return {} |
|
|
| def sig(r): |
| return (r.get("ts"), r.get("entity"), r.get("verdict"), r.get("who")) |
| have = {sig(r) for r in fb.records} |
| new = [r for r in d.get("records", []) if sig(r) not in have] |
| if new: |
| fb.records = sorted(fb.records + new, key=lambda r: float(r.get("ts", 0) or 0)) |
| if fb.path: |
| with contextlib.suppress(OSError), open(fb.path, "a", encoding="utf-8") as f: |
| for r in new: |
| f.write(json.dumps(r, ensure_ascii=False, default=_jdefault) + "\n") |
| allowed = 0 |
| for e, u in d.get("allow", []): |
| if isinstance(u, (int, float)) and u > fb.allow.get(e, 0): |
| fb.allow[e] = u |
| allowed += 1 |
| return {"records_added": len(new), "allow": allowed} |
|
|
|
|
| def _restore_rules(L, items: list) -> dict: |
| from dibaguard.shield.rules import Rule |
|
|
| def sig(text, scope, when, action): |
| return (text, scope, _canon(when), action) |
| have = {sig(r.text, r.scope, r.when, r.action) for r in L.rules.rules} |
| taken = {r.id for r in L.rules.rules} |
| added = 0 |
| for d in items: |
| s = sig(d.get("text", ""), d.get("scope", "source"), d.get("when", {}), d.get("action", "monitor")) |
| if s in have: |
| continue |
| d = dict(d) |
| if not d.get("id") or d["id"] in taken: |
| d["id"] = f"R{1 + max((int(x[1:]) for x in taken if re.fullmatch(r'R\d+', x)), default=0):03d}" |
| try: |
| r = Rule(d) |
| except ValueError: |
| continue |
| L.rules.rules.append(r) |
| taken.add(r.id) |
| have.add(s) |
| added += 1 |
| return {"added": added, "total": len(L.rules.rules)} |
|
|
|
|
| def _reapply_block(backend, ip: str, ttl: float) -> bool: |
| """put a restored block back into the host firewall for its remaining time (iptables: only if not there already)""" |
| from dibaguard.shield.enforce import Iptables |
| if isinstance(backend, Iptables) and backend.run(["iptables", "-C", "DIBAGUARD", "-s", ip, "-j", "DROP"]): |
| return True |
| return bool(backend.block(ip, max(1.0, float(ttl)))) |
|
|
|
|
| def _restore_enforcement(L, d: dict, now: float, reapply: bool) -> dict: |
| """decisions keep their absolute expiry; expired ones are dropped; a stronger / later local decision wins""" |
| E = L.enforcer |
| out = {"restored": 0, "expired": 0, "kept_local": 0, "allowlisted": 0, "reapplied": 0} |
| for ip, st in d.get("active", []): |
| a = st.get("action") |
| if a not in RANK or a == "monitor": |
| continue |
| until = float(st.get("until", 0)) |
| if until <= now: |
| out["expired"] += 1 |
| continue |
| if E.policy.allowed(ip): |
| out["allowlisted"] += 1 |
| continue |
| cur = E.active.get(ip) |
| if cur is not None: |
| rc, ra = RANK.get(cur["action"], 0), RANK[a] |
| if rc > ra or (rc == ra and float(cur["until"]) >= until): |
| out["kept_local"] += 1 |
| continue |
| if rc == ra: |
| cur["until"] = until |
| out["restored"] += 1 |
| continue |
| if reapply and E.policy.enforcing and a == "block" and not (cur is not None and cur["action"] == "block"): |
| out["reapplied"] += int(_reapply_block(E.backend, ip, until - now)) |
| E.active[ip] = dict(st, until=until) |
| out["restored"] += 1 |
| audit = list(d.get("audit", [])) |
| if audit: |
| E.audit[:0] = audit |
| if len(E.audit) > 100_000: |
| del E.audit[: len(E.audit) - 100_000] |
| return out |
|
|
|
|
| def restore(layer, path: str, *, now: float | None = None, force: bool = False, parts=None, key=None, reapply: bool = True, |
| policy_mode: str = "merge", replica: str | None = None) -> dict: |
| """Restore the newest valid snapshot at `path` (directory: newest first, a corrupted file falls back to the |
| previous one; or one .npz file) into a Layer (or a bare Guard). |
| |
| The snapshot must have been made for the same model files and feature layout (fingerprint), else |
| FingerprintMismatch - unless force=True, which restores only the model-independent parts (enforcement, feedback, |
| rules, allowlist). Enforcement keeps absolute expiry times: entries expired at `now` (default: wall clock) are |
| dropped, blocks are put back into the host firewall for their remaining time (reapply). policy_mode: "merge" |
| (union of allowlists, default), "replace" (the snapshot's whole runtime policy), "keep" (ignore it). |
| Everything is decoded before anything is assigned: a failing restore leaves the layer as it was.""" |
| key = _key(key) |
| now = time.time() if now is None else float(now) |
| files = _candidates(path, replica) |
| if not files: |
| raise SnapshotError(f"no snapshot at {path}") |
| skipped = [] |
| for f in files: |
| try: |
| man, arrays = _load(f, key) |
| break |
| except SnapshotError as e: |
| skipped.append(str(e)) |
| else: |
| raise SnapshotError(f"no valid snapshot at {path}: " + " | ".join(skipped)) |
| fp, lay = model_fingerprint(layer), code_layout()["id"] |
| sfp = man.get("fingerprint") or {} |
| same = sfp.get("id") == fp["id"] and man.get("layout") == lay |
| if same: |
| want = [p for p in ALL_PARTS if parts is None or p in parts] |
| else: |
| diff = [k for k in ("tier1", "seq", "http", "dbnet") if sfp.get(k) != fp.get(k)] + ([] if man.get("layout") == lay else ["feature layout"]) |
| if not force: |
| raise FingerprintMismatch(f"{f}: made for other model files ({', '.join(diff) or 'fingerprint'} differ); " |
| f"restore(force=True) brings back only {', '.join(MODEL_INDEPENDENT)}") |
| want = [p for p in MODEL_INDEPENDENT if parts is None or p in parts] |
| try: |
| with _no_gc(): |
| st, rd, ops, done = _decode(arrays, layer, want) |
| except SnapshotError: |
| raise |
| except Exception as e: |
| raise SnapshotError(f"{f}: cannot be decoded by this version ({type(e).__name__}: {e})") from None |
| for obj, attr, val in ops: |
| if obj is None and attr == "mysql_threads": |
| target, threads = val |
| target.clear() |
| target.update(threads) |
| else: |
| setattr(obj, attr, val) |
| L = layer if hasattr(layer, "enforcer") else None |
| rep = {"file": f, "created": man.get("created"), "age_s": round(time.time() - float(man.get("created", time.time())), 3), |
| "replica": man.get("replica"), "fingerprint_match": same, "parts": done, "skipped": skipped, "stats": man.get("stats", {})} |
| if L is not None: |
| if "allowlist" in want and "allowlist" in st: |
| rep["allowlist"] = _restore_allowlist(L, st["allowlist"], policy_mode) |
| rep["parts"].append("allowlist") |
| if "feedback" in want and "feedback" in st: |
| rep["feedback"] = _restore_feedback(L, st["feedback"]) |
| rep["parts"].append("feedback") |
| if "rules" in want and "rules" in st: |
| rep["rules"] = _restore_rules(L, st["rules"]) |
| rep["parts"].append("rules") |
| if "enforcement" in want and "enforcement" in st: |
| rep["enforcement"] = _restore_enforcement(L, st["enforcement"], now, reapply) |
| rep["parts"].append("enforcement") |
| return rep |
|
|
|
|
| def _decode(arrays: dict, layer, want: list) -> tuple: |
| """state JSON + members -> (structure, reader, assignments, model parts decoded); assigns nothing""" |
| st = json.loads(arrays.pop("state").tobytes().decode("ascii")) |
| rd = _R(arrays, json.loads(arrays.pop("keys").tobytes().decode("ascii")) if "keys" in arrays else []) |
| g = getattr(layer, "guard", layer) |
| L = layer if hasattr(layer, "enforcer") else None |
| ops: list = [] |
| done: list = [] |
| if "guard" in want and "guard" in st: |
| ops += _guard_in(rd, st["guard"], g) |
| done.append("guard") |
| if L is not None and "layer" in want and "layer" in st: |
| ops += _layer_in(rd, st["layer"], L) |
| done.append("layer") |
| if L is not None and "db" in want and "db" in st and getattr(L, "dbg", None) is not None: |
| ops += _db_in(rd, st["db"], L) |
| done.append("db") |
| return st, rd, ops, done |
|
|
|
|
| class Snapshotter: |
| """Periodic snapshots for the sidecar: maybe() from its loop (or Housekeeper's thread) writes one every `every_s` |
| seconds; restore_on_start() before serving. The capture holds `lock` (the console lock); the write does not.""" |
|
|
| def __init__(self, layer, directory: str, every_s: float = 300.0, keep: int = 5, lock=None, replica: str | None = None, |
| key=None, **bounds): |
| self.layer, self.directory, self.every_s, self.keep, self.lock = layer, directory, float(every_s), int(keep), lock |
| self.replica = _safe_name(replica or default_replica()) |
| self.key, self.bounds = key, bounds |
| self._last = -math.inf |
| self.last_info: dict | None = None |
| self.errors: list = [] |
|
|
| def restore_on_start(self, now: float | None = None, on_mismatch: str = "independent", **kw) -> dict: |
| """newest valid snapshot of this replica (any replica if it has none: a moved / promoted host). A snapshot of |
| other model files restores its model-independent parts (on_mismatch="independent", the default: blocks and |
| operator decisions survive a model upgrade), or nothing ("skip"), or raises ("raise").""" |
| last_err = None |
| for r in (self.replica, None): |
| if not _candidates(self.directory, r): |
| continue |
| try: |
| with _lock(self.lock): |
| return dict(restore(self.layer, self.directory, now=now, key=self.key, replica=r, **kw), restored=True) |
| except FingerprintMismatch as e: |
| if on_mismatch == "raise": |
| raise |
| if on_mismatch == "skip": |
| return {"restored": False, "reason": str(e)} |
| with _lock(self.lock): |
| return dict(restore(self.layer, self.directory, now=now, key=self.key, replica=r, force=True, **kw), restored=True, mismatch=str(e)) |
| except SnapshotError as e: |
| last_err = str(e) |
| return {"restored": False, "reason": last_err or f"no snapshot in {self.directory}"} |
|
|
| def snapshot_now(self) -> dict: |
| info = snapshot(self.layer, self.directory, keep=self.keep, replica=self.replica, key=self.key, lock=self.lock, **self.bounds) |
| self._last = time.monotonic() |
| self.last_info = info |
| return info |
|
|
| def maybe(self, now=None) -> dict | None: |
| if time.monotonic() - self._last < self.every_s: |
| return None |
| try: |
| return self.snapshot_now() |
| except Exception as e: |
| self._last = time.monotonic() |
| self.errors.append(f"{time.time():.0f} {type(e).__name__}: {e}") |
| del self.errors[:-20] |
| return {"error": str(e)} |
|
|
|
|
| |
| def affinity_key(ip, v4_prefix: int = 24, v6_prefix: int = 64) -> str: |
| """the routing unit of a client: its /24 (IPv4) or /64 (IPv6); "a.b.c.d:port" and "[v6]:port" accepted; anything |
| that is not an address (an API key, a session id) is its own unit""" |
| s = str(ip).strip() |
| if s.startswith("[") and "]" in s: |
| s = s[1:s.index("]")] |
| elif s.count(":") == 1: |
| s = s.split(":", 1)[0] |
| try: |
| a = ipaddress.ip_address(s) |
| except ValueError: |
| return s.lower() |
| if a.version == 6 and a.ipv4_mapped is not None: |
| a = a.ipv4_mapped |
| return str(ipaddress.ip_network(f"{a}/{v4_prefix if a.version == 4 else v6_prefix}", strict=False)) |
|
|
|
|
| def _hrw(key: str, node) -> int: |
| return int.from_bytes(hashlib.blake2b(f"{key}\x00{node}".encode("utf-8", "surrogatepass"), digest_size=8).digest(), "big") |
|
|
|
|
| def affinity(ip, replicas, v4_prefix: int = 24, v6_prefix: int = 64, weights: dict | None = None): |
| """the replica that owns this source: rendezvous (highest-random-weight) hashing of its /24 (/64). Stable across |
| processes and hosts, independent of the order of `replicas`; removing a replica moves only its own sources, |
| adding one moves ~1/n of them. weights: optional {replica: capacity} (weighted rendezvous).""" |
| reps = list(replicas) |
| if not reps: |
| raise ValueError("no replicas") |
| k = affinity_key(ip, v4_prefix, v6_prefix) |
| if not weights: |
| return max(reps, key=lambda r: (_hrw(k, r), str(r))) |
| return max(reps, key=lambda r: (-float(weights.get(r, 1.0)) / math.log((_hrw(k, r) + 0.5) / 2.0 ** 64), str(r))) |
|
|
|
|
| |
| class View: |
| """the merged cluster view at one moment: decisions (ip -> record), lifts (ip -> newest lift), allows (entity -> until)""" |
|
|
| def __init__(self, decisions: dict, lifts: dict, allows: dict): |
| self.decisions, self.lifts, self.allows = decisions, lifts, allows |
|
|
|
|
| def merge(slots, now: float) -> View: |
| """Every replica's latest record per (kind, key) -> the cluster view. A lift cancels every decision (from any |
| replica) published at or before it; of the live decisions on one source the strongest action wins and the expiry |
| is the latest of them (as one Enforcer would have done had it seen every verdict); ties are broken by replica id, |
| so every replica computes the same view from the same records.""" |
| lifts: dict = {} |
| allows: dict = {} |
| cands: dict = {} |
| ok = [] |
| for r in slots: |
| try: |
| k = r.get("k") |
| if k == "d": |
| r["u"], r["ts"] = float(r["u"]), float(r["ts"]) |
| if r.get("a") not in RANK or not isinstance(r.get("ip"), str): |
| continue |
| elif k == "l": |
| r["ts"], r["h"] = float(r["ts"]), float(r.get("h", math.inf)) |
| if not isinstance(r.get("ip"), str): |
| continue |
| elif k == "a": |
| r["u"] = float(r["u"]) |
| if not isinstance(r.get("e"), str): |
| continue |
| else: |
| continue |
| except (KeyError, TypeError, ValueError, AttributeError): |
| continue |
| ok.append(r) |
| for r in ok: |
| k = r["k"] |
| if k == "l" and r["h"] > now: |
| cur = lifts.get(r["ip"]) |
| if cur is None or r["ts"] > cur["ts"]: |
| lifts[r["ip"]] = r |
| elif k == "a" and r["u"] > now: |
| allows[r["e"]] = max(allows.get(r["e"], 0.0), r["u"]) |
| for r in ok: |
| if r["k"] != "d" or r["u"] <= now: |
| continue |
| lf = lifts.get(r["ip"]) |
| if lf is not None and r["ts"] <= lf["ts"]: |
| continue |
| cands.setdefault(r["ip"], []).append(r) |
| dec = {} |
| for ip, cs in cands.items(): |
| best = max(cs, key=lambda r: (RANK[r["a"]], r["u"], str(r.get("r")))) |
| d = dict(best) |
| d["u"] = max(r["u"] for r in cs) |
| d["by"] = sorted({str(r.get("r")) for r in cs}) |
| dec[ip] = d |
| return View(dec, lifts, allows) |
|
|
|
|
| class SharedState: |
| """pluggable backend of the cluster view. Records: decision {"k": "d", ip, a(ction), u(ntil), ts, r(eplica), f, s, |
| why}; lift {"k": "l", ip, ts, r, h(old until), reset}; allow {"k": "a", e(ntity), u, ts, r}. Each replica writes |
| only its own records; the latest record per (replica, kind, key) is its current word on that key.""" |
| replica: str = "replica" |
|
|
| def publish(self, records: list) -> None: |
| raise NotImplementedError |
|
|
| def slots(self) -> list: |
| raise NotImplementedError |
|
|
| def view(self, now: float) -> View: |
| return merge(self.slots(), now) |
|
|
| def publish_evidence(self, summary: dict) -> None: |
| raise NotImplementedError |
|
|
| def evidence(self) -> list: |
| raise NotImplementedError |
|
|
| def compact(self, now: float) -> None: |
| pass |
|
|
| def heartbeat(self, info: dict | None = None) -> None: |
| pass |
|
|
|
|
| def _live(rec: dict, now: float) -> bool: |
| k = rec.get("k") |
| if k == "d" or k == "a": |
| return float(rec.get("u", 0)) > now |
| if k == "l": |
| return float(rec.get("h", 0)) > now |
| return False |
|
|
|
|
| def _slot(rec: dict) -> tuple: |
| return (rec.get("k"), _canon(rec.get("ip", rec.get("e")))) |
|
|
|
|
| class FileSharedState(SharedState): |
| """A directory on shared storage (NFS / SMB / a replicated volume). No locks are needed - each replica appends only |
| to its own journal and replaces only its own files: |
| enforce/<replica>.jsonl header {"gen"} + append-only records; compaction rewrites it (gen+1, atomic rename) |
| evidence/<replica>.json the replica's latest evidence summary (atomic replace) |
| replicas/<replica>.json heartbeat (host, pid, time) - two hosts must never share a replica id |
| Readers keep a per-journal offset and read only what was appended; a new generation means re-read from the top.""" |
|
|
| def __init__(self, root: str, replica: str | None = None, durable: bool = True, stale_s: float = 60.0): |
| self.root, self.replica, self.durable = root, _safe_name(replica or default_replica()), durable |
| for sub in ("enforce", "evidence", "replicas"): |
| os.makedirs(os.path.join(root, sub), exist_ok=True) |
| self._own = os.path.join(root, "enforce", f"{self.replica}.jsonl") |
| self._readers: dict = {} |
| self._ev_cache: dict = {} |
| hb = os.path.join(root, "replicas", f"{self.replica}.json") |
| with contextlib.suppress(OSError, ValueError): |
| with open(hb, encoding="utf-8") as fh: |
| prev = json.load(fh) |
| if prev.get("host") != socket.gethostname() and time.time() - float(prev.get("ts", 0)) < stale_s: |
| raise RuntimeError(f"replica id {self.replica!r} is in use by host {prev.get('host')!r} (heartbeat " |
| f"{time.time() - float(prev.get('ts', 0)):.0f} s ago): give every replica its own id") |
| self._gen = self._own_gen() |
| self.heartbeat() |
|
|
| def _own_gen(self) -> int: |
| try: |
| with open(self._own, "rb") as f: |
| return int(json.loads(f.readline())["gen"]) |
| except (OSError, ValueError, KeyError, TypeError): |
| gen = int(time.time() * 1000) |
| _atomic_bytes(self._own, _canon({"dibaguard-journal": 1, "replica": self.replica, "gen": gen}) + b"\n", self.durable) |
| return gen |
|
|
| def publish(self, records: list) -> None: |
| if not records: |
| return |
| data = b"".join(_canon(dict(r, r=self.replica)) + b"\n" for r in records) |
| with open(self._own, "ab+") as f: |
| size = f.seek(0, 2) |
| if size: |
| f.seek(size - 1) |
| if f.read(1) != b"\n": |
| data = b"\n" + data |
| f.write(data) |
| f.flush() |
| if self.durable: |
| os.fsync(f.fileno()) |
|
|
| def _read(self, r: str, path: str) -> None: |
| st = self._readers.setdefault(r, {"gen": None, "off": 0, "slots": {}, "bad": 0}) |
| try: |
| f = open(path, "rb") |
| except OSError: |
| st["slots"] = {} |
| st["gen"] = None |
| return |
| with f: |
| head = f.readline() |
| try: |
| gen = int(json.loads(head)["gen"]) |
| except (ValueError, KeyError, TypeError): |
| return |
| size = f.seek(0, 2) |
| if gen != st["gen"] or size < st["off"]: |
| st.update(gen=gen, off=len(head), slots={}) |
| f.seek(st["off"]) |
| data = f.read() |
| end = data.rfind(b"\n") |
| if end < 0: |
| return |
| for line in data[: end + 1].splitlines(): |
| try: |
| rec = json.loads(line) |
| if not isinstance(rec, dict) or rec.get("k") not in ("d", "l", "a"): |
| raise ValueError |
| except ValueError: |
| st["bad"] += 1 |
| continue |
| rec["r"] = r |
| sk = _slot(rec) |
| old = st["slots"].get(sk) |
| if old is None or float(rec.get("ts", 0)) >= float(old.get("ts", 0)): |
| st["slots"][sk] = rec |
| st["off"] += end + 1 |
|
|
| def slots(self) -> list: |
| d = os.path.join(self.root, "enforce") |
| present = set() |
| for n in os.listdir(d): |
| if n.endswith(".jsonl"): |
| r = n[:-6] |
| present.add(r) |
| self._read(r, os.path.join(d, n)) |
| for r in list(self._readers): |
| if r not in present: |
| del self._readers[r] |
| return [rec for st in self._readers.values() for rec in st["slots"].values()] |
|
|
| def compact(self, now: float) -> None: |
| """rewrite this replica's journal with only its live records (new generation, atomic rename)""" |
| self._read(self.replica, self._own) |
| live = [rec for rec in self._readers.get(self.replica, {}).get("slots", {}).values() if _live(rec, now)] |
| gen = max(self._gen + 1, int(time.time() * 1000)) |
| data = _canon({"dibaguard-journal": 1, "replica": self.replica, "gen": gen}) + b"\n" + b"".join(_canon(r) + b"\n" for r in live) |
| try: |
| _atomic_bytes(self._own, data, self.durable) |
| except PermissionError: |
| return |
| self._gen = gen |
|
|
| def publish_evidence(self, summary: dict) -> None: |
| with contextlib.suppress(PermissionError): |
| _atomic_bytes(os.path.join(self.root, "evidence", f"{self.replica}.json"), _canon(dict(summary, r=self.replica)), False) |
|
|
| def evidence(self) -> list: |
| d = os.path.join(self.root, "evidence") |
| out = [] |
| for n in os.listdir(d): |
| if not n.endswith(".json") or n[:-5] == self.replica: |
| continue |
| p = os.path.join(d, n) |
| try: |
| s = os.stat(p) |
| ck = (s.st_mtime_ns, s.st_size) |
| hit = self._ev_cache.get(n) |
| if hit is None or hit[0] != ck: |
| with open(p, "rb") as f: |
| hit = (ck, json.loads(f.read())) |
| self._ev_cache[n] = hit |
| out.append(dict(hit[1], r=n[:-5])) |
| except (OSError, ValueError): |
| continue |
| return out |
|
|
| def heartbeat(self, info: dict | None = None) -> None: |
| with contextlib.suppress(OSError): |
| _atomic_bytes(os.path.join(self.root, "replicas", f"{self.replica}.json"), |
| _canon(dict(info or {}, replica=self.replica, host=socket.gethostname(), pid=os.getpid(), ts=time.time())), False) |
|
|
| def alive(self, within_s: float = 30.0) -> list: |
| """replica ids with a fresh heartbeat (for affinity tables / dashboards)""" |
| d, out = os.path.join(self.root, "replicas"), [] |
| for n in os.listdir(d): |
| with contextlib.suppress(OSError, ValueError): |
| with open(os.path.join(d, n), encoding="utf-8") as f: |
| hb = json.load(f) |
| if time.time() - float(hb.get("ts", 0)) <= within_s: |
| out.append(hb.get("replica", n[:-5])) |
| return sorted(out) |
|
|
|
|
| def _s(x) -> str: |
| return x.decode("utf-8") if isinstance(x, (bytes, bytearray)) else str(x) |
|
|
|
|
| class RedisSharedState(SharedState): |
| """The same records in Redis (optional; `pip install redis` - nothing else in DibaGuard needs it): |
| {ns}:replicas set of replica ids |
| {ns}:slots:{replica} hash "<kind>|<key>" -> record JSON (HSET replaces: the latest word per slot) |
| {ns}:ev:{replica} string evidence summary, with a TTL |
| {ns}:hb:{replica} string heartbeat, with a TTL |
| Every replica writes only its own keys; the merge runs client-side (merge()), exactly as with the file backend. |
| `client` may be any object with the redis-py methods used here (sadd, smembers, hset, hgetall, hdel, set, get).""" |
|
|
| def __init__(self, url: str = "redis://127.0.0.1:6379/0", replica: str | None = None, ns: str = "dibaguard", client=None, |
| evidence_ttl_s: int = 600): |
| if client is None: |
| try: |
| import redis |
| except ImportError as e: |
| raise RuntimeError("RedisSharedState needs the 'redis' package (pip install redis); FileSharedState needs nothing") from e |
| client = redis.Redis.from_url(url) |
| self.c, self.ns, self.replica = client, ns, _safe_name(replica or default_replica()) |
| self.ev_ttl = int(evidence_ttl_s) |
| self.c.sadd(f"{ns}:replicas", self.replica) |
|
|
| def _hk(self, r: str) -> str: |
| return f"{self.ns}:slots:{r}" |
|
|
| def publish(self, records: list) -> None: |
| if records: |
| m = {} |
| for rec in records: |
| rec = dict(rec, r=self.replica) |
| m[f"{rec['k']}|{_canon(rec.get('ip', rec.get('e'))).decode()}"] = _canon(rec).decode() |
| self.c.hset(self._hk(self.replica), mapping=m) |
|
|
| def _replicas(self) -> list: |
| return sorted(_s(r) for r in (self.c.smembers(f"{self.ns}:replicas") or ())) |
|
|
| def slots(self) -> list: |
| out = [] |
| for r in self._replicas(): |
| for v in (self.c.hgetall(self._hk(r)) or {}).values(): |
| try: |
| rec = json.loads(_s(v)) |
| except ValueError: |
| continue |
| if isinstance(rec, dict) and rec.get("k") in ("d", "l", "a"): |
| rec["r"] = r |
| out.append(rec) |
| return out |
|
|
| def compact(self, now: float) -> None: |
| h = self.c.hgetall(self._hk(self.replica)) or {} |
| dead = [] |
| for f, v in h.items(): |
| try: |
| if not _live(json.loads(_s(v)), now): |
| dead.append(f) |
| except ValueError: |
| dead.append(f) |
| if dead: |
| self.c.hdel(self._hk(self.replica), *dead) |
|
|
| def publish_evidence(self, summary: dict) -> None: |
| self.c.set(f"{self.ns}:ev:{self.replica}", _canon(dict(summary, r=self.replica)).decode(), ex=self.ev_ttl) |
|
|
| def evidence(self) -> list: |
| out = [] |
| for r in self._replicas(): |
| if r == self.replica: |
| continue |
| v = self.c.get(f"{self.ns}:ev:{r}") |
| if v: |
| with contextlib.suppress(ValueError): |
| out.append(dict(json.loads(_s(v)), r=r)) |
| return out |
|
|
| def heartbeat(self, info: dict | None = None) -> None: |
| hb = _canon(dict(info or {}, host=socket.gethostname(), pid=os.getpid(), ts=time.time())).decode() |
| self.c.set(f"{self.ns}:hb:{self.replica}", hb, ex=self.ev_ttl) |
|
|
|
|
| |
| def evidence_thresholds(g) -> tuple[float, float]: |
| """the engine's own source-evidence thresholds (fast, slow): alpha/4 each, conformal on benign source-minute maxima""" |
| a = g.alpha / 4 |
| need = ConformalFPR.n_needed(a, 0.5) |
| tf = g.ev_ref.threshold(a) if g.ev_ref.n() >= need else math.inf |
| ts = g.ev_slow_ref.threshold(a) if g.ev_slow_ref.n() >= max(200, need) else math.inf |
| return tf, ts |
|
|
|
|
| def _decayed(se, k, now: float) -> float: |
| st = se.e.get(k) |
| return 0.0 if st is None else st[0] * math.exp(-max(0.0, now - st[1]) / se.tau) |
|
|
|
|
| def _families(log: list, wanted: set, scan: int = 5000) -> dict: |
| """source -> its most frequent named family among the recent alarms ("" when none is named)""" |
| c: dict = {} |
| for e in itertools.islice(reversed(log), scan): |
| s = e.get("src") |
| if s in wanted and _NAMED(e.get("family")): |
| c.setdefault(s, Counter())[e["family"]] += 1 |
| return {s: cc.most_common(1)[0][0] for s, cc in c.items()} |
|
|
|
|
| class ReplicaSync: |
| """Keeps one replica's Layer coherent with the others through a SharedState. sync(now) (or maybe() every |
| `every_s`) does, in order: |
| 1. publish what changed here: new / escalated / extended decisions (with their absolute expiry), operator lifts |
| (an entry gone before its expiry, or a reject / benign_source verdict - with a risk / evidence reset), and |
| allowlist entries; |
| 2. read the merged view and apply it: remote lifts remove what they cancel (and reset that source's evidence on |
| a reject), stronger or later decisions are imported into the Enforcer (so this proxy's fast path enforces |
| them, and the host firewall gets the block), allowlist entries join this replica's feedback allowlist; |
| 3. evidence federation: publish this replica's top-K sources by accumulated evidence and judge every source |
| that has evidence on >= `min_contributors` replicas on its total, against the engine's own threshold; |
| federated alarms go to the alarm log and through the policy like any other verdict; |
| 4. compact this replica's journal now and then. |
| Shared-storage I/O runs outside `lock`; reading and changing the layer runs under it.""" |
|
|
| def __init__(self, layer, shared: SharedState, *, every_s: float = 5.0, lock=None, federate: bool = True, top_k: int = 1000, |
| min_contributors: int = 2, min_share: float = 0.1, stale_s: float = 600.0, compact_every_s: float = 300.0): |
| self.layer, self.shared, self.replica = layer, shared, shared.replica |
| self.every_s, self.lock = float(every_s), lock |
| self.federate, self.top_k, self.min_contributors, self.min_share, self.stale_s = federate, top_k, min_contributors, min_share, stale_s |
| self.compact_every_s = compact_every_s |
| fb = getattr(layer, "feedback", None) |
| self._fb_i = len(fb.records) if fb is not None else 0 |
| self._pub: dict = {} |
| self._imp: dict = {} |
| self._known: dict = {} |
| self._allow_pub: dict = {} |
| self._lift_seen: dict = {} |
| self._fed: dict = {} |
| self._outbox: list = [] |
| self._last_sync = -math.inf |
| self._last_compact: float | None = None |
| self.last: dict = {} |
|
|
| def _publish(self, recs: list) -> None: |
| self._outbox.extend(recs) |
| if len(self._outbox) > 200_000: |
| del self._outbox[: len(self._outbox) - 200_000] |
| if self._outbox: |
| self.shared.publish(self._outbox) |
| self._outbox = [] |
|
|
| def _max_ttl(self) -> float: |
| p = self.layer.policy |
| return max([3600.0] + [p.ttl(a) for a in ("block", "challenge", "rate_limit")]) |
|
|
| def _local_decisions(self, now: float) -> list: |
| out = [] |
| for ip, st in self.layer.enforcer.active.items(): |
| cur = (st["action"], float(st["until"])) |
| if cur[1] <= now: |
| continue |
| imp = self._imp.get(ip) |
| if imp is not None and imp[:2] == cur: |
| continue |
| pub = self._pub.get(ip) |
| if pub is not None and pub[:2] == cur: |
| continue |
| |
| |
| ts = min(now, float(st.get("since", now))) if pub is None and imp is None else now |
| out.append({"k": "d", "ip": ip, "a": cur[0], "u": cur[1], "ts": ts, "r": self.replica, "f": str(st.get("family", "")), |
| "s": int(st.get("score", 0) or 0), "why": str(st.get("reason", ""))[:200]}) |
| self._pub[ip] = (cur[0], cur[1], ts) |
| self._imp.pop(ip, None) |
| return out |
|
|
| def _collect(self, now: float) -> tuple[list, dict | None]: |
| L, recs, lifts = self.layer, [], {} |
| E, fb = L.enforcer, getattr(L, "feedback", None) |
| if fb is not None: |
| for r in fb.records[self._fb_i:]: |
| e = r.get("entity") |
| if r.get("verdict") in ("reject", "benign_source") and e and e != "site:drift": |
| lifts[str(e)] = True |
| self._fb_i = len(fb.records) |
| for e, u in fb.allow.items(): |
| if float(u) > now and self._allow_pub.get(e) != u: |
| recs.append({"k": "a", "e": e, "u": float(u), "ts": now, "r": self.replica}) |
| self._allow_pub[e] = u |
| for ip, (_, u) in self._known.items(): |
| if ip not in E.active and u > now: |
| lifts.setdefault(ip, False) |
| hold = now + self._max_ttl() + 300.0 |
| for ip, reset in lifts.items(): |
| if ip in E.active: |
| continue |
| recs.append({"k": "l", "ip": ip, "ts": now, "r": self.replica, "h": hold, "reset": bool(reset)}) |
| self._pub.pop(ip, None) |
| self._imp.pop(ip, None) |
| self._lift_seen[ip] = now |
| recs += self._local_decisions(now) |
| return recs, (self._summary(now) if self.federate else None) |
|
|
| def _summary(self, now: float) -> dict | None: |
| g = self.layer.guard |
| tf, ts = evidence_thresholds(g) |
| rows: dict = {} |
| for se, thr, j in ((g.evidence, tf, 0), (g.evidence_slow, ts, 1)): |
| if not math.isfinite(thr): |
| continue |
| floor, tau = self.min_share * thr, se.tau |
| top = heapq.nlargest(self.top_k, se.e.items(), key=lambda kv, tau=tau: kv[1][0] * math.exp(-max(0.0, now - kv[1][1]) / tau)) |
| for k, v in top: |
| e = v[0] * math.exp(-max(0.0, now - v[1]) / tau) |
| if e < floor: |
| break |
| rows.setdefault(k, [0.0, 0.0])[j] = e |
| fams = _families(g.log, {str(k) for k in rows}) if rows else {} |
| return {"ts": now, "tau": g.evidence.tau, "tau_slow": g.evidence_slow.tau, |
| "sources": [[_jkey(k), r[0], r[1], fams.get(str(k), "")] for k, r in rows.items()]} |
|
|
| def _decision_ts(self, ip: str, now: float) -> float: |
| if ip in self._pub: |
| return self._pub[ip][2] |
| if ip in self._imp: |
| return self._imp[ip][3] |
| return now |
|
|
| def _apply(self, view: View, others: list, now: float) -> dict: |
| L = self.layer |
| E, g, fb = L.enforcer, L.guard, getattr(L, "feedback", None) |
| st = {"imported": 0, "extended": 0, "lifted": 0, "reset": 0, "allowed": 0, "federated": 0} |
| for ip, lf in view.lifts.items(): |
| lts = float(lf["ts"]) |
| if lts > self._lift_seen.get(ip, -math.inf): |
| if lf.get("reset") and lf.get("r") != self.replica: |
| g.entities.s.pop(ip, None) |
| g.evidence.e.pop(ip, None) |
| st["reset"] += 1 |
| self._lift_seen[ip] = lts |
| cur = E.active.get(ip) |
| if cur is None or self._decision_ts(ip, now) > lts: |
| continue |
| E.active.pop(ip) |
| if E.policy.enforcing and cur["action"] == "block": |
| E.backend.unblock(ip) |
| E._record({"ip": ip, "action": "lift", "family": cur.get("family", ""), "score": cur.get("score", 0), |
| "reason": f"lifted on replica {lf.get('r')}", "taken": E.policy.enforcing}) |
| self._pub.pop(ip, None) |
| self._imp.pop(ip, None) |
| st["lifted"] += 1 |
| for ip, d in view.decisions.items(): |
| a, u = d["a"], float(d["u"]) |
| if a == "monitor" or E.policy.allowed(ip) or (fb is not None and fb.is_allowed(ip, now)): |
| continue |
| cur = E.active.get(ip) |
| if cur is not None: |
| rc, ra = RANK.get(cur["action"], 0), RANK[a] |
| if rc > ra or (rc == ra and float(cur["until"]) >= u): |
| continue |
| if rc == ra: |
| cur["until"] = u |
| self._imp[ip] = (a, u, d.get("r"), float(d["ts"])) |
| st["extended"] += 1 |
| continue |
| taken = False |
| if E.policy.enforcing and a == "block" and (cur is None or cur["action"] != "block"): |
| taken = bool(E.backend.block(ip, max(1.0, u - now))) |
| E.active[ip] = {"action": a, "until": u, "family": d.get("f") or "replica", "score": int(d.get("s", 0) or 0), |
| "reason": f"replica {d.get('r')}: {d.get('why', '')}"[:300], "since": now, "replica": d.get("r")} |
| E._record({"ip": ip, "action": a, "family": d.get("f") or "replica", "score": int(d.get("s", 0) or 0), |
| "reason": f"imported from replica {d.get('r')}", "taken": taken, "ttl_s": round(u - now, 1)}) |
| self._imp[ip] = (a, u, d.get("r"), float(d["ts"])) |
| st["imported"] += 1 |
| if fb is not None: |
| for e, u in view.allows.items(): |
| if fb.allow.get(e, 0) < u: |
| fb.allow[e] = u |
| self._allow_pub[e] = u |
| st["allowed"] += 1 |
| if self.federate: |
| st["federated"] = len(self._federate(others, now)) |
| return st |
|
|
| def _federate(self, others: list, now: float) -> list: |
| L = self.layer |
| g, E, fb = L.guard, L.enforcer, getattr(L, "feedback", None) |
| tf, ts = evidence_thresholds(g) |
| if not (math.isfinite(tf) or math.isfinite(ts)): |
| return [] |
| tot: dict = {} |
| for s in others: |
| age = max(0.0, now - float(s.get("ts", now))) |
| if age > self.stale_s: |
| continue |
| kf, ks = math.exp(-age / g.evidence.tau), math.exp(-age / g.evidence_slow.tau) |
| for row in s.get("sources", []): |
| k = _ukey(row[0]) |
| a = tot.setdefault(k, [0.0, 0.0, set(), Counter()]) |
| a[0] += float(row[1]) * kf |
| a[1] += float(row[2]) * ks |
| a[2].add(s.get("r")) |
| if _NAMED(row[3]): |
| a[3][row[3]] += 1 |
| acted = [] |
| local_fams = _families(g.log, {str(k) for k in tot}) if tot else {} |
| for k, (ef, es, reps, fams) in tot.items(): |
| lf, ls = _decayed(g.evidence, k, now), _decayed(g.evidence_slow, k, now) |
| if lf > 0 or ls > 0: |
| reps = reps | {self.replica} |
| if local_fams.get(str(k)): |
| fams[local_fams[str(k)]] += 1 |
| if len(reps) < self.min_contributors: |
| continue |
| F, S = ef + lf, es + ls |
| if not (F > tf or S > ts): |
| continue |
| src = str(k) |
| if fb is not None and fb.is_allowed(src, now): |
| continue |
| fam = fams.most_common(1)[0][0] if fams else "attack_generic" |
| pv = min(float(g.ev_ref.pvalues([F])[0]) if math.isfinite(tf) else 1.0, float(g.ev_slow_ref.pvalues([S])[0]) if math.isfinite(ts) else 1.0) |
| score = event_score(1.0 - pv, severity(fam)) |
| action = L.policy.action_for(fam, score, src) |
| cur = E.active.get(src) |
| if cur is not None and RANK.get(cur["action"], 0) >= RANK.get(action, 0): |
| continue |
| if now - self._fed.get(src, -math.inf) < 60.0 and action == "monitor": |
| continue |
| self._fed[src] = now |
| why = (f"federated evidence over replicas {', '.join(sorted(map(str, reps)))}: {F:.1f} (fast) / {S:.1f} (slow) " |
| f"> threshold {tf:.1f} / {ts:.1f}") |
| g._log("federated", now, src, "*", 0, fam, int(score), 4, evidence=why, soft=not _NAMED(fam)) |
| acted.append(E.apply(src, action, fam, int(score), reason=why, now=now)) |
| return acted |
|
|
| def sync(self, now: float | None = None) -> dict: |
| now = time.time() if now is None else float(now) |
| if self._last_compact is None: |
| self._last_compact = now |
| with _lock(self.lock): |
| recs, summary = self._collect(now) |
| self._publish(recs) |
| if summary is not None: |
| self.shared.publish_evidence(summary) |
| self.shared.heartbeat() |
| view = self.shared.view(now) |
| others = self.shared.evidence() if self.federate else [] |
| with _lock(self.lock): |
| st = self._apply(view, others, now) |
| recs2 = self._local_decisions(now) |
| self._known = {ip: (s["action"], float(s["until"])) for ip, s in self.layer.enforcer.active.items()} |
| self._publish(recs2) |
| if now - self._last_compact >= self.compact_every_s: |
| self.shared.compact(now) |
| self._last_compact = now |
| self._last_sync = time.monotonic() |
| st.update(published=len(recs) + len(recs2), cluster_decisions=len(view.decisions)) |
| self.last = st |
| return st |
|
|
| def maybe(self, now: float | None = None) -> dict | None: |
| if time.monotonic() - self._last_sync < self.every_s: |
| return None |
| return self.sync(now) |
|
|
|
|
| |
| class Housekeeper: |
| """One daemon thread for the sidecar: replica sync, periodic snapshots and (optionally) Layer.tick, every |
| `period_s`. stop() writes a final snapshot (a clean shutdown loses nothing).""" |
|
|
| def __init__(self, layer, lock=None, snapshotter: Snapshotter | None = None, sync: ReplicaSync | None = None, tick: bool = False, |
| period_s: float = 1.0): |
| self.layer, self.lock, self.snap, self.sync, self.tick, self.period_s = layer, lock, snapshotter, sync, tick, period_s |
| self._stop = threading.Event() |
| self._t: threading.Thread | None = None |
| self.errors: list = [] |
| self.restored: dict | None = None |
|
|
| def step(self, now: float | None = None) -> dict: |
| out = {} |
|
|
| def tick(): |
| with _lock(self.lock): |
| return self.layer.tick(now) |
| for name, fn in (("sync", lambda: self.sync.maybe(now) if self.sync else None), |
| ("snapshot", lambda: self.snap.maybe() if self.snap else None), |
| ("tick", tick if self.tick else lambda: None)): |
| try: |
| r = fn() |
| if r is not None: |
| out[name] = r |
| except Exception as e: |
| self.errors.append(f"{time.time():.0f} {name}: {type(e).__name__}: {e}") |
| del self.errors[:-50] |
| return out |
|
|
| def _run(self) -> None: |
| while not self._stop.wait(self.period_s): |
| self.step() |
|
|
| def start(self) -> "Housekeeper": |
| if self._t is None: |
| self._t = threading.Thread(target=self._run, name="dibaguard-ha", daemon=True) |
| self._t.start() |
| return self |
|
|
| def stop(self, final_snapshot: bool = True) -> dict | None: |
| self._stop.set() |
| if self._t is not None: |
| self._t.join(timeout=30) |
| return self.snap.snapshot_now() if (final_snapshot and self.snap) else None |
|
|
|
|
| def from_env(layer, lock=None, env=None, start: bool = True) -> Housekeeper | None: |
| """HA from environment variables (None when none is set): |
| DIBAGUARD_SNAPSHOT_DIR snapshots here (restore on start, then every DIBAGUARD_SNAPSHOT_EVERY_S, default 300; |
| keep DIBAGUARD_SNAPSHOT_KEEP, default 5; DIBAGUARD_RESTORE=0 to start empty) |
| DIBAGUARD_SNAPSHOT_KEY HMAC key for snapshots (else sha256 only) |
| DIBAGUARD_SHARED_DIR FileSharedState root (active-active), or DIBAGUARD_REDIS_URL for RedisSharedState |
| DIBAGUARD_REPLICA this replica's id (default: host name); DIBAGUARD_SYNC_EVERY_S (default 5) |
| DIBAGUARD_TICK=1 also lift expired actions from the background thread""" |
| env = os.environ if env is None else env |
| snap_dir, shared_dir, redis_url = env.get("DIBAGUARD_SNAPSHOT_DIR"), env.get("DIBAGUARD_SHARED_DIR"), env.get("DIBAGUARD_REDIS_URL") |
| if not (snap_dir or shared_dir or redis_url): |
| return None |
| replica = _safe_name(env.get("DIBAGUARD_REPLICA") or default_replica()) |
| snap = Snapshotter(layer, snap_dir, every_s=float(env.get("DIBAGUARD_SNAPSHOT_EVERY_S", 300)), keep=int(env.get("DIBAGUARD_SNAPSHOT_KEEP", 5)), |
| lock=lock, replica=replica) if snap_dir else None |
| restored = snap.restore_on_start() if snap is not None and env.get("DIBAGUARD_RESTORE", "1") != "0" else None |
| shared = RedisSharedState(redis_url, replica=replica) if redis_url else FileSharedState(shared_dir, replica=replica) if shared_dir else None |
| sync = ReplicaSync(layer, shared, every_s=float(env.get("DIBAGUARD_SYNC_EVERY_S", 5)), lock=lock) if shared is not None else None |
| hk = Housekeeper(layer, lock, snap, sync, tick=env.get("DIBAGUARD_TICK", "0") == "1") |
| hk.restored = restored |
| return hk.start() if start else hk |
|
|
