Blame for runtime/python/dibaguard/decide.py

"""Decision layer: calibrated probabilities -> guaranteed alarm budget, drift detection, entity risk, site pressure.

* Temperature: per-group temperature scaling of logits (fitted by NLL on a calibration split).
* ConformalFPR (Mondrian): per group (site / endpoint group / source type) the alarm threshold is the conformal
  quantile of that group's benign scores, so P(alarm | benign) <= alpha holds for exchangeable benign traffic with a
  finite-sample guarantee - not a vendor-chosen cut-off. pvalue() gives the conformal p-value of any score.
* DriftMartingale: power martingale over the benign p-values; by Ville's inequality it exceeds 1/delta with
  probability <= delta while traffic stays exchangeable, so crossing it means "recalibrate" (normal traffic changed).
* EntityRisk: per-entity risk in [0, 1] that rises with each suspicious event (r <- 1 - (1-r)(1-s), s = p x severity)
  and decays with half-life `half_life_s`; bounded memory (LRU over `max_entities`).
* SitePressure: exponentially weighted share of attack traffic, overall and per family.
All O(1) or O(log n) per event, numpy only.
"""
from __future__ import annotations

import math
from bisect import bisect_left, insort
from collections import OrderedDict

import numpy as np


class Temperature:
    def __init__(self):
        self.t: dict = {}

    @staticmethod
    def _nll(logits: np.ndarray, y: np.ndarray, t: float) -> float:
        z = logits / t
        if z.ndim == 1:   # binary: one logit per row
            return float(np.mean(np.logaddexp(0, z) - y * z))
        z = z - z.max(1, keepdims=True)
        return float(np.mean(np.log(np.exp(z).sum(1)) - z[np.arange(len(y)), y]))

    def fit(self, group: str, logits, y) -> float:
        logits, y = np.asarray(logits, np.float64), np.asarray(y)
        grid = np.exp(np.linspace(math.log(0.05), math.log(20), 121))
        best = min(grid, key=lambda t: self._nll(logits, y, t))
        lo, hi = best / 1.06, best * 1.06   # refine around the grid point
        for _ in range(40):
            m1, m2 = lo + (hi - lo) / 3, hi - (hi - lo) / 3
            if self._nll(logits, y, m1) < self._nll(logits, y, m2):
                hi = m2
            else:
                lo = m1
        self.t[group] = float((lo + hi) / 2)
        return self.t[group]

    def __call__(self, group: str, logits):
        return np.asarray(logits, np.float64) / self.t.get(group, 1.0)


class ConformalFPR:
    """Benign calibration scores per group; higher score = more suspicious."""

    def __init__(self, max_per_group: int = 20000, min_n: int = 200):
        self.cal: dict = {}
        self.max_per_group, self.min_n = max_per_group, min_n

    def add(self, group: str, scores) -> None:
        arr = self.cal.setdefault(group, [])
        new = np.asarray(scores, np.float64).ravel()
        if len(new) > 64:   # bulk: one sorted merge instead of len(new) list insertions
            arr = self.cal[group] = np.sort(np.concatenate([np.asarray(arr, np.float64), new])).tolist()
        else:
            for s in new:
                insort(arr, float(s))
        if len(arr) > self.max_per_group:   # keep an evenly thinned sorted sample (quantiles preserved)
            self.cal[group] = arr[:: math.ceil(len(arr) / self.max_per_group)]

    def n(self, group: str) -> int:
        return len(self.cal.get(group, ()))

    def _group(self, group: str, fallback: str | None) -> list:
        arr = self.cal.get(group)
        if (arr is None or len(arr) < self.min_n) and fallback is not None:
            arr = self.cal.get(fallback)
        if not arr:
            raise KeyError(f"no calibration scores for group {group!r}")
        return arr

    def pvalue(self, group: str, score: float, fallback: str | None = "*") -> float:
        """(1 + #{benign >= score}) / (n + 1): small = unlike the group's benign traffic."""
        arr = self._group(group, fallback)
        return (1 + len(arr) - bisect_left(arr, float(score) - _tol(score))) / (len(arr) + 1)

    def threshold(self, group: str, alpha: float, fallback: str | None = "*", confidence: float | None = None) -> float:
        """Alarm threshold for false-alarm rate alpha on benign traffic.

        confidence=None: marginal guarantee E[FPR] <= alpha (the ceil((n+1)(1-alpha))-th order statistic).
        confidence=c   : PAC guarantee P(FPR <= alpha) >= c over the draw of the calibration set. With the k-th order
                         statistic, FPR > alpha exactly when at most n-k calibration scores lie above the true
                         (1-alpha)-quantile, a Binomial(n, alpha) event; k is the smallest order with that probability
                         <= 1-c. Returns inf when n is too small (see n_needed)."""
        arr = self._group(group, fallback)
        n = len(arr)
        if confidence is None:
            k = math.ceil((n + 1) * (1 - alpha))
            if k > n and n >= ConformalFPR.n_needed(alpha, 0.5):
                k = n                     # the maximum holds FPR <= alpha with probability >= 1/2 (see n_needed)
        else:
            m = _binom_max_m(n, alpha, 1 - confidence)
            if m < 0:
                return float("inf")
            k = n - m
        return float("inf") if k > n else arr[k - 1] + _tol(arr[k - 1])

    @staticmethod
    def n_needed(alpha: float, confidence: float) -> int:
        """Fewest benign calibration scores for which the PAC threshold exists: (1-alpha)^n <= 1-confidence."""
        return math.ceil(math.log(1 - confidence) / math.log(1 - alpha))

    def alarm(self, group: str, score: float, alpha: float, fallback: str | None = "*") -> bool:
        return self.pvalue(group, score, fallback) <= alpha


