| """FlowGuard: tier-1 flow scoring + the decision layer, numpy only. |
| |
| g = FlowGuard.load("models/flows_v1", view="nf", variant="robust") |
| g.calibrate(benign_flows) # optional: this site's own normal traffic (no attack labels needed) |
| out = g.score(flows, entity=src_ips, t=timestamps) |
| |
| score() returns per flow: p_attack (temperature-calibrated), alarm (conformal threshold: this site's own when it has |
| enough benign scores for the requested alpha, else the shipped one), family + severity (family head runs on alarms |
| only), event score 0-100, entity risk (per source, decaying) and the site's attack pressure. |
| """ |
| from __future__ import annotations |
|
|
| import os |
|
|
| import numpy as np |
|
|
| from dibaguard.decide import ConformalFPR, EntityRisk, SitePressure, event_score |
| from dibaguard.families import severity |
| from dibaguard.flows import schema |
| from dibaguard.flows.trees import Oblivious |
|
|
|
|
| class FlowGuard: |
| def __init__(self, detector: Oblivious, family: Oblivious | None = None, alpha: float = 1e-3, |
| entity_half_life_s: float = 600.0, pressure_window_s: float = 300.0): |
| self.det, self.fam, self.alpha = detector, family, alpha |
| self.view, self.robust = detector.meta["view"], detector.meta["variant"] == "robust" |
| self.temperature = float(detector.meta.get("temperature", 1.0)) |
| self.shipped = {float(k): float(v) for k, v in detector.meta.get("thresholds", {}).items()} |
| self.site = ConformalFPR(max_per_group=200_000, min_n=1) |
| self.entities = EntityRisk(half_life_s=entity_half_life_s) |
| self.pressure = SitePressure(window_s=pressure_window_s) |
|
|
| @classmethod |
| def load(cls, model_dir: str, view: str = "nf", variant: str = "robust", **kw) -> "FlowGuard": |
| det = Oblivious.load(os.path.join(model_dir, f"{view}_{variant}_detector.npz")) |
| fp = os.path.join(model_dir, f"{view}_{variant}_family.npz") |
| return cls(det, Oblivious.load(fp) if os.path.exists(fp) else None, **kw) |
|
|
| def features(self, flows) -> np.ndarray: |
| return schema.matrix(self.view, flows, robust=self.robust) |
|
|
| def calibrate(self, benign_flows) -> int: |
| """Add this site's benign traffic; returns how many scores the site threshold now rests on.""" |
| self.site.add("site", self.det.raw(self.features(benign_flows))[:, 0]) |
| return self.site.n("site") |
|
|
| def threshold(self) -> tuple[float, str]: |
| if self.site.n("site") >= ConformalFPR.n_needed(self.alpha, 0.5): |
| return self.site.threshold("site", self.alpha, fallback=None), "site" |
| if self.alpha in self.shipped: |
| return self.shipped[self.alpha], "shipped" |
| return float(min(self.shipped.values(), default=0.0)), "shipped-nearest" |
|
|
| def score(self, flows, entity=None, t=None) -> dict: |
| x = self.features(flows) |
| z = self.det.raw(x)[:, 0] |
| p = 1.0 / (1.0 + np.exp(-z / self.temperature)) |
| thr, src = self.threshold() |
| alarm = z > thr |
| fam = np.full(len(x), "benign", dtype=object) |
| fam[alarm] = "attack_generic" |
| if self.fam is not None and alarm.any(): |
| fam[alarm] = np.asarray(self.fam.classes, dtype=object)[self.fam.raw(x[alarm]).argmax(1)] |
| sev = np.array([severity(f) if a else 0.0 for f, a in zip(fam, alarm)]) |
| pe = np.where(alarm, p, np.minimum(p, 0.5)) |
| score = np.array([event_score(float(a), float(b)) for a, b in zip(pe, sev)], np.int16) |
| risk = None |
| if entity is not None: |
| tt = np.zeros(len(x)) if t is None else np.asarray(t, np.float64) |
| risk = np.array([self.entities.update(str(e), float(ti), float(pi), float(si), bool(al)) |
| for e, ti, pi, si, al in zip(entity, tt, pe, sev, alarm)]) |
| for ti, pi, fi, al in zip(tt, pe, fam, alarm): |
| self.pressure.update(float(ti), float(pi) if al else 0.0, fi if al else None) |
| return {"p_attack": p, "alarm": alarm, "family": fam, "severity": sev, "score": score, |
| "entity_risk": risk, "threshold": thr, "threshold_source": src, |
| "site_pressure": self.pressure.pressure, "pressure_by_family": self.pressure.by_family()} |
|
|
