feat: initial EDR system — agent + backend + frontend + ensemble ML
This commit is contained in:
kilyabin
+224
@@ -0,0 +1,224 @@
|
||||
"""Anomaly detection — voting ensemble of three unsupervised models.
|
||||
|
||||
Pipeline:
|
||||
IsolationForest + Local Outlier Factor + One-Class SVM
|
||||
A process is flagged anomalous when at least 2 of 3 models agree.
|
||||
|
||||
Training data (in order of preference):
|
||||
1. UNSW-NB15 dataset CSV (data/UNSW_NB15_training-set.csv)
|
||||
2. Synthetic normal-behaviour data (fallback, 2000 rows)
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import joblib
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
MODEL_PATH = Path(os.getenv("MODEL_PATH", "data/model.pkl"))
|
||||
DATASET_PATH = Path(os.getenv("DATASET_PATH", "data/UNSW_NB15_training-set.csv"))
|
||||
|
||||
# Trusted desktop applications — ML classification skipped.
|
||||
# Threshold alerts (CPU > 80 %, RAM > 3 GB) still fire from the backend.
|
||||
TRUSTED_PROCESSES: frozenset[str] = frozenset({
|
||||
"firefox", "firefox-bin", "Web Content", "Isolated Web Co",
|
||||
"chrome", "chromium", "chromium-browser",
|
||||
"brave", "brave-browser", "opera",
|
||||
"electron", "Electron",
|
||||
"telegram", "Telegram", "telegram-desktop",
|
||||
"slack", "discord", "discord-bin",
|
||||
"code", "code-oss", "vscodium", "cursor",
|
||||
"zoom", "zoom-bin", "zoomus",
|
||||
"obsidian",
|
||||
"spotify", "vlc", "mpv",
|
||||
"libreoffice", "soffice",
|
||||
"thunderbird",
|
||||
"gnome-shell", "plasmashell", "kwin_wayland", "kwin_x11",
|
||||
"Xorg", "Xwayland",
|
||||
})
|
||||
|
||||
# Feature vector order: [cpu_percent, memory_mb, open_files, connections]
|
||||
FEATURE_COLS = ["cpu_percent", "memory_mb", "open_files", "connections"]
|
||||
|
||||
|
||||
class EnsembleDetector:
|
||||
"""Majority-vote ensemble: IsolationForest + LOF + One-Class SVM.
|
||||
|
||||
Each model casts a vote (+1 normal, -1 anomaly).
|
||||
A process is anomalous when the vote sum is ≤ -1 (i.e. 2 or 3 models agree).
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.iforest = None
|
||||
self.lof = None
|
||||
self.ocsvm = None
|
||||
|
||||
# ── training ─────────────────────────────────────────────────────────────
|
||||
|
||||
def train(self, X: np.ndarray) -> None:
|
||||
"""Fit all three models on X and persist the ensemble to disk."""
|
||||
from sklearn.ensemble import IsolationForest
|
||||
from sklearn.neighbors import LocalOutlierFactor
|
||||
from sklearn.svm import OneClassSVM
|
||||
from sklearn.preprocessing import RobustScaler
|
||||
|
||||
logger.info("Training ensemble on %d samples …", X.shape[0])
|
||||
|
||||
# RobustScaler is median/IQR-based — less sensitive to extreme outliers
|
||||
# than StandardScaler, which matters when training data may contain
|
||||
# occasional noisy spikes.
|
||||
scaler = RobustScaler()
|
||||
X_scaled = scaler.fit_transform(X)
|
||||
|
||||
self.iforest = IsolationForest(
|
||||
n_estimators=100,
|
||||
contamination=0.01,
|
||||
random_state=42,
|
||||
n_jobs=-1,
|
||||
)
|
||||
self.iforest.fit(X_scaled)
|
||||
logger.info(" IsolationForest trained")
|
||||
|
||||
# LOF in novelty=True mode so it can predict on new samples (not just
|
||||
# the training set). n_neighbors=20 is a robust default for datasets
|
||||
# of ~1000–5000 rows.
|
||||
self.lof = LocalOutlierFactor(
|
||||
n_neighbors=20,
|
||||
contamination=0.01,
|
||||
novelty=True,
|
||||
n_jobs=-1,
|
||||
)
|
||||
self.lof.fit(X_scaled)
|
||||
logger.info(" LocalOutlierFactor trained")
|
||||
|
||||
# One-Class SVM with RBF kernel. nu=0.01 ≈ expected outlier fraction.
|
||||
# gamma="scale" adapts to feature variance automatically.
|
||||
self.ocsvm = OneClassSVM(kernel="rbf", nu=0.01, gamma="scale")
|
||||
self.ocsvm.fit(X_scaled)
|
||||
logger.info(" OneClassSVM trained")
|
||||
|
||||
MODEL_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||||
joblib.dump({"iforest": self.iforest, "lof": self.lof,
|
||||
"ocsvm": self.ocsvm, "scaler": scaler}, MODEL_PATH)
|
||||
logger.info("Ensemble saved → %s", MODEL_PATH)
|
||||
|
||||
# Keep the scaler attached for inference
|
||||
self._scaler = scaler
|
||||
|
||||
def load_or_train(self) -> None:
|
||||
"""Load persisted ensemble or train a new one."""
|
||||
if MODEL_PATH.exists():
|
||||
bundle = joblib.load(MODEL_PATH)
|
||||
# Support legacy single-model pkl (no "scaler" key)
|
||||
if not isinstance(bundle, dict):
|
||||
logger.info("Legacy single-model pkl detected — retraining ensemble")
|
||||
MODEL_PATH.unlink()
|
||||
self.load_or_train()
|
||||
return
|
||||
self.iforest = bundle["iforest"]
|
||||
self.lof = bundle["lof"]
|
||||
self.ocsvm = bundle["ocsvm"]
|
||||
self._scaler = bundle["scaler"]
|
||||
logger.info("Ensemble loaded from %s", MODEL_PATH)
|
||||
return
|
||||
logger.info("No saved model found — training now")
|
||||
X = self._load_dataset()
|
||||
self.train(X)
|
||||
|
||||
# ── data helpers ──────────────────────────────────────────────────────────
|
||||
|
||||
def _load_dataset(self) -> np.ndarray:
|
||||
if DATASET_PATH.exists():
|
||||
return self._load_unsw()
|
||||
logger.warning("Dataset not found — using synthetic training data")
|
||||
return self._synthetic_data()
|
||||
|
||||
def _load_unsw(self) -> np.ndarray:
|
||||
import pandas as pd
|
||||
df = pd.read_csv(DATASET_PATH)
|
||||
numeric = df.select_dtypes(include=[np.number]).dropna()
|
||||
X = numeric.iloc[:, :4].values.astype(np.float32)
|
||||
logger.info("Loaded UNSW-NB15: %d rows", X.shape[0])
|
||||
return X
|
||||
|
||||
def _synthetic_data(self) -> np.ndarray:
|
||||
"""2000 rows of realistic normal Linux process behaviour.
|
||||
|
||||
Distribution:
|
||||
75 % idle/background (CPU ≈ 0 %, low RAM)
|
||||
20 % light workers (moderate CPU, moderate RAM)
|
||||
5 % heavier services (higher RAM, some connections)
|
||||
"""
|
||||
rng = np.random.default_rng(42)
|
||||
n = 2000
|
||||
|
||||
cpu_idle = rng.exponential(0.3, int(n * 0.75))
|
||||
cpu_light = rng.uniform(0.5, 15, int(n * 0.20))
|
||||
cpu_heavy = rng.uniform(15, 60, int(n * 0.05))
|
||||
cpu = np.concatenate([cpu_idle, cpu_light, cpu_heavy])
|
||||
|
||||
mem_small = rng.uniform(0.5, 50, int(n * 0.50))
|
||||
mem_medium = rng.uniform(50, 300, int(n * 0.30))
|
||||
mem_large = rng.uniform(300, 800, int(n * 0.12))
|
||||
mem_desktop = rng.uniform(800, 2000, int(n * 0.08))
|
||||
mem = np.concatenate([mem_small, mem_medium, mem_large, mem_desktop])
|
||||
|
||||
files = np.concatenate([
|
||||
rng.integers(0, 20, int(n * 0.70)).astype(float),
|
||||
rng.integers(20, 80, int(n * 0.30)).astype(float),
|
||||
])
|
||||
|
||||
conns = np.concatenate([
|
||||
rng.integers(0, 2, int(n * 0.70)).astype(float),
|
||||
rng.integers(2, 20, int(n * 0.30)).astype(float),
|
||||
])
|
||||
|
||||
idx = rng.permutation(n)
|
||||
return np.column_stack([cpu[idx], mem[idx], files[idx], conns[idx]])
|
||||
|
||||
# ── inference ─────────────────────────────────────────────────────────────
|
||||
|
||||
def predict(self, processes: list[dict]) -> list[dict]:
|
||||
"""Annotate each process dict with is_anomaly=True/False.
|
||||
|
||||
Trusted processes skip ML — they can only be flagged by the
|
||||
backend's threshold checks (CPU > 80 %, RAM > 3 GB).
|
||||
"""
|
||||
if not processes or self.iforest is None:
|
||||
for p in processes:
|
||||
p["is_anomaly"] = False
|
||||
return processes
|
||||
|
||||
trusted_idx = {i for i, p in enumerate(processes)
|
||||
if p["name"] in TRUSTED_PROCESSES}
|
||||
candidates = [p for i, p in enumerate(processes) if i not in trusted_idx]
|
||||
|
||||
if candidates:
|
||||
X_raw = np.array(
|
||||
[[p["cpu_percent"], p["memory_mb"], p["open_files"], p["connections"]]
|
||||
for p in candidates],
|
||||
dtype=np.float32,
|
||||
)
|
||||
X = self._scaler.transform(X_raw)
|
||||
|
||||
# Each model returns +1 (normal) or -1 (anomaly)
|
||||
v_if = self.iforest.predict(X) # IsolationForest
|
||||
v_lof = self.lof.predict(X) # Local Outlier Factor
|
||||
v_ocsvm = self.ocsvm.predict(X) # One-Class SVM
|
||||
|
||||
# Majority vote: flag anomaly when vote_sum ≤ -1 (≥ 2 of 3 agree)
|
||||
vote_sum = v_if + v_lof + v_ocsvm
|
||||
for proc, vs in zip(candidates, vote_sum):
|
||||
proc["is_anomaly"] = bool(vs <= -1)
|
||||
|
||||
for i in trusted_idx:
|
||||
processes[i]["is_anomaly"] = False
|
||||
return processes
|
||||
|
||||
|
||||
# Module-level alias so existing import `from model import AnomalyDetector`
|
||||
# continues to work without changes in main.py.
|
||||
AnomalyDetector = EnsembleDetector
|
||||
Reference in New Issue
Block a user