def _binom_max_m(n: int, p: float, delta: float) -> int:
    """Largest m with P(Binomial(n, p) <= m) <= delta (-1 if none), by summing the pmf in log space."""
    lp, lq = math.log(p), math.log1p(-p)
    cdf, m = 0.0, -1
    for j in range(n + 1):
        cdf += math.exp(math.lgamma(n + 1) - math.lgamma(j + 1) - math.lgamma(n - j + 1) + j * lp + (n - j) * lq)
        if cdf > delta:
            return m
        m = j
    return m


class DriftMartingale:
    """Power martingale M_t = prod eps * p_i^(eps-1) over p-values of traffic believed benign."""

    def __init__(self, eps: float = 0.92, delta: float = 1e-3):
        self.eps, self.delta, self.log_m, self.n = eps, delta, 0.0, 0

    def update(self, p: float) -> bool:
        self.log_m += math.log(self.eps) + (self.eps - 1) * math.log(max(p, 1e-12))
        self.log_m = max(self.log_m, 0.0)    # restart from 1 when evidence fades (detects later shifts too)
        self.n += 1
        return self.drifted

    @property
    def drifted(self) -> bool:
        return self.log_m >= math.log(1 / self.delta)

    def reset(self) -> None:
        self.log_m, self.n = 0.0, 0


def _tol(v):
    """scores equal up to float32 rounding are equal: model scores are float32 and come out a unit in the last place
    apart depending on the batch they were computed in. Against a reference full of identical values (every new
    client's first request scores the same) a one-ulp difference was an alarm (p = 1 / (n + 1))"""
    return 1e-5 * np.maximum(1.0, np.abs(v))


def _late(clock: float, t: float, tau: float) -> float:
    """weight of an event at time t in a decayed counter whose clock is already at `clock` (1 when not late): an
    out-of-order event is counted as what it is worth now - neither rewinding the clock nor counted as fresh"""
    return 1.0 if t >= clock else math.exp(-(clock - t) / tau)


def evict(od: "OrderedDict", keep, look: int = 8) -> None:
    """Drop one entry of a bounded per-source table, least recently used first - except that an entry which has earned
    state (keep(value)) goes to the back of the queue instead, up to `look` times. A spray of fresh one-event sources
    (spoofed, or addresses rotated inside an IPv6 /64) then evicts other fresh sources, not the attacker's record it
    was sent to flush; a table full of earned entries still evicts its oldest."""
    for _ in range(look):
        k, v = od.popitem(last=False)
        if not keep(v):
            return
        od[k] = v
    od.popitem(last=False)


