From 9f92aa3a38d1fd02620eadb8efa65c6b13b4ba0b Mon Sep 17 00:00:00 2001 From: Aleksei Fedorov Date: Sun, 9 Aug 2026 14:56:52 +0000 Subject: [PATCH] ML stages --- kernel_ai/ml/collectors/base.py | 104 ++++++++++++++++++++------------ kernel_ai/ml/config.py | 4 +- kernel_ai/ml/sequence.py | 27 ++++++++- kernel_ai/ml/worker.py | 43 ++++++++++--- tests/test_ml_collectors.py | 42 ++++++++++++- 5 files changed, 169 insertions(+), 51 deletions(-) diff --git a/kernel_ai/ml/collectors/base.py b/kernel_ai/ml/collectors/base.py index 6065ec7..40ca16d 100644 --- a/kernel_ai/ml/collectors/base.py +++ b/kernel_ai/ml/collectors/base.py @@ -3,6 +3,8 @@ from __future__ import annotations import json +import os +import platform from dataclasses import asdict, dataclass from typing import Iterable, Iterator @@ -43,45 +45,13 @@ } ) -# Linux x86_64 syscall numbers for the allowlist (audit logs emit numbers). -ALLOWED_SYSCALL_NR: frozenset[int] = frozenset( - { - 56, # clone - 57, # fork - 58, # vfork - 59, # execve - 322, # execveat - 435, # clone3 - 41, # socket (not scored alone; kept out — connect/accept matter more) - 42, # connect - 43, # accept - 49, # bind - 50, # listen - 288, # accept4 - 2, # open - 257, # openat - 85, # creat - 263, # unlinkat - 264, # renameat - 316, # renameat2 - 9, # mmap - 10, # mprotect - 330, # pkey_mprotect - 105, # setuid - 113, # setreuid - 117, # setresuid - 106, # setgid - 114, # setregid - 119, # setresgid - 101, # ptrace - 310, # process_vm_writev - 319, # memfd_create - 323, # userfaultfd - } -) +# audit ARCH_* bitmasks (see linux/audit.h) — numbers collide across arches +# (e.g. x86_64 56=clone vs aarch64 56=openat), so maps must stay separate. +AUDIT_ARCH_X86_64 = 0xC000003E +AUDIT_ARCH_AARCH64 = 0xC00000B7 -# Minimal nr → name map for allowlisted calls (collector / audit parser). -SYSCALL_NR_TO_NAME: dict[int, str] = { +# Linux x86_64 syscall numbers for the allowlist (audit logs emit numbers). +SYSCALL_NR_TO_NAME_X86_64: dict[int, str] = { 56: "clone", 57: "fork", 58: "vfork", @@ -114,6 +84,64 @@ 323: "userfaultfd", } +# Linux aarch64 — fork/vfork/open/creat are not separate syscalls. +SYSCALL_NR_TO_NAME_AARCH64: dict[int, str] = { + 220: "clone", + 221: "execve", + 281: "execveat", + 435: "clone3", + 203: "connect", + 202: "accept", + 200: "bind", + 201: "listen", + 242: "accept4", + 56: "openat", + 35: "unlinkat", + 38: "renameat", + 276: "renameat2", + 222: "mmap", + 226: "mprotect", + 288: "pkey_mprotect", + 146: "setuid", + 145: "setreuid", + 147: "setresuid", + 144: "setgid", + 143: "setregid", + 149: "setresgid", + 117: "ptrace", + 271: "process_vm_writev", + 279: "memfd_create", + 282: "userfaultfd", +} + + +def _host_nr_map() -> dict[int, str]: + machine = (platform.machine() or os.uname().machine or "").lower() + if machine in {"aarch64", "arm64"}: + return SYSCALL_NR_TO_NAME_AARCH64 + return SYSCALL_NR_TO_NAME_X86_64 + + +def nr_map_for_audit_arch(arch: int | None) -> dict[int, str]: + """Pick nr→name table from audit ``arch=`` field (or host default).""" + if arch == AUDIT_ARCH_AARCH64: + return SYSCALL_NR_TO_NAME_AARCH64 + if arch == AUDIT_ARCH_X86_64: + return SYSCALL_NR_TO_NAME_X86_64 + return _host_nr_map() + + +def resolve_syscall_name(nr: int, arch: int | None = None) -> str | None: + """Return allowlisted name for ``nr``, or None if not in the arch map.""" + return nr_map_for_audit_arch(arch).get(nr) + + +# Back-compat: host-arch map + union of known numbers (filter only; name via resolve). +SYSCALL_NR_TO_NAME: dict[int, str] = dict(_host_nr_map()) +ALLOWED_SYSCALL_NR: frozenset[int] = frozenset( + set(SYSCALL_NR_TO_NAME_X86_64) | set(SYSCALL_NR_TO_NAME_AARCH64) +) + @dataclass(frozen=True) class SyscallEvent: diff --git a/kernel_ai/ml/config.py b/kernel_ai/ml/config.py index ad1bd59..ca24702 100644 --- a/kernel_ai/ml/config.py +++ b/kernel_ai/ml/config.py @@ -135,7 +135,9 @@ class MLConfig: seq_n: int = _env_int("KERNEL_AI_ML_SEQ_N", 3) # n-gram size seq_max_pids: int = _env_int("KERNEL_AI_ML_SEQ_MAX_PIDS", 512) # pids sampled/tick seq_window: int = _env_int("KERNEL_AI_ML_SEQ_WINDOW", 400) # rolling n-grams scored - seq_min_window: int = _env_int("KERNEL_AI_ML_SEQ_MIN_WINDOW", 120) # before scoring + seq_min_window: int = _env_int("KERNEL_AI_ML_SEQ_MIN_WINDOW", 120) # global fallback + # Per-pid STIDE window (avoids dilution by high-volume connect/setuid spam). + seq_pid_min_window: int = _env_int("KERNEL_AI_ML_SEQ_PID_MIN_WINDOW", 24) # 2s tick sampling only catches processes *parked* in a syscall. A short burst # of sub-samples per tick captures real syscall transitions (better sequences). seq_subsamples: int = _env_int("KERNEL_AI_ML_SEQ_SUBSAMPLES", 4) diff --git a/kernel_ai/ml/sequence.py b/kernel_ai/ml/sequence.py index 09a0d23..6698d2f 100644 --- a/kernel_ai/ml/sequence.py +++ b/kernel_ai/ml/sequence.py @@ -80,11 +80,18 @@ def __init__(self, n: int = 3, window: int = 400, max_pids: int = 4096) -> None: self.window = window self.max_pids = max_pids self._hist: dict[int, deque[str]] = {} - # Rolling window of recent n-gram keys used for live scoring. + # Global rolling window (profile growth / Stage 8 / fallback scoring). self._recent: deque[str] = deque(maxlen=window) + # Per-pid rolling windows — STIDE scores these so connect-spam from one + # daemon cannot dilute a hostile short chain on another pid. + self._recent_by_pid: dict[int, deque[str]] = {} # Counts of every n-gram observed since the last flush (profile growth). self._pending: dict[str, int] = {} + def _drop_pid(self, pid: int) -> None: + self._hist.pop(pid, None) + self._recent_by_pid.pop(pid, None) + def _append(self, pid: int, name: str) -> None: hist = self._hist.get(pid) if hist is None: @@ -95,12 +102,17 @@ def _append(self, pid: int, name: str) -> None: key = _SEP.join(hist) self._recent.append(key) self._pending[key] = self._pending.get(key, 0) + 1 + pid_win = self._recent_by_pid.get(pid) + if pid_win is None: + pid_win = deque(maxlen=self.window) + self._recent_by_pid[pid] = pid_win + pid_win.append(key) def update(self, samples: dict[int, str]) -> None: # Drop histories for pids that vanished to bound memory. if len(self._hist) > self.max_pids: for dead in [p for p in self._hist if p not in samples]: - self._hist.pop(dead, None) + self._drop_pid(dead) for pid, name in samples.items(): self._append(int(pid), str(name)) @@ -122,12 +134,21 @@ def update_stream(self, events) -> int: # Bound memory: drop oldest pid histories opportunistically. overflow = len(self._hist) - self.max_pids for dead in list(self._hist.keys())[:overflow]: - self._hist.pop(dead, None) + self._drop_pid(dead) return n def recent(self) -> list[str]: return list(self._recent) + def recent_by_pid(self, *, min_len: int) -> list[tuple[int, list[str]]]: + """Return ``(pid, ngram_window)`` for pids with enough recent n-grams.""" + need = max(1, int(min_len)) + out: list[tuple[int, list[str]]] = [] + for pid, dq in self._recent_by_pid.items(): + if len(dq) >= need: + out.append((pid, list(dq))) + return out + def drain_pending(self) -> dict[str, int]: """Return + clear n-gram counts accumulated since the last drain.""" pending = self._pending diff --git a/kernel_ai/ml/worker.py b/kernel_ai/ml/worker.py index e4f5b03..145a298 100644 --- a/kernel_ai/ml/worker.py +++ b/kernel_ai/ml/worker.py @@ -92,8 +92,15 @@ def _build_isoforest_anomaly(score: float, scores: dict[str, Score], cfg: MLConf } -def _build_sequence_anomaly(mismatch: float, misses: int, window_len: int, - top_unseen: list[str], cfg: MLConfig) -> dict: +def _build_sequence_anomaly( + mismatch: float, + misses: int, + window_len: int, + top_unseen: list[str], + cfg: MLConfig, + *, + pid: int | None = None, +) -> dict: """Build a mutation from a STIDE syscall-sequence verdict (Stage 4). Unlike Stages 1-2 (which judge *magnitudes*), this fires when the recent @@ -102,6 +109,10 @@ def _build_sequence_anomaly(mismatch: float, misses: int, window_len: int, """ severity = "high" if mismatch >= cfg.seq_mismatch_crit else "medium" cause = ("; novel: " + ", ".join(top_unseen)) if top_unseen else "" + pid_bit = f" pid={pid}" if pid is not None else "" + meta = {"stage": 4, "mismatch": round(mismatch, 4), "window": window_len} + if pid is not None: + meta["pid"] = int(pid) return { "source": "stage4_sequence", "feature": "syscall_seq", @@ -114,10 +125,10 @@ def _build_sequence_anomaly(mismatch: float, misses: int, window_len: int, "baseline_std": None, "position": 0.22, "message": ( - f"Unusual syscall sequencing: {mismatch * 100:.0f}% of recent " + f"Unusual syscall sequencing{pid_bit}: {mismatch * 100:.0f}% of recent " f"{window_len} n-grams are novel ({misses} unseen){cause}" ), - "meta": {"stage": 4, "mismatch": round(mismatch, 4), "window": window_len}, + "meta": meta, } @@ -282,17 +293,33 @@ def _tick_sequence(self) -> dict | None: if self.seq_model is None: return None - window = self.seq_tracker.recent() - if len(window) < self.cfg.seq_min_window: + + # Prefer per-pid windows (classic STIDE): high-volume connect/setuid + # spam from one process must not dilute a hostile chain on another. + best: tuple[float, int, list[str], int | None] | None = None + for pid, window in self.seq_tracker.recent_by_pid( + min_len=self.cfg.seq_pid_min_window + ): + mismatch, misses = self.seq_model.score_window(window) + if best is None or mismatch > best[0]: + best = (mismatch, misses, window, pid) + if best is None: + window = self.seq_tracker.recent() + if len(window) >= self.cfg.seq_min_window: + mismatch, misses = self.seq_model.score_window(window) + best = (mismatch, misses, window, None) + if best is None: return None - mismatch, misses = self.seq_model.score_window(window) + mismatch, misses, window, pid = best if mismatch < self.cfg.seq_mismatch_warn: return None if (now - self._last_seq_emit) < self.cfg.seq_cooldown_sec: return None self._last_seq_emit = now top = self.seq_model.top_unseen(window, limit=3) - return _build_sequence_anomaly(mismatch, misses, len(window), top, self.cfg) + return _build_sequence_anomaly( + mismatch, misses, len(window), top, self.cfg, pid=pid + ) def _tick_stage8(self) -> dict | None: """Stage 8 stub: score the Stage 4 rolling window if a model is ready.""" diff --git a/tests/test_ml_collectors.py b/tests/test_ml_collectors.py index a5a4a33..b1de8c7 100644 --- a/tests/test_ml_collectors.py +++ b/tests/test_ml_collectors.py @@ -1,9 +1,24 @@ """Tests for Stage 6 syscall event contract + n-gram stream ingest.""" -from kernel_ai.ml.collectors.base import SyscallEvent, decode_events, encode_events +from kernel_ai.ml.collectors.base import ( + AUDIT_ARCH_AARCH64, + AUDIT_ARCH_X86_64, + SyscallEvent, + decode_events, + encode_events, + resolve_syscall_name, +) from kernel_ai.ml.sequence import NgramTracker +def test_resolve_syscall_nr_arch_collision(): + """Same numeric id must not cross-map between x86_64 and aarch64.""" + assert resolve_syscall_name(56, AUDIT_ARCH_X86_64) == "clone" + assert resolve_syscall_name(56, AUDIT_ARCH_AARCH64) == "openat" + assert resolve_syscall_name(117, AUDIT_ARCH_X86_64) == "setresuid" + assert resolve_syscall_name(117, AUDIT_ARCH_AARCH64) == "ptrace" + + def test_encode_decode_roundtrip(): events = [ SyscallEvent(ts=1.0, pid=10, uid=0, comm="bash", syscall="clone"), @@ -30,3 +45,28 @@ def test_ngram_tracker_update_stream_builds_trigrams(): assert "openat|execve|connect" in recent pending = tracker.drain_pending() assert pending["clone|openat|execve"] == 1 + + +def test_ngram_tracker_recent_by_pid_not_diluted(): + """Hostile short chain on pid B must remain visible vs spam on pid A.""" + from kernel_ai.ml.sequence import StideModel + + tracker = NgramTracker(n=3, window=80) + spam = [ + SyscallEvent(ts=float(i), pid=1, uid=0, comm="nginx", syscall="connect") + for i in range(200) + ] + novel = [ + SyscallEvent(ts=100 + i * 0.01, pid=2, uid=0, comm="evil", syscall=name) + for i, name in enumerate( + ["ptrace", "memfd_create", "execve", "connect", "setuid"] * 8 + ) + ] + tracker.update_stream(spam + novel) + by_pid = dict(tracker.recent_by_pid(min_len=24)) + assert 2 in by_pid + assert any(g.startswith("ptrace|") for g in by_pid[2]) + model = StideModel(n=3, ngrams={"connect|connect|connect"}) + global_m, _ = model.score_window(tracker.recent()) + pid_m, _ = model.score_window(by_pid[2]) + assert pid_m > global_m