| """Site profile: every feature re-expressed relative to *this network's* normal traffic, compiled into the trees. |
| |
| A flow model trained on a few testbeds learns absolute values ("> 1 200 bytes", "TTL 64") that mean something |
| different on another network, which is why detectors collapse on networks they have not seen. Here each feature |
| can also be given to the model as its percentile code inside the site's own benign traffic (SiteProfile: a grid of |
| ~100 quantiles per feature, learned from a window of the site's traffic - no attack labels). |
| |
| The percentile code is a monotone step function of the raw value, so a tree test `code(x_f) > b` is exactly |
| `x_f > c` for one raw threshold c. compile() rewrites every such test of a trained Oblivious model into its raw |
| equivalent for one site: the deployed model reads raw features, has the same size and the same speed, and is |
| bit-identical to the model applied to the percentile codes. Adapting to a new network = recompiling the borders. |
| """ |
| from __future__ import annotations |
|
|
| import numpy as np |
|
|
| from dibaguard.flows.trees import Oblivious |
|
|
| LEVELS = np.r_[1e-4, 1e-3, np.arange(0.005, 1.0, 0.01), 0.999, 0.9999] |
|
|
|
|
| class SiteProfile: |
| def __init__(self, grid: np.ndarray): |
| self.grid = np.asarray(grid, np.float32) |
|
|
| @classmethod |
| def fit(cls, x_benign: np.ndarray, levels=LEVELS) -> "SiteProfile": |
| x = np.asarray(x_benign, np.float32) |
| return cls(np.quantile(x, levels, axis=0, method="inverted_cdf").astype(np.float32)) |
|
|
| def code(self, x: np.ndarray) -> np.ndarray: |
| """Percentile code in [0, 1] of every value (mid-rank among the grid points, ties share their middle).""" |
| x = np.asarray(x, np.float32) |
| G = self.grid.shape[0] |
| out = np.empty(x.shape, np.float32) |
| for f in range(x.shape[1]): |
| g = self.grid[:, f] |
| lo = np.searchsorted(g, x[:, f], "left") |
| hi = np.searchsorted(g, x[:, f], "right") |
| out[:, f] = ((lo + hi) / np.float32(2 * G)).astype(np.float32) |
| return out |
|
|
| def raw_border(self, f: int, b: float) -> float: |
| """Raw threshold c with (code(x)[f] > b) <=> (x > c) for every float32 x.""" |
| g = np.unique(self.grid[:, f]) |
| col = self.grid[:, f] |
| G = len(col) |
| lo = np.searchsorted(col, g, "left") |
| hi = np.searchsorted(col, g, "right") |
| at = ((lo + hi) / np.float32(2 * G)).astype(np.float32) |
| above = (hi / np.float32(G)).astype(np.float32) |
| b = np.float32(b) |
| if np.float32(0) > b: |
| return -np.inf |
| for j in range(len(g)): |
| if at[j] > b: |
| return float(np.nextafter(g[j], np.float32(-np.inf))) |
| if above[j] > b: |
| return float(g[j]) |
| return np.inf |
|
|
| def compile(self, model: Oblivious, code_of: dict) -> Oblivious: |
| """model was trained on features where column i is the percentile code of raw column code_of[i] |
| (other columns raw). Returns the same model over raw columns, borders rewritten for this site.""" |
| feat, border = model.feat.copy(), model.border.copy() |
| for s in range(len(feat) - 1): |
| i = int(feat[s]) |
| if i in code_of: |
| feat[s] = code_of[i] |
| border[s] = self.raw_border(code_of[i], float(border[s])) |
| meta = dict(model.meta, site_compiled=True) |
| return Oblivious(feat, border, model.tree_tests, model.leaves, model.bias, model.names, model.classes, meta) |
|
|