class EntityRisk:
    def __init__(self, half_life_s: float = 600.0, max_entities: int = 1_000_000):
        self.tau = half_life_s / math.log(2)
        self.max = max_entities
        self.s: "OrderedDict[str, list]" = OrderedDict()   # key -> [risk, last_t, events, alerts]

    def update(self, key: str, t: float, p_attack: float, severity: float = 1.0, alert: bool = False) -> float:
        st = self.s.get(key)
        if st is None:
            st = [0.0, t, 0, 0]
            self.s[key] = st
            if len(self.s) > self.max:
                evict(self.s, lambda v: v[0] > 0.2 or v[3] > 0)
        else:
            self.s.move_to_end(key)
            st[0] *= math.exp(-max(0.0, t - st[1]) / self.tau)
        sv = min(1.0, max(0.0, p_attack * severity)) * _late(st[1], t, self.tau)
        st[0] = 1 - (1 - st[0]) * (1 - sv)
        st[1], st[2], st[3] = max(st[1], t), st[2] + 1, st[3] + int(alert)     # a late (older) event never rewinds the clock
        return st[0]

    def risk(self, key: str, t: float) -> float:
        st = self.s.get(key)
        return 0.0 if st is None else st[0] * math.exp(-max(0.0, t - st[1]) / self.tau)

    def top(self, t: float, k: int = 10) -> list:
        return sorted(((self.risk(k_, t), k_) for k_ in self.s), reverse=True)[:k]


class SitePressure:
    """Exponentially weighted share of attack traffic (time constant `window_s`), overall and per family."""

    def __init__(self, window_s: float = 300.0):
        self.tau, self.t, self.w, self.a, self.fam = window_s, None, 0.0, 0.0, {}

    def update(self, t: float, p_attack: float, family: str | None = None) -> float:
        if self.t is not None:
            k = math.exp(-max(0.0, t - self.t) / self.tau)
            self.w *= k
            self.a *= k
            for f in self.fam:
                self.fam[f] *= k
        wt = 1.0 if self.t is None else _late(self.t, t, self.tau)
        self.t = t if self.t is None else max(self.t, t)
        self.w += wt
        self.a += wt * p_attack
        if family:
            self.fam[family] = self.fam.get(family, 0.0) + wt * p_attack
        return self.pressure

    @property
    def pressure(self) -> float:
        return self.a / self.w if self.w else 0.0

    def by_family(self) -> dict:
        return {f: v / self.w for f, v in sorted(self.fam.items(), key=lambda x: -x[1])} if self.w else {}


def event_score(p_attack: float, severity: float) -> int:
    """0-100 'how much attack' for one event: probability weighted toward severe families."""
    return int(round(100 * p_attack * (0.5 + 0.5 * severity)))


