File size: 9,233 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 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 | """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)
|