DibaAi's picture
DibaGuard: public licence-clean release (Dibachain)
5d97275 verified
Raw History Blame Contribute Delete
9.23 kB
"""Tier-2 sequence network: one source's stream of flows -> per-flow attack logit + family, in numpy.
What one flow cannot show and even decayed counters only summarise, the *order* of a source's recent flows does:
a scan walks ports, brute force repeats one (host, port) with short failed connects, a DDoS bot hammers one target
with no replies, a beacon returns at a fixed gap. The network reads the stream event by event (a causal GRU, so
the runtime keeps one hidden vector per source and pays one step per flow) and is conditioned on a *site context*:
a vector pooled from a sample of the site's own benign stream. Training sees many networks with their own
contexts, so at a new site the model gets its benign reference the same way - no attack labels, no retraining.
Event features are view-independent (nf and lu alike): scale-free per-flow numbers (logs of sizes, packets,
duration, bytes per packet, no-reply, TCP flag bits) plus the entity stream features of dibaguard.flows.entity
(logs of counts, shares as they are). Ports go through a small vocabulary (exact for 0..1023, hashed above) and
protocol through a 5-way id. The torch model lives in guard_jobs/seq_train.py and must stay in step with
`event_features`, `SeqNet.forward` here (parity is checked at export).
"""
from __future__ import annotations
import json
import numpy as np
from dibaguard.flows import entity as ent
PORT_VOCAB = 1024 + 1 + 255 # 0 = unknown, 1..1024 = port 0..1023, then hashed buckets
PROTO_VOCAB = 5 # 0 unknown, 1 tcp, 2 udp, 3 icmp, 4 other
FLAG_BITS = (0x01, 0x02, 0x04, 0x08, 0x10, 0x20) # FIN SYN RST PSH ACK URG
NUMERIC = (["log_bytes_in", "log_bytes_out", "log_pkts_in", "log_pkts_out", "log_duration_ms", "log_bpp_in",
"log_bpp_out", "pkt_ratio", "noreply"]
+ [f"{k}_{b}" for k in ("flags", "cflags", "sflags") for b in ("fin", "syn", "rst", "psh", "ack", "urg")]
+ ["flags_known"] + [f"ent_{n}" for n in ent.feature_names()] + ["gap_known"])
_LOG_ENT = {n for n in ent.feature_names() if any(n.endswith(s) for s in ("flows_h1", "flows_h2", "flows_h3", "bytes_in_h1",
"bytes_in_h2", "bytes_in_h3", "pkts_in_h1", "pkts_in_h2", "pkts_in_h3", "gap_s", "gap_ewma_s", "dst_flows_h1",
"dst_flows_h2"))}
def port_id(dport) -> np.ndarray:
p = np.asarray(dport, np.int64)
out = np.where((p >= 0) & (p < 1024), p + 1, 1025 + (p - 1024) % 255)
return np.where(p < 0, 0, out).astype(np.int32)
def proto_id(proto) -> np.ndarray:
p = np.asarray(proto, np.int64)
return np.select([p == 6, p == 17, p == 1, p >= 0], [1, 2, 3, 4], 0).astype(np.int32)
def _bits(v) -> np.ndarray:
v = np.asarray(v, np.int64)
return np.stack([(v & b) != 0 for b in FLAG_BITS], 1).astype(np.float32)
def event_features(cols, view: str) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""cols: mapping column name -> array (an entity_v1 frame) -> (numeric float32 [n, len(NUMERIC)], port ids, proto ids)."""
if view == "lu":
bi, bo, pi, po = (np.asarray(cols[c], np.float64) for c in ("bytes_in", "bytes_out", "num_pkts_in", "num_pkts_out"))
dur = np.asarray(cols["duration"], np.float64) * 1000.0
proto = cols["proto"] if "proto" in cols else np.full(len(bi), -1)
dport = np.full(len(bi), -1)
fl = cf = sf = np.zeros(len(bi), np.int64)
known = 0.0
else:
bi, bo, pi, po, dur = (np.asarray(cols[c], np.float64) for c in ("IN_BYTES", "OUT_BYTES", "IN_PKTS", "OUT_PKTS",
"FLOW_DURATION_MILLISECONDS"))
proto, dport = cols["PROTOCOL"], cols["L4_DST_PORT"]
fl, cf, sf = (np.asarray(cols[c], np.float64).astype(np.int64) for c in ("TCP_FLAGS", "CLIENT_TCP_FLAGS", "SERVER_TCP_FLAGS"))
known = 1.0
bi, bo, pi, po, dur = (np.maximum(np.nan_to_num(v, nan=0.0), 0.0) for v in (bi, bo, pi, po, dur))
num = [np.log1p(bi), np.log1p(bo), np.log1p(pi), np.log1p(po), np.log1p(dur), np.log1p(bi / np.maximum(pi, 1)),
np.log1p(bo / np.maximum(po, 1)), np.minimum(po / np.maximum(pi, 1), 8.0), (po <= 0).astype(np.float64)]
x = [np.stack(num, 1).astype(np.float32), _bits(fl), _bits(cf), _bits(sf), np.full((len(bi), 1), known, np.float32)]
e = np.nan_to_num(np.column_stack([np.asarray(cols[n], np.float64) for n in ent.feature_names()]), nan=-1.0)
names = ent.feature_names()
for j, n in enumerate(names):
if n in _LOG_ENT:
e[:, j] = np.log1p(np.maximum(e[:, j], 0.0))
gap_known = (np.asarray(cols["gap_s"], np.float64) >= 0).astype(np.float32)[:, None]
x += [e.astype(np.float32), gap_known]
return np.hstack(x), port_id(np.nan_to_num(np.asarray(dport, np.float64), nan=-1)), proto_id(np.nan_to_num(np.asarray(proto, np.float64), nan=-1))
def _sigmoid(x):
return 1.0 / (1.0 + np.exp(-x))
class SeqNet:
"""numpy forward of the torch SeqNet (guard_jobs.seq_train). Weights: dict name -> float32 array."""
def __init__(self, w: dict, info: dict):
self.w, self.info = w, info
self.d = int(w["gru.l0.w_hh"].shape[1])
self.layers = sum(1 for k in w if k.startswith("gru.l") and k.endswith(".w_hh"))
self.families = info.get("families", [])
self.ctx = "ctx.w" in w
# -- pieces
def embed(self, x_num: np.ndarray, port: np.ndarray, proto: np.ndarray) -> np.ndarray:
"""[n, NUMERIC] + ids -> encoder output [n, d]."""
w = self.w
z = np.concatenate([x_num, w["emb_port"][port], w["emb_proto"][proto]], 1)
h = z @ w["enc.0.w"].T + w["enc.0.b"]
h = 0.5 * h * (1.0 + np.tanh(0.7978845608 * (h + 0.044715 * h ** 3))) # gelu (tanh approximation)
return h @ w["enc.1.w"].T + w["enc.1.b"]
def gru_step(self, x: np.ndarray, state: np.ndarray) -> np.ndarray:
"""x [n, d], state [layers, n, d] -> new state (modified copy)."""
w, d = self.w, self.d
new = np.empty_like(state)
inp = x
for layer in range(self.layers):
h = state[layer]
gi = inp @ w[f"gru.l{layer}.w_ih"].T + w[f"gru.l{layer}.b_ih"]
gh = h @ w[f"gru.l{layer}.w_hh"].T + w[f"gru.l{layer}.b_hh"]
r = _sigmoid(gi[:, :d] + gh[:, :d])
zg = _sigmoid(gi[:, d:2 * d] + gh[:, d:2 * d])
n = np.tanh(gi[:, 2 * d:] + r * gh[:, 2 * d:])
inp = (1.0 - zg) * n + zg * h
new[layer] = inp
return new
def context(self, h_pool: np.ndarray | None) -> tuple[np.ndarray, np.ndarray]:
"""Mean of final hidden states over benign reference windows [d] -> (gamma, beta) [d] each."""
if not self.ctx or h_pool is None:
return np.zeros(self.d, np.float32), np.zeros(self.d, np.float32)
g = np.tanh(h_pool @ self.w["ctx.w"].T + self.w["ctx.b"])
return g[: self.d], g[self.d:]
def head(self, h: np.ndarray, gamma, beta) -> np.ndarray:
"""top-layer hidden [n, d] -> logits [n, 1 + n_families]."""
w = self.w
h = h * (1.0 + gamma) + beta
z = h @ w["head.0.w"].T + w["head.0.b"]
z = 0.5 * z * (1.0 + np.tanh(0.7978845608 * (z + 0.044715 * z ** 3)))
return z @ w["head.1.w"].T + w["head.1.b"]
# -- whole windows (training parity / batch scoring)
def windows(self, x_num, port, proto, idx: np.ndarray, h_pool=None) -> np.ndarray:
"""idx [n_win, K] row indices (-1 = pad) -> logits [n_win, K, 1 + n_fam] (pads = 0)."""
n_win, K = idx.shape
e = self.embed(x_num, port, proto)
state = np.zeros((self.layers, n_win, self.d), np.float32)
gamma, beta = self.context(h_pool)
out = np.zeros((n_win, K, 1 + len(self.families)), np.float32)
for k in range(K):
rows = idx[:, k]
live = rows >= 0
if not live.any():
break
nxt = self.gru_step(e[np.where(live, rows, 0)], state)
state = np.where(live[None, :, None], nxt, state)
out[live, k] = self.head(state[-1][live], gamma, beta)
return out
def pool(self, x_num, port, proto, idx: np.ndarray) -> np.ndarray:
"""Benign reference windows -> mean final top-layer hidden [d] (the site context input)."""
n_win, K = idx.shape
e = self.embed(x_num, port, proto)
state = np.zeros((self.layers, n_win, self.d), np.float32)
for k in range(K):
rows = idx[:, k]
live = rows >= 0
if not live.any():
break
nxt = self.gru_step(e[np.where(live, rows, 0)], state)
state = np.where(live[None, :, None], nxt, state)
return state[-1].mean(0)
def proba_attack(self, logits: np.ndarray) -> np.ndarray:
return _sigmoid(logits[..., 0] / float(self.info.get("temperature", 1.0)))
# -- io
def save(self, path: str) -> None:
np.savez(path, **self.w, info=np.frombuffer(json.dumps(self.info).encode(), np.uint8))
@classmethod
def load(cls, path: str) -> "SeqNet":
z = np.load(path)
info = json.loads(z["info"].tobytes().decode())
return cls({k: z[k] for k in z.files if k != "info"}, info)