File size: 4,311 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 | """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(): # family head only for alarms
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)) # below threshold: never reported as a likely attack
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()}
|