Download runtime/python/dibaguard/flows/entity.py from Dibachain/DibaGuard: direct link, hf CLI and curl.
- Browser
- Download file 14.1 kB
-
https://huggingface.co/Dibachain/DibaGuard/resolve/main/runtime/python/dibaguard/flows/entity.py
- Command line
-
hf download hf://Dibachain/DibaGuard/runtime/python/dibaguard/flows/entity.py
-
curl -L -o entity.py https://huggingface.co/Dibachain/DibaGuard/resolve/main/runtime/python/dibaguard/flows/entity.py
14.1 kB
| """Per-entity stream state (tier 0): what one flow cannot show, the source's recent behaviour can. | |
| For every source (and every destination) a small record of exponentially decayed counters is kept - flows, | |
| bytes, packets; share of flows that got no reply, that were SYN-only, that were reset, that were very short; | |
| destination fan-out (new port / new host within memory); repetition of the same (host, port); time since the | |
| last flow. Every counter decays with half-life H, so "flows" reads as events in the last ~1.4 H seconds and all | |
| shares are scale-free - they transfer between networks far better than absolute per-flow values. | |
| Three horizons (default 10 s / 2 min / 30 min) see bursts, sessions and slow scans. Per flow the cost is O(1): | |
| two dict lookups and ~60 multiply-adds. The state is bounded (LRU over `max_entities`), so a load-balancer | |
| replica holds a few hundred MB at most. featurize() runs the same update over a time-sorted table (training), so | |
| training features and runtime features are the same numbers. | |
| """ | |
| from __future__ import annotations | |
| import math | |
| from collections import OrderedDict | |
| import numpy as np | |
| from dibaguard.decide import evict | |
| HORIZONS_S = (10.0, 120.0, 1800.0) | |
| # authentication services: where repeated connections mean password guessing (also used by the campaign tier) | |
| AUTH_PORTS = [21, 22, 23, 25, 110, 143, 389, 445, 465, 587, 993, 995, 1433, 1521, 3306, 3389, 5432, 5900, 6379, 9200, 27017] | |
| # where a fast stream of repeated sessions *without* failure signs is still password guessing (directory and file services - | |
| # LDAP, SMB - are chatty by nature: there, only failures make it brute force) | |
| GUESSING_PORTS = [p for p in AUTH_PORTS if p not in (389, 445)] | |
| PER_H = ["flows", "bytes_in", "pkts_in", "noreply", "synonly", "rst", "newport", "newhost", "repeat", "short"] | |
| K = len(PER_H) | |
| TAIL = ["gap_s", "gap_ewma_s", "dst_flows_h1", "dst_noreply_h1", "dst_newsrc_h1", "dst_flows_h2"] | |
| MEMORY = 512 # (host, port) keys remembered per source; sources remembered per destination | |
| # attack-direction features for the site-only anomaly tier: more of any of these is more suspicious (upper tail only); | |
| # sizes and gaps are left out - benign bulk transfers and idle periods must not raise alarms; destination-side | |
| # features are left out too - a surge of new clients on one host would flag every one of them (that signal belongs | |
| # to the site pressure, not to a per-flow verdict) | |
| ANOMALY = [f"src_{n}_h{i + 1}" for i in range(len(HORIZONS_S)) | |
| for n in ("flows", "noreply", "synonly", "rst", "newport", "newhost", "repeat", "short")] | |
| # the site-only anomaly tier leaves 'repeat' out: a benign client's repeat share climbs for hours as it revisits the | |
| # same endpoints (it is an attack signal only together with failures, which the family rules check) | |
| # the site-only anomaly tier judges the *shape* of a source's behaviour, not its volume: pure rate (src_flows_*) is | |
| # left out - a busy legitimate client (an API integration, a heavy user, a busier hour) is not an attack by itself; | |
| # volume attacks are the monotone prior's, the destination-flood tier's and the global tiers' job | |
| ANOMALY_SITE = [f for f in ANOMALY if "repeat" not in f and "_flows_" not in f] | |
| SHORT_MS = 10.0 # a flow shorter than this with <= 2 packets counts as 'short' (probe / failed connect) | |
| def feature_names(horizons=HORIZONS_S) -> list[str]: | |
| out = [] | |
| for i in range(len(horizons)): | |
| out += [f"src_{n}_h{i + 1}" for n in PER_H] | |
| return out + TAIL | |
| class _Rec: | |
| __slots__ = ("t", "c", "seen", "hosts", "gap_ewma", "n") | |
| def __init__(self, t: float, n: int): | |
| self.t, self.c, self.seen, self.hosts, self.gap_ewma, self.n = t, [0.0] * n, OrderedDict(), {}, -1.0, 0 | |
| class EntityState: | |
| def __init__(self, horizons=HORIZONS_S, max_entities: int = 2_000_000, gap_alpha: float = 0.2): | |
| self.tau = [h / math.log(2) for h in horizons] | |
| self.nh, self.max, self.ga = len(horizons), max_entities, gap_alpha | |
| self.src: "OrderedDict[object, _Rec]" = OrderedDict() | |
| self.dst: "OrderedDict[object, _Rec]" = OrderedDict() | |
| def _get(self, table, key, t) -> tuple[_Rec, float, float]: | |
| """-> (record decayed to t, gap since its last event (-1 when new), lateness). A late record (older than the | |
| record's clock: a long flow exported at its end, another exporter's skewed clock) never rewinds the clock - | |
| that would make the next event decay everything accumulated - and is counted with the weight it has at the | |
| record's clock, exp(-lateness / tau): the exact update of a decayed counter for an out-of-order event""" | |
| r = table.get(key) | |
| if r is None: | |
| r = _Rec(t, self.nh * K) | |
| table[key] = r | |
| if len(table) > self.max: | |
| # a source with a history (8+ flows) is passed over while fresh one-flow sources can go instead | |
| evict(table, lambda v: v.n >= 8) | |
| return r, -1.0, 0.0 | |
| table.move_to_end(key) | |
| dt = t - r.t | |
| if dt > 0: | |
| c = r.c | |
| for i, tau in enumerate(self.tau): | |
| f = math.exp(-dt / tau) | |
| for j in range(i * K, i * K + K): | |
| c[j] *= f | |
| r.t = t | |
| return r, dt, 0.0 | |
| return r, 0.0, -dt | |
| def _remember(rec: _Rec, key, host) -> bool: | |
| """Remember key; returns whether `host` was already known (before this flow).""" | |
| known = host in rec.hosts | |
| s = rec.seen | |
| if key in s: | |
| s.move_to_end(key) | |
| else: | |
| s[key] = host | |
| rec.hosts[host] = rec.hosts.get(host, 0) + 1 | |
| if len(s) > MEMORY: | |
| _, h = s.popitem(last=False) | |
| n = rec.hosts[h] - 1 | |
| if n: | |
| rec.hosts[h] = n | |
| else: | |
| del rec.hosts[h] | |
| return known | |
| def update(self, t: float, src, dst, dport: int, bytes_in: float, pkts_in: float, pkts_out: float, | |
| syn_only: bool, rst: bool, short: bool, out: list) -> None: | |
| """Append this flow's entity features (len(feature_names())) to `out`, then absorb the flow.""" | |
| r, gap, late = self._get(self.src, src, t) | |
| d, _, dlate = self._get(self.dst, dst, t) | |
| r.n += 1 # true age of the source (flows ever seen), not a decayed count | |
| key = (dst, dport) | |
| seen = key in r.seen | |
| new_host = dst not in r.hosts | |
| new_src = src not in d.seen | |
| noreply = pkts_out <= 0 | |
| ev = (1.0, bytes_in, pkts_in, float(noreply), float(syn_only), float(rst), float(not seen), float(new_host), | |
| float(seen), float(short)) | |
| c = r.c | |
| if late > 0 or dlate > 0: # out of order: each horizon counts it at its decayed weight | |
| self._update_late(r, d, ev, noreply, new_src, gap, late, dlate, out) | |
| else: | |
| for i in range(self.nh): | |
| b = i * K | |
| n = c[b] + 1.0 | |
| out.append(n) | |
| out.append(c[b + 1] + bytes_in) | |
| out.append(c[b + 2] + pkts_in) | |
| for j in range(3, K): | |
| out.append((c[b + j] + ev[j]) / n) | |
| out.append(gap) | |
| out.append(r.gap_ewma) | |
| dc = d.c | |
| n1 = dc[0] + 1.0 | |
| out.append(n1) | |
| out.append((dc[3] + noreply) / n1) | |
| out.append((dc[6] + new_src) / n1) | |
| out.append(dc[K] + 1.0) | |
| for i in range(self.nh): | |
| b = i * K | |
| for j in range(K): | |
| c[b + j] += ev[j] | |
| dc[b] += 1.0 | |
| dc[b + 3] += noreply | |
| dc[b + 6] += new_src | |
| self._remember(r, key, dst) | |
| self._remember(d, src, src) | |
| if gap >= 0: | |
| r.gap_ewma = gap if r.gap_ewma < 0 else (1 - self.ga) * r.gap_ewma + self.ga * gap | |
| def _update_late(self, r: _Rec, d: _Rec, ev, noreply, new_src, gap, late, dlate, out: list) -> None: | |
| w = [max(math.exp(-late / tau), 1e-9) for tau in self.tau] | |
| dw = [max(math.exp(-dlate / tau), 1e-9) for tau in self.tau] | |
| c, dc = r.c, d.c | |
| for i in range(self.nh): | |
| b, wi = i * K, w[i] | |
| n = c[b] + wi | |
| out.append(n) | |
| out.append(c[b + 1] + wi * ev[1]) | |
| out.append(c[b + 2] + wi * ev[2]) | |
| for j in range(3, K): | |
| out.append((c[b + j] + wi * ev[j]) / n) | |
| out.append(gap) | |
| out.append(r.gap_ewma) | |
| n1 = dc[0] + dw[0] | |
| out.append(n1) | |
| out.append((dc[3] + dw[0] * noreply) / n1) | |
| out.append((dc[6] + dw[0] * new_src) / n1) | |
| out.append(dc[K] + dw[1]) | |
| for i in range(self.nh): | |
| b = i * K | |
| for j in range(K): | |
| c[b + j] += w[i] * ev[j] | |
| dc[b] += dw[i] | |
| dc[b + 3] += dw[i] * noreply | |
| dc[b + 6] += dw[i] * new_src | |
| def featurize(self, t, src, dst, dport, bytes_in, pkts_in, pkts_out, syn_only, rst, duration_ms) -> np.ndarray: | |
| """Column arrays, already sorted by t -> float32 [n, len(feature_names())]. Runs the stream update.""" | |
| n = len(t) | |
| short = (np.asarray(duration_ms, np.float64) < SHORT_MS) & (np.asarray(pkts_in, np.float64) <= 2) | |
| cols = [np.asarray(a).tolist() for a in (t, src, dst, dport, bytes_in, pkts_in, pkts_out, syn_only, rst, short)] | |
| out: list = [] | |
| up = self.update | |
| for row in zip(*cols): | |
| up(*row, out) | |
| return np.asarray(out, np.float32).reshape(n, -1) | |
| def tcp_syn_only(client_flags, server_flags) -> np.ndarray: | |
| """NetFlow TCP_FLAGS bitmask: client sent SYN (0x02) and no ACK (0x10), server nothing.""" | |
| cf = np.asarray(client_flags, np.int64) | |
| return ((cf & 0x02) != 0) & ((cf & 0x10) == 0) & (np.asarray(server_flags, np.int64) == 0) | |
| def tcp_rst(flags) -> np.ndarray: | |
| return (np.asarray(flags, np.int64) & 0x04) != 0 | |
| WELL_KNOWN = (20, 21, 22, 23, 25, 53, 67, 80, 110, 123, 135, 137, 139, 143, 161, 389, 443, 445, 465, 587, 993, 995, 1433, 1883, | |
| 3306, 3389, 5432, 5900, 6379, 8080, 8443, 8883) | |
| _WK = {p: i + 1 for i, p in enumerate(WELL_KNOWN)} | |
| def service_key(dport, proto) -> np.ndarray: | |
| """Mondrian calibration group of a flow: protocol x destination service (well-known port exact, else 'other'). | |
| A site's thresholds are kept per service, so a new or bursty service cannot move the others' thresholds.""" | |
| dp = np.asarray(dport, np.int64) | |
| pr = np.asarray(proto, np.int64) | |
| svc = np.array([_WK.get(int(p), 0) if p >= 0 else 0 for p in dp], np.int64) | |
| svc = np.where((dp >= 1024), len(WELL_KNOWN) + 1, svc) | |
| return (np.clip(pr, 0, 255) * (len(WELL_KNOWN) + 2) + svc).astype(np.int32) | |
| _F = {n: i for i, n in enumerate(feature_names())} | |
| def family_rules(x, port=None) -> np.ndarray: | |
| """Transparent attack family from the entity stream features of alarmed flows ([n, 36] -> object array). | |
| Ranked by what the counters show over the 2-minute horizon: port walking, repeated failed connects to one | |
| service, floods at one target, host walking; 'attack_generic' when none of the signatures is clear.""" | |
| x = np.asarray(x, np.float64) | |
| h = "h2" | |
| flows, nop, syn, newp, newh, rep, short = (x[:, _F[f"src_{k}_{h}"]] for k in ("flows", "noreply", "synonly", "newport", "newhost", "repeat", "short")) | |
| dflows, dnew = x[:, _F["dst_flows_h1"]], x[:, _F["dst_newsrc_h1"]] | |
| rate10 = x[:, _F["src_flows_h1"]] # flows in the last ~10 s | |
| out = np.full(len(x), "attack_generic", dtype=object) | |
| # the same (host, port) again and again - *and* failing (no reply / cut short / reset) or at an attacker's rate: | |
| # steady successful polling (telemetry, Kerberos / LDAP renewals, API clients, NTP) repeats too, and is not brute force | |
| rst = x[:, _F[f"src_rst_{h}"]] | |
| # (a fast repeater without failures is brute force only on an authentication service - many connections to one | |
| # proxy, API or broker are a busy client) | |
| auth = np.isin(np.asarray(port).astype(np.int64), GUESSING_PORTS) if port is not None else np.ones(len(x), bool) | |
| # ('short' - < 10 ms, <= 2 packets - is not a failure sign here: every DNS / NTP exchange over UDP is that short, and a | |
| # guessed password completes its TCP handshake; refused or unanswered connects show as reset / no reply) | |
| out[(rep > 0.6) & (flows >= 10) & ((nop > 0.2) | (rst > 0.2) | ((flows >= 30) & auth))] = "bruteforce" | |
| out[(newh > 0.3) & (flows >= 10) & (newp <= 0.5) & ((nop > 0.5) | (short > 0.5))] = "recon" | |
| # host walking on the services lateral movement uses (SMB, RPC, WinRM, SSH, RDP): one host reaching many others | |
| # where it has no business - answered or not - is a compromised host spreading, not a survey | |
| if port is not None: | |
| lat = np.isin(np.asarray(port).astype(np.int64), [445, 135, 139, 5985, 5986, 22, 3389]) | |
| out[lat & (newh > 0.3) & (flows >= 10) & (newp <= 0.5)] = "lateral" | |
| # (a scan's probes go unanswered, half-open or refused; 'short' alone is every UDP DNS / NTP exchange - a resolver | |
| # asking many authoritative servers is not port walking) | |
| out[(newp > 0.5) & (flows >= 10) & ((syn > 0.5) | (nop > 0.5) | (rst > 0.5))] = "portscan" | |
| # a slow scan walks ports just the same over the 30-minute horizon (a probe every half minute never makes 10 flows | |
| # in 2 minutes): named when nothing else on the short horizon names it (playbook A10) | |
| newp3, flows3 = x[:, _F["src_newport_h3"]], x[:, _F["src_flows_h3"]] | |
| fail3 = (x[:, _F["src_synonly_h3"]] > 0.5) | (x[:, _F["src_noreply_h3"]] > 0.5) | (x[:, _F["src_short_h3"]] > 0.5) | |
| out[(out == "attack_generic") & (newp3 > 0.5) & (flows3 >= 10) & fail3] = "portscan" | |
| flood = (rate10 >= 100) & (newp <= 0.2) & (newh <= 0.2) # >= 10 flows/s at one service for the whole horizon | |
| out[flood] = "dos" | |
| out[flood & (dnew > 0.5) & (dflows >= 500)] = "ddos" | |
| return out | |