| """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 |
| PROTO_VOCAB = 5 |
| FLAG_BITS = (0x01, 0x02, 0x04, 0x08, 0x10, 0x20) |
| 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 |
|
|
| |
| 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))) |
| 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"] |
|
|
| |
| 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))) |
|
|
| |
| 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) |
|
|