class RollingReference:
    """Deployment-faithful calibration reference: the last `keep` time chunks of `window_s` seconds of traffic that
    was *not* alarmed (self-cleaning, no labels), as sorted score arrays. Thresholds / p-values come from the
    merged chunks; the chunk being filled is never used (it would include the row under test). For feature-level
    anomaly (quantile p-values), the same object stores per-feature sorted columns."""

    def __init__(self, window_s: float = 1800.0, keep: int = 2, max_rows: int = 200_000):
        self.window_s, self.keep, self.max_rows = window_s, keep, max_rows
        self.done: list = []            # finished chunks: sorted arrays (scores) or [n, F] column-sorted arrays
        self.cur: list = []             # raw rows of the chunk being filled
        self.cur_k: int | None = None
        self._merged = None
        self.base = None                # the site's one-off calibration (diversity); recent chunks add recency

    def add(self, t: float, value) -> None:
        k = int(t // self.window_s)
        if self.cur_k is None:
            self.cur_k = k
        if k != self.cur_k:
            self._close()
            self.cur_k = k
        if len(self.cur) < self.max_rows:
            self.cur.append(value)

    def _close(self) -> None:
        if self.cur:
            a = np.asarray(self.cur, np.float64)
            self.done.append(np.sort(a, axis=0))
            self.done = self.done[-self.keep:]
            self._merged = None
        self.cur = []

    def set_base(self, values) -> None:
        self.base = np.sort(np.asarray(values, np.float64), axis=0) if len(values) else None
        self._merged = None

    def ref(self):
        parts = ([self.base] if self.base is not None else []) + self.done
        if self._merged is None and parts:
            self._merged = np.sort(np.concatenate(parts, axis=0), axis=0)
        return self._merged

    def n(self) -> int:
        r = self.ref()
        return 0 if r is None else len(r)

    def threshold(self, alpha: float) -> float:
        r = self.ref()
        if r is None or r.ndim != 1:
            return float("inf")
        k = math.ceil((len(r) + 1) * (1 - alpha))
        if k > len(r) and len(r) >= ConformalFPR.n_needed(alpha, 0.5):
            # fewer than 1/alpha scores, but enough for the maximum to hold FPR <= alpha with probability >= 1/2: the
            # engine treats such a reference as usable, so it must not answer 'no threshold' (inf = never alarm)
            k = len(r)
        return float("inf") if k > len(r) else float(r[k - 1]) + _tol(r[k - 1])

    def pvalues(self, scores) -> np.ndarray:
        r = self.ref()
        s = np.asarray(scores, np.float64)
        if r is None:
            return np.ones_like(s)
        return (1 + len(r) - np.searchsorted(r, s - _tol(s), side="left")) / (len(r) + 1)

    def feature_pvalues(self, x) -> np.ndarray:
        """x [n, F] -> upper-tail per-feature empirical p-values (1 + #ref >= x) / (n + 1), Bonferroni over F, min -> [n]."""
        return feature_pvalues(self.ref(), x)

    def feature_scores(self, x) -> np.ndarray:
        """Fisher-style anomaly score: sum over features of -log upper-tail p (conformally calibrated by Guard)."""
        return feature_scores(self.ref(), x)


def feature_scores(ref, x) -> np.ndarray:
    """x [n, F] vs a column-sorted reference [m, F]: sum_j -log((1 + #ref_j >= x_j) / (m + 1)). A source that is
    extreme on many attack-direction features scores far above a benign one that is unusual on a single feature -
    the reason the min-p rule (one rare event = alarm) is not used."""
    x = np.asarray(x, np.float64)
    if ref is None or ref.ndim != 2:
        return np.zeros(len(x))
    m, F = ref.shape
    p = (1 + m - np.stack([np.searchsorted(ref[:, j], x[:, j] - _tol(x[:, j]), side="left") for j in range(F)], 1)) / (m + 1)
    return -np.log(p).sum(1)


def feature_pvalues(ref, x) -> np.ndarray:
    x = np.asarray(x, np.float64)
    if ref is None or ref.ndim != 2:
        return np.ones(len(x))
    n, F = ref.shape
    hi = (1 + n - np.stack([np.searchsorted(ref[:, j], x[:, j] - _tol(x[:, j]), side="left") for j in range(F)], 1)) / (n + 1)
    return np.minimum(1.0, hi * F).min(1)


class SourceEvidence:
    """Per-source accumulated evidence: E <- E * exp(-dt / tau) + (-log p) over the source's flows. One scanner, one
    brute-forcer or one flood crosses a threshold within a handful of flows; a benign client, whose p-values are
    uniform, never accumulates. The unit judged is the (source, minute): the reference is the maximum E inside
    each benign source-minute of the site's normal traffic (conformal threshold on those maxima)."""

    def __init__(self, tau_s: float = 120.0, max_entities: int = 500_000):
        self.tau, self.max = tau_s, max_entities
        self.e: "OrderedDict[object, list]" = OrderedDict()     # src -> [E, last_t, minute, minute_max]

    def update(self, src, t: float, logp: float) -> float:
        st = self.e.get(src)
        if st is None:
            st = [0.0, t, int(t // 60), 0.0]
            self.e[src] = st
            if len(self.e) > self.max:
                evict(self.e, lambda v: v[0] > 1.0)
        else:
            self.e.move_to_end(src)
            st[0] *= math.exp(-max(0.0, t - st[1]) / self.tau)
        st[0] = max(0.0, st[0] + logp * _late(st[1], t, self.tau))   # CUSUM: floored at 0 (a benign source drifts back)
        if t < st[1]:                         # a late (older) event: counted at its decayed weight, never rewinding the
            t = st[1]                         # clock - a rewind would make the next event decay everything accumulated
        st[1] = t
        m = int(t // 60)
        if m != st[2]:
            st[2], st[3] = m, 0.0
        st[3] = max(st[3], st[0])
        return st[0]

    def evidence(self, src) -> float:
        st = self.e.get(src)
        return 0.0 if st is None else st[0]


class GlobalRate:
    """Site-wide failure rate per key (an HTTP path, or a (dst, port) service), as an exponentially-decayed count.
    Distributed campaigns - credential stuffing or scanning spread over hundreds of low-volume source IPs, where no
    single source accumulates evidence - show up only in the aggregate: one key suddenly fails far more than its
    calibrated baseline. The response is a challenge on the contributing cohort, never a block, so benign users of a
    briefly-hot path are only challenged, not blocked."""

    def __init__(self, tau_s: float = 120.0, max_keys: int = 200_000, base_tau_s: float = 3600.0, rel: float = 4.0,
                 hot_learn: float = 0.1):
        self.tau, self.max = tau_s, max_keys
        # the key's own baseline: the same failures counted over a long horizon. A busy login path that always sees
        # some failures (or a large site whose 0.5% failed connects add up to a few per second) is hot only well above
        # its own normal, not above an absolute floor (measured: 96% of a busy site's benign failed connects alarmed
        # with the floor alone). While hot, the baseline learns at `hot_learn` of the rate: a campaign is not adopted
        # within minutes, a lasting benign level shift is, within hours
        self.base_tau, self.rel, self.hot_learn = base_tau_s, rel, hot_learn
        self.c: "OrderedDict[object, list]" = OrderedDict()   # key -> [rate, last_t, long-horizon count]
        self.ref = RollingReference(1800.0)                   # benign per-minute peak failure rate per key

    def _bump(self, key, t: float, inc: float, floor: float | None = None) -> tuple[float, bool]:
        """-> (decayed count on the short horizon, whether it is over floor + rel x the key's own expected count)"""
        st = self.c.get(key)
        if st is None:
            st = [0.0, t, 0.0]
            self.c[key] = st
            if len(self.c) > self.max:
                self.c.popitem(last=False)
        else:
            self.c.move_to_end(key)
            if len(st) < 3:                                    # restored from a snapshot without the baseline
                st.append(st[0] * self.base_tau / self.tau)
            dt = max(0.0, t - st[1])
            st[0] *= math.exp(-dt / self.tau)
            st[2] *= math.exp(-dt / self.base_tau)
        if t < st[1]:                                          # a late event: its weight at the key's clock
            inc *= _late(st[1], t, self.tau)
        st[0] += inc
        st[1] = max(st[1], t)
        hot = floor is not None and inc > 0 and st[0] > floor + self.rel * st[2] * self.tau / self.base_tau
        st[2] += inc * (self.hot_learn if hot else 1.0)
        return st[0], hot

    def expected(self, key, t: float) -> float:
        """the key's own normal on the short horizon (its long-horizon count scaled to tau), decayed to t"""
        st = self.c.get(key)
        if st is None or len(st) < 3:
            return 0.0
        return st[2] * math.exp(-max(0.0, t - st[1]) / self.base_tau) * self.tau / self.base_tau

    def calibrate(self, keys, t, failed) -> None:
        """benign window -> the reference of per-(key, minute) peak failure rates, and every key's own baseline"""
        import numpy as np
        peaks: dict = {}
        tmp = GlobalRate(self.tau, self.max, self.base_tau, self.rel, self.hot_learn)
        for i in range(len(t)):
            r, _ = tmp._bump(keys[i], float(t[i]), 1.0 if failed[i] else 0.0)
            k = (keys[i], int(t[i] // 60))
            peaks[k] = max(peaks.get(k, 0.0), r)
        self.ref.set_base(np.fromiter(peaks.values(), float) if peaks else np.zeros(1))
        for key, st in tmp.c.items():
            self.c[key] = st
            self.c.move_to_end(key)
        while len(self.c) > self.max:
            self.c.popitem(last=False)

    def score(self, keys, t, failed, floor: float = 10.0, rates: bool = False):
        """-> boolean array: this failed request is to a key whose site-wide decayed failure rate exceeds `floor` plus
        `rel` times the key's own expected rate (a distributed campaign). Challenge-only: a benign path rarely sees that
        many failures within ~tau, a credential-stuffing / scanning cohort does. rates=True also returns the decayed
        rate at each event (for callers that combine the absolute floor with what they know of the source)."""
        import numpy as np
        hot = np.zeros(len(t), bool)
        r = np.zeros(len(t))
        for i in range(len(t)):
            r[i], hot[i] = self._bump(keys[i], float(t[i]), 1.0 if failed[i] else 0.0, floor)
        return (hot, r) if rates else hot