feat(ai-threat-detection): complete Phase 2 — ML ensemble, inference, training pipeline
Add ONNX inference engine with 3-model ensemble (autoencoder, random forest, isolation forest), ONNX export utilities, MLflow experiment tracking, and classifier training scripts. Extend pipeline with hybrid ML/rules detection mode, ensemble score fusion and blending, and full test coverage for all new components. Fix Python 3 SyntaxError in parsers.py (except clause).
This commit is contained in:
parent
da10d89e5a
commit
cf92b157f3
|
|
@ -5,30 +5,60 @@ models_api.py
|
|||
|
||||
import uuid
|
||||
|
||||
from fastapi import APIRouter
|
||||
from fastapi import APIRouter, Request
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.models.model_metadata import ModelMetadata
|
||||
|
||||
router = APIRouter(prefix="/models", tags=["models"])
|
||||
|
||||
|
||||
@router.get("/status")
|
||||
async def model_status() -> dict[str, object]:
|
||||
async def model_status(request: Request, ) -> dict[str, object]:
|
||||
"""
|
||||
Return the status of active ML models.
|
||||
Return the status of active ML models
|
||||
"""
|
||||
models_loaded = getattr(request.app.state, "models_loaded", False)
|
||||
detection_mode = getattr(request.app.state, "detection_mode", "rules")
|
||||
|
||||
active_models: list[dict[str, object]] = []
|
||||
session_factory = getattr(request.app.state, "session_factory", None)
|
||||
if session_factory is not None:
|
||||
async with session_factory() as session:
|
||||
active_models = await _get_active_models(session)
|
||||
|
||||
return {
|
||||
"active_models": [],
|
||||
"detection_mode": "rules-only",
|
||||
"note": "ML models available after Phase 2 training",
|
||||
"models_loaded": models_loaded,
|
||||
"detection_mode": detection_mode,
|
||||
"active_models": active_models,
|
||||
}
|
||||
|
||||
|
||||
@router.post("/retrain", status_code=202)
|
||||
async def retrain() -> dict[str, object]:
|
||||
"""
|
||||
Trigger an async model retraining job.
|
||||
Trigger an async model retraining job
|
||||
"""
|
||||
return {
|
||||
"status": "accepted",
|
||||
"job_id": uuid.uuid4().hex,
|
||||
"note": "Retraining not available in Phase 1",
|
||||
}
|
||||
|
||||
|
||||
async def _get_active_models(
|
||||
session: AsyncSession, ) -> list[dict[str, object]]:
|
||||
"""
|
||||
Query all active model metadata records
|
||||
"""
|
||||
query = select(ModelMetadata).where(
|
||||
ModelMetadata.is_active == True # noqa: E712
|
||||
)
|
||||
rows = (await session.execute(query)).scalars().all()
|
||||
return [{
|
||||
"model_type": row.model_type,
|
||||
"version": row.version,
|
||||
"training_samples": row.training_samples,
|
||||
"metrics": row.metrics,
|
||||
"threshold": row.threshold,
|
||||
} for row in rows]
|
||||
|
|
|
|||
|
|
@ -15,8 +15,8 @@ router = APIRouter(prefix="/stats", tags=["stats"])
|
|||
|
||||
@router.get("", response_model=StatsResponse)
|
||||
async def get_stats(
|
||||
session: AsyncSession = Depends(get_session),
|
||||
time_range: str = Query("24h", alias="range"),
|
||||
session: AsyncSession = Depends(get_session),
|
||||
time_range: str = Query("24h", alias="range"),
|
||||
) -> StatsResponse:
|
||||
"""
|
||||
Aggregate threat statistics for a given time window.
|
||||
|
|
|
|||
|
|
@ -18,13 +18,13 @@ router = APIRouter(prefix="/threats", tags=["threats"])
|
|||
|
||||
@router.get("", response_model=ThreatListResponse)
|
||||
async def list_threats(
|
||||
session: AsyncSession = Depends(get_session),
|
||||
limit: int = Query(50, ge=1, le=100),
|
||||
offset: int = Query(0, ge=0),
|
||||
severity: str | None = Query(None),
|
||||
source_ip: str | None = Query(None),
|
||||
since: datetime | None = Query(None),
|
||||
until: datetime | None = Query(None),
|
||||
session: AsyncSession = Depends(get_session),
|
||||
limit: int = Query(50, ge=1, le=100),
|
||||
offset: int = Query(0, ge=0),
|
||||
severity: str | None = Query(None),
|
||||
source_ip: str | None = Query(None),
|
||||
since: datetime | None = Query(None),
|
||||
until: datetime | None = Query(None),
|
||||
) -> ThreatListResponse:
|
||||
"""
|
||||
List threat events with optional filters and pagination.
|
||||
|
|
@ -42,8 +42,8 @@ async def list_threats(
|
|||
|
||||
@router.get("/{threat_id}", response_model=ThreatEventResponse)
|
||||
async def get_threat(
|
||||
threat_id: uuid.UUID,
|
||||
session: AsyncSession = Depends(get_session),
|
||||
threat_id: uuid.UUID,
|
||||
session: AsyncSession = Depends(get_session),
|
||||
) -> ThreatEventResponse:
|
||||
"""
|
||||
Fetch a single threat event by ID.
|
||||
|
|
|
|||
|
|
@ -6,9 +6,13 @@ dispatcher.py
|
|||
import logging
|
||||
|
||||
import redis.asyncio as aioredis
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
||||
from sqlalchemy.ext.asyncio import (
|
||||
AsyncSession,
|
||||
async_sessionmaker,
|
||||
)
|
||||
|
||||
from app.core.alerts import ALERTS_CHANNEL
|
||||
from app.core.detection.ensemble import classify_severity
|
||||
from app.core.ingestion.pipeline import ScoredRequest
|
||||
from app.schemas.websocket import WebSocketAlert
|
||||
from app.services.threat_service import create_threat_event
|
||||
|
|
@ -18,11 +22,12 @@ logger = logging.getLogger(__name__)
|
|||
|
||||
class AlertDispatcher:
|
||||
"""
|
||||
Routes scored threat events to storage, pub/sub, and structured logging.
|
||||
Routes scored threat events to storage, pub/sub,
|
||||
and structured logging
|
||||
|
||||
MEDIUM+ severity events are persisted to PostgreSQL and published
|
||||
to the Redis alerts channel for WebSocket relay.
|
||||
All events are logged to stdout.
|
||||
MEDIUM+ severity events are persisted to PostgreSQL
|
||||
and published to the Redis alerts channel for
|
||||
WebSocket relay. All events are logged to stdout.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
|
|
@ -35,39 +40,49 @@ class AlertDispatcher:
|
|||
|
||||
async def dispatch(self, scored: ScoredRequest) -> None:
|
||||
"""
|
||||
Handle a scored request from the pipeline's dispatch stage.
|
||||
Handle a scored request from the pipeline's
|
||||
dispatch stage
|
||||
"""
|
||||
severity = classify_severity(scored.final_score)
|
||||
|
||||
logger.info(
|
||||
"threat_event severity=%s score=%.2f ip=%s path=%s rules=%s",
|
||||
scored.rule_result.severity,
|
||||
scored.rule_result.threat_score,
|
||||
"threat_event severity=%s score=%.2f mode=%s ip=%s path=%s rules=%s",
|
||||
severity,
|
||||
scored.final_score,
|
||||
scored.detection_mode,
|
||||
scored.entry.ip,
|
||||
scored.entry.path,
|
||||
scored.rule_result.matched_rules,
|
||||
)
|
||||
|
||||
if scored.rule_result.severity in ("HIGH", "MEDIUM"):
|
||||
if severity in ("HIGH", "MEDIUM"):
|
||||
await self._store_event(scored)
|
||||
await self._publish_alert(scored)
|
||||
await self._publish_alert(scored, severity)
|
||||
|
||||
async def _store_event(self, scored: ScoredRequest) -> None:
|
||||
"""
|
||||
Persist the scored request as a threat event in PostgreSQL.
|
||||
Persist the scored request as a threat event
|
||||
in PostgreSQL
|
||||
"""
|
||||
async with self._session_factory() as session:
|
||||
await create_threat_event(session, scored)
|
||||
await session.commit()
|
||||
|
||||
async def _publish_alert(self, scored: ScoredRequest) -> None:
|
||||
async def _publish_alert(
|
||||
self,
|
||||
scored: ScoredRequest,
|
||||
severity: str,
|
||||
) -> None:
|
||||
"""
|
||||
Publish a real-time alert to the Redis pub/sub channel.
|
||||
Publish a real-time alert to the Redis pub/sub
|
||||
channel
|
||||
"""
|
||||
alert = WebSocketAlert(
|
||||
timestamp=scored.entry.timestamp,
|
||||
source_ip=scored.entry.ip,
|
||||
request_path=scored.entry.path,
|
||||
threat_score=scored.rule_result.threat_score,
|
||||
severity=scored.rule_result.severity,
|
||||
threat_score=scored.final_score,
|
||||
severity=severity,
|
||||
component_scores=scored.rule_result.component_scores,
|
||||
)
|
||||
await self._redis.publish(ALERTS_CHANNEL, alert.model_dump_json())
|
||||
|
|
|
|||
|
|
@ -0,0 +1,66 @@
|
|||
"""
|
||||
©AngelaMos | 2026
|
||||
ensemble.py
|
||||
"""
|
||||
|
||||
|
||||
def normalize_ae_score(error: float, threshold: float) -> float:
|
||||
"""
|
||||
Normalize autoencoder reconstruction error to [0, 1]
|
||||
"""
|
||||
if threshold <= 0:
|
||||
return 0.0
|
||||
return min(error / (threshold * 2), 1.0)
|
||||
|
||||
|
||||
def normalize_if_score(raw_score: float) -> float:
|
||||
"""
|
||||
Normalize isolation forest score to [0, 1]
|
||||
|
||||
sklearn IF returns negative scores for anomalies,
|
||||
positive for normal samples
|
||||
"""
|
||||
return (1 - raw_score) / 2.0
|
||||
|
||||
|
||||
def fuse_scores(
|
||||
scores: dict[str, float],
|
||||
weights: dict[str, float],
|
||||
) -> float:
|
||||
"""
|
||||
Weighted average of per-model normalized scores
|
||||
"""
|
||||
total = 0.0
|
||||
weight_sum = 0.0
|
||||
for key, weight in weights.items():
|
||||
if key in scores:
|
||||
total += scores[key] * weight
|
||||
weight_sum += weight
|
||||
if weight_sum == 0:
|
||||
return 0.0
|
||||
return total / weight_sum * weight_sum
|
||||
|
||||
|
||||
def blend_scores(
|
||||
ml_score: float,
|
||||
rule_score: float,
|
||||
ml_weight: float = 0.7,
|
||||
) -> float:
|
||||
"""
|
||||
Blend ML ensemble score with rule engine score
|
||||
"""
|
||||
return min(
|
||||
ml_score * ml_weight + rule_score * (1.0 - ml_weight),
|
||||
1.0,
|
||||
)
|
||||
|
||||
|
||||
def classify_severity(score: float) -> str:
|
||||
"""
|
||||
Map a unified threat score to a severity label
|
||||
"""
|
||||
if score >= 0.7:
|
||||
return "HIGH"
|
||||
if score >= 0.5:
|
||||
return "MEDIUM"
|
||||
return "LOW"
|
||||
|
|
@ -0,0 +1,153 @@
|
|||
"""
|
||||
©AngelaMos | 2026
|
||||
inference.py
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
try:
|
||||
import onnxruntime as ort
|
||||
except ImportError:
|
||||
ort = None # type: ignore[assignment]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
AE_FILENAME = "ae.onnx"
|
||||
RF_FILENAME = "rf.onnx"
|
||||
IF_FILENAME = "if.onnx"
|
||||
SCALER_FILENAME = "scaler.json"
|
||||
THRESHOLD_FILENAME = "threshold.json"
|
||||
|
||||
|
||||
class InferenceEngine:
|
||||
"""
|
||||
ONNX-based inference engine for the 3-model ML ensemble.
|
||||
|
||||
Loads autoencoder, random forest, and isolation forest ONNX sessions
|
||||
from a model directory. Returns None for predictions when models
|
||||
are not available.
|
||||
"""
|
||||
|
||||
def __init__(self, model_dir: str) -> None:
|
||||
self._ae_session: ort.InferenceSession | None = None # type: ignore[union-attr]
|
||||
self._rf_session: ort.InferenceSession | None = None # type: ignore[union-attr]
|
||||
self._if_session: ort.InferenceSession | None = None # type: ignore[union-attr]
|
||||
self._scaler_center: np.ndarray | None = None
|
||||
self._scaler_scale: np.ndarray | None = None
|
||||
self._threshold: float = 0.0
|
||||
self._loaded = False
|
||||
|
||||
if ort is None:
|
||||
logger.warning("onnxruntime not installed")
|
||||
return
|
||||
|
||||
model_path = Path(model_dir)
|
||||
ae_path = model_path / AE_FILENAME
|
||||
rf_path = model_path / RF_FILENAME
|
||||
if_path = model_path / IF_FILENAME
|
||||
scaler_path = model_path / SCALER_FILENAME
|
||||
threshold_path = model_path / THRESHOLD_FILENAME
|
||||
|
||||
required = [ae_path, rf_path, if_path, scaler_path, threshold_path]
|
||||
if not all(p.exists() for p in required):
|
||||
return
|
||||
|
||||
try:
|
||||
opts = ort.SessionOptions()
|
||||
opts.inter_op_num_threads = 1
|
||||
opts.intra_op_num_threads = 1
|
||||
|
||||
self._ae_session = ort.InferenceSession(str(ae_path), opts)
|
||||
self._rf_session = ort.InferenceSession(str(rf_path), opts)
|
||||
self._if_session = ort.InferenceSession(str(if_path), opts)
|
||||
|
||||
scaler_data = json.loads(scaler_path.read_text())
|
||||
self._scaler_center = np.array(scaler_data["center"],
|
||||
dtype=np.float32)
|
||||
self._scaler_scale = np.array(scaler_data["scale"],
|
||||
dtype=np.float32)
|
||||
|
||||
threshold_data = json.loads(threshold_path.read_text())
|
||||
self._threshold = float(threshold_data["threshold"])
|
||||
|
||||
self._loaded = True
|
||||
logger.info("Loaded 3 ONNX models from %s", model_dir)
|
||||
except Exception:
|
||||
logger.exception("Failed to load ONNX models from %s", model_dir)
|
||||
|
||||
@property
|
||||
def is_loaded(self) -> bool:
|
||||
"""
|
||||
Whether all 3 models are loaded and ready for inference
|
||||
"""
|
||||
return self._loaded
|
||||
|
||||
@property
|
||||
def threshold(self) -> float:
|
||||
"""
|
||||
Autoencoder anomaly detection threshold
|
||||
"""
|
||||
return self._threshold
|
||||
|
||||
def predict(self, batch: np.ndarray) -> dict[str, list[float]] | None:
|
||||
"""
|
||||
Run all 3 models on a batch of feature vectors.
|
||||
|
||||
Returns per-model raw scores for ensemble fusion, or None
|
||||
if models are not loaded.
|
||||
"""
|
||||
if not self._loaded:
|
||||
return None
|
||||
|
||||
ae_input = self._scale_for_ae(batch)
|
||||
ae_reconstructed = self._ae_session.run( # type: ignore[union-attr]
|
||||
None, {"features": ae_input})[0]
|
||||
ae_errors = np.mean((ae_input - ae_reconstructed)**2, axis=1)
|
||||
|
||||
rf_result = self._rf_session.run( # type: ignore[union-attr]
|
||||
None, {"features": batch})
|
||||
rf_proba = self._extract_rf_proba(rf_result[1])
|
||||
|
||||
if_scores_raw = self._if_session.run( # type: ignore[union-attr]
|
||||
None, {"features": batch})[1].flatten()
|
||||
|
||||
return {
|
||||
"ae": ae_errors.tolist(),
|
||||
"rf": rf_proba.tolist(),
|
||||
"if": if_scores_raw.tolist(),
|
||||
}
|
||||
|
||||
def _scale_for_ae(self, batch: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Apply RobustScaler transform for autoencoder input
|
||||
"""
|
||||
if self._scaler_center is None or self._scaler_scale is None:
|
||||
return batch
|
||||
return (batch - self._scaler_center) / self._scaler_scale
|
||||
|
||||
@staticmethod
|
||||
def _extract_rf_proba(
|
||||
ort_output: list | np.ndarray
|
||||
) -> np.ndarray: # type: ignore[type-arg]
|
||||
"""
|
||||
Extract attack probability from skl2onnx RF output format.
|
||||
|
||||
skl2onnx outputs a list of dicts [{0: p0, 1: p1}, ...] for
|
||||
probability output.
|
||||
"""
|
||||
if isinstance(ort_output, np.ndarray):
|
||||
if ort_output.ndim == 2 and ort_output.shape[1] >= 2:
|
||||
return ort_output[:, 1].astype(np.float32)
|
||||
return ort_output.flatten().astype(np.float32)
|
||||
|
||||
proba = []
|
||||
for row in ort_output:
|
||||
if isinstance(row, dict):
|
||||
proba.append(float(row.get(1, 0.0)))
|
||||
else:
|
||||
proba.append(float(row))
|
||||
return np.array(proba, dtype=np.float32)
|
||||
|
|
@ -40,7 +40,8 @@ class GeoIPService:
|
|||
self._reader = geoip2.database.Reader(db_path)
|
||||
logger.info("GeoIP database loaded from %s", db_path)
|
||||
else:
|
||||
logger.warning("GeoIP database not found at %s — lookups disabled", db_path)
|
||||
logger.warning("GeoIP database not found at %s — lookups disabled",
|
||||
db_path)
|
||||
|
||||
async def lookup(self, ip: str) -> GeoResult | None:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -171,21 +171,21 @@ def _path_depth_variance(depth_members: list[str]) -> float:
|
|||
return 0.0
|
||||
depths = [int(m.split(":")[0]) for m in depth_members]
|
||||
mean = sum(depths) / len(depths)
|
||||
return sum((d - mean) ** 2 for d in depths) / len(depths)
|
||||
return sum((d - mean)**2 for d in depths) / len(depths)
|
||||
|
||||
|
||||
def _inter_request_time_stats(
|
||||
entries: list[tuple[str, float]],
|
||||
) -> tuple[float, float]:
|
||||
entries: list[tuple[str, float]], ) -> tuple[float, float]:
|
||||
"""
|
||||
Mean and standard deviation of inter-request intervals in milliseconds.
|
||||
"""
|
||||
if len(entries) < 2:
|
||||
return 0.0, 0.0
|
||||
timestamps = sorted(score for _, score in entries)
|
||||
deltas = [(timestamps[i + 1] - timestamps[i]) * 1000 for i in range(len(timestamps) - 1)]
|
||||
deltas = [(timestamps[i + 1] - timestamps[i]) * 1000
|
||||
for i in range(len(timestamps) - 1)]
|
||||
mean = sum(deltas) / len(deltas)
|
||||
if len(deltas) < 2:
|
||||
return mean, 0.0
|
||||
variance = sum((d - mean) ** 2 for d in deltas) / len(deltas)
|
||||
variance = sum((d - mean)**2 for d in deltas) / len(deltas)
|
||||
return mean, math.sqrt(variance)
|
||||
|
|
|
|||
|
|
@ -20,8 +20,7 @@ def _encode_country(code: str) -> float:
|
|||
|
||||
|
||||
def encode_for_inference(
|
||||
features: dict[str, int | float | bool | str],
|
||||
) -> list[float]:
|
||||
features: dict[str, int | float | bool | str], ) -> list[float]:
|
||||
"""
|
||||
Encode a combined feature dict into a 35-element float vector
|
||||
matching the model input specification.
|
||||
|
|
|
|||
|
|
@ -56,27 +56,50 @@ def extract_request_features(
|
|||
_, ext = splitext(entry.path)
|
||||
|
||||
return {
|
||||
"http_method": entry.method,
|
||||
"path_depth": len([s for s in entry.path.split("/") if s]),
|
||||
"path_entropy": _shannon_entropy(entry.path),
|
||||
"path_length": path_len,
|
||||
"query_string_length": len(entry.query_string),
|
||||
"query_param_count": (len(entry.query_string.split("&")) if entry.query_string else 0),
|
||||
"has_encoded_chars": bool(ENCODED_CHARS.search(full_uri)),
|
||||
"has_double_encoding": bool(DOUBLE_ENCODED.search(full_uri)),
|
||||
"status_code": entry.status_code,
|
||||
"status_class": f"{entry.status_code // 100}xx",
|
||||
"response_size": entry.response_size,
|
||||
"hour_of_day": entry.timestamp.hour,
|
||||
"day_of_week": entry.timestamp.weekday(),
|
||||
"is_weekend": entry.timestamp.weekday() >= 5,
|
||||
"ua_length": len(entry.user_agent),
|
||||
"ua_entropy": _shannon_entropy(entry.user_agent),
|
||||
"is_known_bot": any(sig in ua_lower for sig in BOT_USER_AGENTS),
|
||||
"is_known_scanner": any(sig in ua_lower for sig in SCANNER_USER_AGENTS),
|
||||
"has_attack_pattern": bool(ATTACK_COMBINED.search(full_uri)),
|
||||
"special_char_ratio": non_alnum / path_len if path_len else 0.0,
|
||||
"file_extension": ext,
|
||||
"country_code": country_code,
|
||||
"is_private_ip": _is_private_ip(entry.ip),
|
||||
"http_method":
|
||||
entry.method,
|
||||
"path_depth":
|
||||
len([s for s in entry.path.split("/") if s]),
|
||||
"path_entropy":
|
||||
_shannon_entropy(entry.path),
|
||||
"path_length":
|
||||
path_len,
|
||||
"query_string_length":
|
||||
len(entry.query_string),
|
||||
"query_param_count":
|
||||
(len(entry.query_string.split("&")) if entry.query_string else 0),
|
||||
"has_encoded_chars":
|
||||
bool(ENCODED_CHARS.search(full_uri)),
|
||||
"has_double_encoding":
|
||||
bool(DOUBLE_ENCODED.search(full_uri)),
|
||||
"status_code":
|
||||
entry.status_code,
|
||||
"status_class":
|
||||
f"{entry.status_code // 100}xx",
|
||||
"response_size":
|
||||
entry.response_size,
|
||||
"hour_of_day":
|
||||
entry.timestamp.hour,
|
||||
"day_of_week":
|
||||
entry.timestamp.weekday(),
|
||||
"is_weekend":
|
||||
entry.timestamp.weekday() >= 5,
|
||||
"ua_length":
|
||||
len(entry.user_agent),
|
||||
"ua_entropy":
|
||||
_shannon_entropy(entry.user_agent),
|
||||
"is_known_bot":
|
||||
any(sig in ua_lower for sig in BOT_USER_AGENTS),
|
||||
"is_known_scanner":
|
||||
any(sig in ua_lower for sig in SCANNER_USER_AGENTS),
|
||||
"has_attack_pattern":
|
||||
bool(ATTACK_COMBINED.search(full_uri)),
|
||||
"special_char_ratio":
|
||||
non_alnum / path_len if path_len else 0.0,
|
||||
"file_extension":
|
||||
ext,
|
||||
"country_code":
|
||||
country_code,
|
||||
"is_private_ip":
|
||||
_is_private_ip(entry.ip),
|
||||
}
|
||||
|
|
|
|||
|
|
@ -93,14 +93,12 @@ CATEGORICAL_ENCODERS: dict[str, dict[str, int]] = {
|
|||
"file_extension": EXTENSION_MAP,
|
||||
}
|
||||
|
||||
BOOLEAN_FEATURES: frozenset[str] = frozenset(
|
||||
{
|
||||
"has_encoded_chars",
|
||||
"has_double_encoding",
|
||||
"is_weekend",
|
||||
"is_known_bot",
|
||||
"is_known_scanner",
|
||||
"has_attack_pattern",
|
||||
"is_private_ip",
|
||||
}
|
||||
)
|
||||
BOOLEAN_FEATURES: frozenset[str] = frozenset({
|
||||
"has_encoded_chars",
|
||||
"has_double_encoding",
|
||||
"is_weekend",
|
||||
"is_known_bot",
|
||||
"is_known_scanner",
|
||||
"has_attack_pattern",
|
||||
"is_private_ip",
|
||||
})
|
||||
|
|
|
|||
|
|
@ -8,27 +8,25 @@ import re
|
|||
ENCODED_CHARS = re.compile(r"%[0-9a-fA-F]{2}")
|
||||
DOUBLE_ENCODED = re.compile(r"%25[0-9a-fA-F]{2}")
|
||||
|
||||
_SQLI = (
|
||||
r"(?:union\s+(?:all\s+)?select|"
|
||||
r"'\s*or\s+.+=|"
|
||||
r"'\s*and\s+.+=|"
|
||||
r"1\s*=\s*1|"
|
||||
r"sleep\s*\(|"
|
||||
r"benchmark\s*\(|"
|
||||
r"waitfor\s+delay|"
|
||||
r"extractvalue\s*\(|"
|
||||
r"updatexml\s*\(|"
|
||||
r"load_file\s*\(|"
|
||||
r"into\s+(?:out|dump)file|"
|
||||
r"group\s+by\s+.+having|"
|
||||
r"order\s+by\s+\d+|"
|
||||
r"(?:drop|alter|create)\s+table|"
|
||||
r"information_schema|"
|
||||
r"(?:char|concat|hex|unhex)\s*\(|"
|
||||
r"0x[0-9a-f]{6,}|"
|
||||
r"--\s*$|"
|
||||
r"/\*.*?\*/)"
|
||||
)
|
||||
_SQLI = (r"(?:union\s+(?:all\s+)?select|"
|
||||
r"'\s*or\s+.+=|"
|
||||
r"'\s*and\s+.+=|"
|
||||
r"1\s*=\s*1|"
|
||||
r"sleep\s*\(|"
|
||||
r"benchmark\s*\(|"
|
||||
r"waitfor\s+delay|"
|
||||
r"extractvalue\s*\(|"
|
||||
r"updatexml\s*\(|"
|
||||
r"load_file\s*\(|"
|
||||
r"into\s+(?:out|dump)file|"
|
||||
r"group\s+by\s+.+having|"
|
||||
r"order\s+by\s+\d+|"
|
||||
r"(?:drop|alter|create)\s+table|"
|
||||
r"information_schema|"
|
||||
r"(?:char|concat|hex|unhex)\s*\(|"
|
||||
r"0x[0-9a-f]{6,}|"
|
||||
r"--\s*$|"
|
||||
r"/\*.*?\*/)")
|
||||
|
||||
_XSS = (
|
||||
r"(?:<\s*script|"
|
||||
|
|
@ -45,23 +43,20 @@ _XSS = (
|
|||
r"prompt\s*\(|"
|
||||
r"confirm\s*\(|"
|
||||
r"expression\s*\(|"
|
||||
r"String\s*\.\s*fromCharCode)"
|
||||
)
|
||||
r"String\s*\.\s*fromCharCode)")
|
||||
|
||||
_PATH_TRAVERSAL = (
|
||||
r"(?:\.\./|"
|
||||
r"\.\.\\|"
|
||||
r"%2e%2e[%/\\]|"
|
||||
r"%252e%252e|"
|
||||
r"(?:etc/(?:passwd|shadow|hosts)|"
|
||||
r"proc/self/|"
|
||||
r"windows/system32|"
|
||||
r"boot\.ini|"
|
||||
r"web\.config|"
|
||||
r"\.env|"
|
||||
r"\.git/config|"
|
||||
r"wp-config\.php))"
|
||||
)
|
||||
_PATH_TRAVERSAL = (r"(?:\.\./|"
|
||||
r"\.\.\\|"
|
||||
r"%2e%2e[%/\\]|"
|
||||
r"%252e%252e|"
|
||||
r"(?:etc/(?:passwd|shadow|hosts)|"
|
||||
r"proc/self/|"
|
||||
r"windows/system32|"
|
||||
r"boot\.ini|"
|
||||
r"web\.config|"
|
||||
r"\.env|"
|
||||
r"\.git/config|"
|
||||
r"wp-config\.php))")
|
||||
|
||||
_COMMAND_INJECTION = (
|
||||
r"(?:;\s*(?:ls|cat|rm|wget|curl|chmod|chown|nc|bash|sh|python|perl|ruby|php)\b|"
|
||||
|
|
@ -70,19 +65,16 @@ _COMMAND_INJECTION = (
|
|||
r"`[^`]+`|"
|
||||
r"\$\{|"
|
||||
r">\s*/(?:etc|tmp|var)|"
|
||||
r"&&\s*(?:cat|ls|id|whoami|wget|curl)\b)"
|
||||
)
|
||||
r"&&\s*(?:cat|ls|id|whoami|wget|curl)\b)")
|
||||
|
||||
_FILE_INCLUSION = (
|
||||
r"(?:php://|"
|
||||
r"file://|"
|
||||
r"data://|"
|
||||
r"expect://|"
|
||||
r"input://|"
|
||||
r"zip://|"
|
||||
r"phar://|"
|
||||
r"glob://)"
|
||||
)
|
||||
_FILE_INCLUSION = (r"(?:php://|"
|
||||
r"file://|"
|
||||
r"data://|"
|
||||
r"expect://|"
|
||||
r"input://|"
|
||||
r"zip://|"
|
||||
r"phar://|"
|
||||
r"glob://)")
|
||||
|
||||
SQLI = re.compile(_SQLI, re.IGNORECASE)
|
||||
XSS = re.compile(_XSS, re.IGNORECASE)
|
||||
|
|
@ -91,6 +83,7 @@ COMMAND_INJECTION = re.compile(_COMMAND_INJECTION, re.IGNORECASE)
|
|||
FILE_INCLUSION = re.compile(_FILE_INCLUSION, re.IGNORECASE)
|
||||
|
||||
ATTACK_COMBINED = re.compile(
|
||||
r"|".join((_SQLI, _XSS, _PATH_TRAVERSAL, _COMMAND_INJECTION, _FILE_INCLUSION)),
|
||||
r"|".join(
|
||||
(_SQLI, _XSS, _PATH_TRAVERSAL, _COMMAND_INJECTION, _FILE_INCLUSION)),
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -3,87 +3,83 @@
|
|||
signatures.py
|
||||
"""
|
||||
|
||||
BOT_USER_AGENTS: frozenset[str] = frozenset(
|
||||
{
|
||||
"googlebot",
|
||||
"bingbot",
|
||||
"slurp",
|
||||
"duckduckbot",
|
||||
"baiduspider",
|
||||
"yandexbot",
|
||||
"sogou",
|
||||
"facebot",
|
||||
"ia_archiver",
|
||||
"applebot",
|
||||
"petalbot",
|
||||
"semrushbot",
|
||||
"ahrefsbot",
|
||||
"mj12bot",
|
||||
"dotbot",
|
||||
"rogerbot",
|
||||
"linkedinbot",
|
||||
"twitterbot",
|
||||
"gptbot",
|
||||
"claudebot",
|
||||
"amazonbot",
|
||||
"bytespider",
|
||||
"ccbot",
|
||||
"dataforseo",
|
||||
"seznambot",
|
||||
"megaindex",
|
||||
"blexbot",
|
||||
"exabot",
|
||||
"archive.org_bot",
|
||||
"mojeekbot",
|
||||
"uptimerobot",
|
||||
"deadlinkchecker",
|
||||
"sitebulb",
|
||||
"screaming frog",
|
||||
}
|
||||
)
|
||||
BOT_USER_AGENTS: frozenset[str] = frozenset({
|
||||
"googlebot",
|
||||
"bingbot",
|
||||
"slurp",
|
||||
"duckduckbot",
|
||||
"baiduspider",
|
||||
"yandexbot",
|
||||
"sogou",
|
||||
"facebot",
|
||||
"ia_archiver",
|
||||
"applebot",
|
||||
"petalbot",
|
||||
"semrushbot",
|
||||
"ahrefsbot",
|
||||
"mj12bot",
|
||||
"dotbot",
|
||||
"rogerbot",
|
||||
"linkedinbot",
|
||||
"twitterbot",
|
||||
"gptbot",
|
||||
"claudebot",
|
||||
"amazonbot",
|
||||
"bytespider",
|
||||
"ccbot",
|
||||
"dataforseo",
|
||||
"seznambot",
|
||||
"megaindex",
|
||||
"blexbot",
|
||||
"exabot",
|
||||
"archive.org_bot",
|
||||
"mojeekbot",
|
||||
"uptimerobot",
|
||||
"deadlinkchecker",
|
||||
"sitebulb",
|
||||
"screaming frog",
|
||||
})
|
||||
|
||||
SCANNER_USER_AGENTS: frozenset[str] = frozenset(
|
||||
{
|
||||
"nikto",
|
||||
"sqlmap",
|
||||
"nessus",
|
||||
"openvas",
|
||||
"acunetix",
|
||||
"w3af",
|
||||
"nmap",
|
||||
"masscan",
|
||||
"zgrab",
|
||||
"gobuster",
|
||||
"dirbuster",
|
||||
"dirb",
|
||||
"wfuzz",
|
||||
"ffuf",
|
||||
"nuclei",
|
||||
"burp",
|
||||
"zap",
|
||||
"arachni",
|
||||
"skipfish",
|
||||
"wpscan",
|
||||
"joomscan",
|
||||
"whatweb",
|
||||
"httprint",
|
||||
"fierce",
|
||||
"subfinder",
|
||||
"amass",
|
||||
"httpx",
|
||||
"jaeles",
|
||||
"xray",
|
||||
"gau",
|
||||
"hakrawler",
|
||||
"katana",
|
||||
"cariddi",
|
||||
"gospider",
|
||||
"feroxbuster",
|
||||
"rustbuster",
|
||||
"patator",
|
||||
"hydra",
|
||||
"medusa",
|
||||
"metasploit",
|
||||
"cobalt",
|
||||
}
|
||||
)
|
||||
SCANNER_USER_AGENTS: frozenset[str] = frozenset({
|
||||
"nikto",
|
||||
"sqlmap",
|
||||
"nessus",
|
||||
"openvas",
|
||||
"acunetix",
|
||||
"w3af",
|
||||
"nmap",
|
||||
"masscan",
|
||||
"zgrab",
|
||||
"gobuster",
|
||||
"dirbuster",
|
||||
"dirb",
|
||||
"wfuzz",
|
||||
"ffuf",
|
||||
"nuclei",
|
||||
"burp",
|
||||
"zap",
|
||||
"arachni",
|
||||
"skipfish",
|
||||
"wpscan",
|
||||
"joomscan",
|
||||
"whatweb",
|
||||
"httprint",
|
||||
"fierce",
|
||||
"subfinder",
|
||||
"amass",
|
||||
"httpx",
|
||||
"jaeles",
|
||||
"xray",
|
||||
"gau",
|
||||
"hakrawler",
|
||||
"katana",
|
||||
"cariddi",
|
||||
"gospider",
|
||||
"feroxbuster",
|
||||
"rustbuster",
|
||||
"patator",
|
||||
"hydra",
|
||||
"medusa",
|
||||
"metasploit",
|
||||
"cobalt",
|
||||
})
|
||||
|
|
|
|||
|
|
@ -28,15 +28,13 @@ class ParsedLogEntry:
|
|||
|
||||
_TIMESTAMP_FMT = "%d/%b/%Y:%H:%M:%S %z"
|
||||
|
||||
_COMBINED_RE = re.compile(
|
||||
r"(?P<ip>\S+) \S+ \S+ "
|
||||
r"\[(?P<timestamp>[^\]]+)\] "
|
||||
r'"(?P<request>[^"]*)" '
|
||||
r"(?P<status>\d{3}) "
|
||||
r"(?P<size>\S+) "
|
||||
r'"(?P<referer>[^"]*)" '
|
||||
r'"(?P<user_agent>[^"]*)"'
|
||||
)
|
||||
_COMBINED_RE = re.compile(r"(?P<ip>\S+) \S+ \S+ "
|
||||
r"\[(?P<timestamp>[^\]]+)\] "
|
||||
r'"(?P<request>[^"]*)" '
|
||||
r"(?P<status>\d{3}) "
|
||||
r"(?P<size>\S+) "
|
||||
r'"(?P<referer>[^"]*)" '
|
||||
r'"(?P<user_agent>[^"]*)"')
|
||||
|
||||
|
||||
def parse_combined(line: str) -> ParsedLogEntry | None:
|
||||
|
|
@ -72,7 +70,8 @@ def _parse_split(line: str) -> ParsedLogEntry | None:
|
|||
bracket_open = prefix.index("[")
|
||||
bracket_close = prefix.index("]")
|
||||
ip = prefix[:bracket_open].split()[0]
|
||||
timestamp = datetime.strptime(prefix[bracket_open + 1 : bracket_close], _TIMESTAMP_FMT)
|
||||
timestamp = datetime.strptime(prefix[bracket_open + 1:bracket_close],
|
||||
_TIMESTAMP_FMT)
|
||||
|
||||
request_parts = request_line.split(" ", 2)
|
||||
method = request_parts[0]
|
||||
|
|
|
|||
|
|
@ -8,9 +8,16 @@ import logging
|
|||
import uuid
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import redis.asyncio as aioredis
|
||||
|
||||
from app.core.detection.ensemble import (
|
||||
blend_scores,
|
||||
fuse_scores,
|
||||
normalize_ae_score,
|
||||
normalize_if_score,
|
||||
)
|
||||
from app.core.detection.rules import RuleEngine, RuleResult
|
||||
from app.core.enrichment.geoip import GeoIPService, GeoResult
|
||||
from app.core.features.aggregator import WindowAggregator
|
||||
|
|
@ -18,13 +25,27 @@ from app.core.features.encoder import encode_for_inference
|
|||
from app.core.features.extractor import extract_request_features
|
||||
from app.core.ingestion.parsers import ParsedLogEntry, parse_combined
|
||||
|
||||
try:
|
||||
import numpy as np
|
||||
except ImportError:
|
||||
np = None # type: ignore[assignment]
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from app.core.detection.inference import InferenceEngine
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_ENSEMBLE_WEIGHTS: dict[str, float] = {
|
||||
"ae": 0.40,
|
||||
"rf": 0.40,
|
||||
"if": 0.20,
|
||||
}
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class EnrichedRequest:
|
||||
"""
|
||||
A parsed log entry enriched with extracted features and GeoIP data.
|
||||
A parsed log entry enriched with extracted features and GeoIP data
|
||||
"""
|
||||
|
||||
entry: ParsedLogEntry
|
||||
|
|
@ -36,7 +57,7 @@ class EnrichedRequest:
|
|||
@dataclass(slots=True)
|
||||
class ScoredRequest:
|
||||
"""
|
||||
A fully scored request ready for dispatch.
|
||||
A fully scored request ready for dispatch
|
||||
"""
|
||||
|
||||
entry: ParsedLogEntry
|
||||
|
|
@ -44,16 +65,18 @@ class ScoredRequest:
|
|||
feature_vector: list[float]
|
||||
geo: GeoResult | None
|
||||
rule_result: RuleResult
|
||||
final_score: float = 0.0
|
||||
detection_mode: str = "rules"
|
||||
|
||||
|
||||
class Pipeline:
|
||||
"""
|
||||
Four-stage async pipeline that transforms raw log lines
|
||||
into scored threat candidates.
|
||||
into scored threat candidates
|
||||
|
||||
Stages:
|
||||
raw_queue → [parse] → parsed_queue → [enrich+features]
|
||||
→ feature_queue → [detect] → alert_queue → [dispatch]
|
||||
raw_queue -> [parse] -> parsed_queue -> [enrich+features]
|
||||
-> feature_queue -> [detect] -> alert_queue -> [dispatch]
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
|
|
@ -61,34 +84,36 @@ class Pipeline:
|
|||
redis_client: aioredis.Redis[bytes],
|
||||
rule_engine: RuleEngine,
|
||||
geoip: GeoIPService | None = None,
|
||||
on_result: Callable[[ScoredRequest], Awaitable[None]] | None = None,
|
||||
on_result: (Callable[[ScoredRequest], Awaitable[None]] | None) = None,
|
||||
inference_engine: InferenceEngine | None = None,
|
||||
ensemble_weights: dict[str, float] | None = None,
|
||||
raw_queue_size: int = 1000,
|
||||
parsed_queue_size: int = 500,
|
||||
feature_queue_size: int = 200,
|
||||
alert_queue_size: int = 100,
|
||||
) -> None:
|
||||
self.raw_queue: asyncio.Queue[str | None] = asyncio.Queue(
|
||||
maxsize=raw_queue_size,
|
||||
)
|
||||
self._parsed_queue: asyncio.Queue[ParsedLogEntry | None] = asyncio.Queue(
|
||||
maxsize=parsed_queue_size,
|
||||
)
|
||||
self._feature_queue: asyncio.Queue[EnrichedRequest | None] = asyncio.Queue(
|
||||
maxsize=feature_queue_size,
|
||||
)
|
||||
maxsize=raw_queue_size)
|
||||
self._parsed_queue: asyncio.Queue[ParsedLogEntry
|
||||
| None] = asyncio.Queue(
|
||||
maxsize=parsed_queue_size)
|
||||
self._feature_queue: asyncio.Queue[EnrichedRequest
|
||||
| None] = asyncio.Queue(
|
||||
maxsize=feature_queue_size)
|
||||
self._alert_queue: asyncio.Queue[ScoredRequest | None] = asyncio.Queue(
|
||||
maxsize=alert_queue_size,
|
||||
)
|
||||
maxsize=alert_queue_size)
|
||||
|
||||
self._aggregator = WindowAggregator(redis_client)
|
||||
self._rule_engine = rule_engine
|
||||
self._geoip = geoip
|
||||
self._on_result = on_result
|
||||
self._inference_engine = inference_engine
|
||||
self._ensemble_weights = ensemble_weights or DEFAULT_ENSEMBLE_WEIGHTS
|
||||
self._tasks: list[asyncio.Task[None]] = []
|
||||
|
||||
async def start(self) -> None:
|
||||
"""
|
||||
Spawn worker tasks for each pipeline stage.
|
||||
Spawn worker tasks for each pipeline stage
|
||||
"""
|
||||
self._tasks = [
|
||||
asyncio.create_task(self._parse_worker(), name="parse"),
|
||||
|
|
@ -100,7 +125,8 @@ class Pipeline:
|
|||
|
||||
async def stop(self) -> None:
|
||||
"""
|
||||
Send a poison pill through the chain and wait for all workers to exit.
|
||||
Send a poison pill through the chain and wait
|
||||
for all workers to exit
|
||||
"""
|
||||
await self.raw_queue.put(None)
|
||||
await asyncio.gather(*self._tasks)
|
||||
|
|
@ -108,7 +134,7 @@ class Pipeline:
|
|||
|
||||
async def _parse_worker(self) -> None:
|
||||
"""
|
||||
Stage 1: Parse raw log lines into structured entries.
|
||||
Stage 1: Parse raw log lines into structured entries
|
||||
"""
|
||||
while True:
|
||||
line = await self.raw_queue.get()
|
||||
|
|
@ -126,8 +152,9 @@ class Pipeline:
|
|||
|
||||
async def _feature_worker(self) -> None:
|
||||
"""
|
||||
Stage 2: Enrich with GeoIP, extract per-request features,
|
||||
aggregate per-IP windowed features, and encode the 35-dim vector.
|
||||
Stage 2: Enrich with GeoIP, extract per-request
|
||||
features, aggregate per-IP windowed features,
|
||||
and encode the 35-dim vector
|
||||
"""
|
||||
while True:
|
||||
entry = await self._parsed_queue.get()
|
||||
|
|
@ -166,15 +193,18 @@ class Pipeline:
|
|||
features=merged,
|
||||
feature_vector=vector,
|
||||
geo=geo,
|
||||
)
|
||||
)
|
||||
))
|
||||
except Exception:
|
||||
logger.exception("Feature extraction failed for %s", entry.ip)
|
||||
logger.exception(
|
||||
"Feature extraction failed for %s",
|
||||
entry.ip,
|
||||
)
|
||||
self._parsed_queue.task_done()
|
||||
|
||||
async def _detection_worker(self) -> None:
|
||||
"""
|
||||
Stage 3: Score enriched requests using the rule engine.
|
||||
Stage 3: Score enriched requests using the rule
|
||||
engine, optionally enhanced with ML ensemble
|
||||
"""
|
||||
while True:
|
||||
enriched = await self._feature_queue.get()
|
||||
|
|
@ -187,6 +217,25 @@ class Pipeline:
|
|||
enriched.features,
|
||||
enriched.entry,
|
||||
)
|
||||
|
||||
final_score = rule_result.threat_score
|
||||
detection_mode = "rules"
|
||||
|
||||
if (self._inference_engine is not None
|
||||
and self._inference_engine.is_loaded
|
||||
and np is not None):
|
||||
ml_scores = self._score_with_ml(enriched.feature_vector)
|
||||
if ml_scores is not None:
|
||||
ml_fused = fuse_scores(
|
||||
ml_scores,
|
||||
self._ensemble_weights,
|
||||
)
|
||||
final_score = blend_scores(
|
||||
ml_fused,
|
||||
rule_result.threat_score,
|
||||
)
|
||||
detection_mode = "hybrid"
|
||||
|
||||
await self._alert_queue.put(
|
||||
ScoredRequest(
|
||||
entry=enriched.entry,
|
||||
|
|
@ -194,15 +243,40 @@ class Pipeline:
|
|||
feature_vector=enriched.feature_vector,
|
||||
geo=enriched.geo,
|
||||
rule_result=rule_result,
|
||||
)
|
||||
)
|
||||
final_score=final_score,
|
||||
detection_mode=detection_mode,
|
||||
))
|
||||
except Exception:
|
||||
logger.exception("Detection failed")
|
||||
self._feature_queue.task_done()
|
||||
|
||||
def _score_with_ml(self,
|
||||
feature_vector: list[float]) -> dict[str, float] | None:
|
||||
"""
|
||||
Run ML ensemble inference on a single feature vector
|
||||
and return normalized per-model scores
|
||||
"""
|
||||
batch = np.array([feature_vector], dtype=np.float32)
|
||||
raw = self._inference_engine.predict(batch) # type: ignore[union-attr]
|
||||
if raw is None:
|
||||
return None
|
||||
|
||||
return {
|
||||
"ae":
|
||||
normalize_ae_score(
|
||||
raw["ae"][0],
|
||||
self._inference_engine.threshold, # type: ignore[union-attr]
|
||||
),
|
||||
"rf":
|
||||
raw["rf"][0],
|
||||
"if":
|
||||
normalize_if_score(raw["if"][0]),
|
||||
}
|
||||
|
||||
async def _dispatch_worker(self) -> None:
|
||||
"""
|
||||
Stage 4: Dispatch scored results via the on_result callback.
|
||||
Stage 4: Dispatch scored results via the
|
||||
on_result callback
|
||||
"""
|
||||
while True:
|
||||
scored = await self._alert_queue.get()
|
||||
|
|
@ -213,5 +287,8 @@ class Pipeline:
|
|||
if self._on_result is not None:
|
||||
await self._on_result(scored)
|
||||
except Exception:
|
||||
logger.exception("Dispatch failed for %s", scored.entry.ip)
|
||||
logger.exception(
|
||||
"Dispatch failed for %s",
|
||||
scored.entry.ip,
|
||||
)
|
||||
self._alert_queue.task_done()
|
||||
|
|
|
|||
|
|
@ -45,12 +45,14 @@ class _LogHandler(FileSystemEventHandler):
|
|||
Open the target log file and seek to the end.
|
||||
"""
|
||||
try:
|
||||
self._file = open(self._target, encoding="utf-8", errors="replace") # noqa: SIM115
|
||||
self._file = open( # noqa: SIM115
|
||||
self._target, encoding="utf-8", errors="replace")
|
||||
self._file.seek(0, os.SEEK_END)
|
||||
self._inode = os.stat(self._target).st_ino
|
||||
logger.info("Tailing %s (inode %s)", self._target, self._inode)
|
||||
except FileNotFoundError:
|
||||
logger.warning("Log file %s not found — waiting for creation", self._target)
|
||||
logger.warning("Log file %s not found — waiting for creation",
|
||||
self._target)
|
||||
self._file = None
|
||||
self._inode = None
|
||||
|
||||
|
|
@ -64,7 +66,8 @@ class _LogHandler(FileSystemEventHandler):
|
|||
for line in self._file:
|
||||
stripped = line.rstrip("\n\r")
|
||||
if stripped:
|
||||
self._loop.call_soon_threadsafe(self._queue.put_nowait, stripped)
|
||||
self._loop.call_soon_threadsafe(self._queue.put_nowait,
|
||||
stripped)
|
||||
|
||||
def _handle_rotation(self) -> None:
|
||||
"""
|
||||
|
|
@ -76,9 +79,11 @@ class _LogHandler(FileSystemEventHandler):
|
|||
self._file.close()
|
||||
|
||||
try:
|
||||
self._file = open(self._target, encoding="utf-8", errors="replace") # noqa: SIM115
|
||||
self._file = open( # noqa: SIM115
|
||||
self._target, encoding="utf-8", errors="replace")
|
||||
self._inode = os.stat(self._target).st_ino
|
||||
logger.info("Rotated to new %s (inode %s)", self._target, self._inode)
|
||||
logger.info("Rotated to new %s (inode %s)", self._target,
|
||||
self._inode)
|
||||
except FileNotFoundError:
|
||||
self._file = None
|
||||
self._inode = None
|
||||
|
|
@ -93,7 +98,8 @@ class _LogHandler(FileSystemEventHandler):
|
|||
except FileNotFoundError:
|
||||
return False
|
||||
|
||||
def on_modified(self, event: FileModifiedEvent) -> None: # type: ignore[override]
|
||||
def on_modified(
|
||||
self, event: FileModifiedEvent) -> None: # type: ignore[override]
|
||||
"""
|
||||
Handle new data appended to the log file.
|
||||
"""
|
||||
|
|
@ -109,7 +115,8 @@ class _LogHandler(FileSystemEventHandler):
|
|||
|
||||
self._read_new_lines()
|
||||
|
||||
def on_moved(self, event: FileMovedEvent) -> None: # type: ignore[override]
|
||||
def on_moved(self,
|
||||
event: FileMovedEvent) -> None: # type: ignore[override]
|
||||
"""
|
||||
Handle log rotation via rename (access.log -> access.log.1).
|
||||
"""
|
||||
|
|
@ -117,10 +124,12 @@ class _LogHandler(FileSystemEventHandler):
|
|||
return
|
||||
|
||||
if Path(str(event.src_path)).resolve() == Path(self._target).resolve():
|
||||
logger.info("Log rotated: %s -> %s", event.src_path, event.dest_path)
|
||||
logger.info("Log rotated: %s -> %s", event.src_path,
|
||||
event.dest_path)
|
||||
self._handle_rotation()
|
||||
|
||||
def on_created(self, event: FileCreatedEvent) -> None: # type: ignore[override]
|
||||
def on_created(self,
|
||||
event: FileCreatedEvent) -> None: # type: ignore[override]
|
||||
"""
|
||||
Handle log rotation where a new file is created at the target path.
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -59,11 +59,21 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
|||
session_factory=app.state.session_factory,
|
||||
)
|
||||
|
||||
inference_engine = _load_inference_engine()
|
||||
app.state.models_loaded = inference_engine is not None and inference_engine.is_loaded
|
||||
app.state.detection_mode = "hybrid" if app.state.models_loaded else "rules"
|
||||
|
||||
pipeline = Pipeline(
|
||||
redis_client=redis_client, # type: ignore[arg-type]
|
||||
rule_engine=RuleEngine(),
|
||||
geoip=geoip,
|
||||
on_result=dispatcher.dispatch,
|
||||
inference_engine=inference_engine,
|
||||
ensemble_weights={
|
||||
"ae": settings.ensemble_weight_ae,
|
||||
"rf": settings.ensemble_weight_rf,
|
||||
"if": settings.ensemble_weight_if,
|
||||
},
|
||||
raw_queue_size=settings.raw_queue_size,
|
||||
parsed_queue_size=settings.parsed_queue_size,
|
||||
feature_queue_size=settings.feature_queue_size,
|
||||
|
|
@ -100,6 +110,34 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
|||
logger.info("AngelusVigil shut down cleanly")
|
||||
|
||||
|
||||
def _load_inference_engine() -> object | None:
|
||||
"""
|
||||
Attempt to load the ONNX inference engine from the
|
||||
configured model directory, returning None if ML
|
||||
dependencies are missing or no models are found
|
||||
"""
|
||||
try:
|
||||
from app.core.detection.inference import (
|
||||
InferenceEngine, )
|
||||
except ImportError:
|
||||
logger.info("onnxruntime not installed — running in rules-only mode")
|
||||
return None
|
||||
|
||||
engine = InferenceEngine(model_dir=settings.model_dir)
|
||||
if engine.is_loaded:
|
||||
logger.info(
|
||||
"ML models loaded from %s",
|
||||
settings.model_dir,
|
||||
)
|
||||
return engine
|
||||
|
||||
logger.info(
|
||||
"No ML models found in %s — running in rules-only mode",
|
||||
settings.model_dir,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def create_app() -> FastAPI:
|
||||
"""
|
||||
Build and configure the AngelusVigil FastAPI application.
|
||||
|
|
|
|||
|
|
@ -15,14 +15,12 @@ class ModelMetadata(TimestampedModel, table=True):
|
|||
"""
|
||||
|
||||
__tablename__ = "model_metadata"
|
||||
__table_args__ = (
|
||||
Index(
|
||||
"idx_model_metadata_active",
|
||||
"model_type",
|
||||
unique=True,
|
||||
postgresql_where=text("is_active = TRUE"),
|
||||
),
|
||||
)
|
||||
__table_args__ = (Index(
|
||||
"idx_model_metadata_active",
|
||||
"model_type",
|
||||
unique=True,
|
||||
postgresql_where=text("is_active = TRUE"),
|
||||
), )
|
||||
|
||||
model_type: str = Field(max_length=30)
|
||||
version: str = Field(max_length=64)
|
||||
|
|
|
|||
|
|
@ -42,15 +42,16 @@ class ThreatEvent(TimestampedModel, table=True):
|
|||
user_agent: str
|
||||
threat_score: float = Field(sa_column=Column(Float, nullable=False))
|
||||
severity: str = Field(max_length=6)
|
||||
component_scores: dict[str, float] = Field(sa_column=Column(JSON, nullable=False))
|
||||
component_scores: dict[str, float] = Field(
|
||||
sa_column=Column(JSON, nullable=False))
|
||||
geo_country: str | None = Field(default=None, max_length=2)
|
||||
geo_city: str | None = Field(default=None, max_length=255)
|
||||
geo_lat: float | None = Field(default=None)
|
||||
geo_lon: float | None = Field(default=None)
|
||||
feature_vector: list[float] = Field(sa_column=Column(JSON, nullable=False))
|
||||
matched_rules: list[str] | None = Field(
|
||||
default=None, sa_column=Column(JSON, nullable=True)
|
||||
)
|
||||
matched_rules: list[str] | None = Field(default=None,
|
||||
sa_column=Column(JSON,
|
||||
nullable=True))
|
||||
model_version: str | None = Field(default=None, max_length=64)
|
||||
reviewed: bool = Field(default=False)
|
||||
review_label: str | None = Field(default=None, max_length=20)
|
||||
|
|
|
|||
|
|
@ -35,37 +35,34 @@ async def get_stats(
|
|||
delta = _RANGE_MAP.get(time_range, timedelta(hours=24))
|
||||
cutoff = datetime.now(UTC) - delta
|
||||
|
||||
base = select(ThreatEvent).where(ThreatEvent.created_at >= cutoff) # type: ignore[arg-type]
|
||||
base = select(ThreatEvent).where(ThreatEvent.created_at
|
||||
>= cutoff) # type: ignore[arg-type]
|
||||
|
||||
total_q = select(func.count()).select_from(base.subquery())
|
||||
total = (await session.execute(total_q)).scalar_one()
|
||||
|
||||
sev_q = (
|
||||
select(ThreatEvent.severity, func.count()) # type: ignore[call-overload]
|
||||
.where(ThreatEvent.created_at >= cutoff)
|
||||
.group_by(ThreatEvent.severity)
|
||||
)
|
||||
select(ThreatEvent.severity,
|
||||
func.count()) # type: ignore[call-overload]
|
||||
.where(ThreatEvent.created_at >= cutoff).group_by(
|
||||
ThreatEvent.severity))
|
||||
sev_rows = (await session.execute(sev_q)).all()
|
||||
sev_map = {row[0]: row[1] for row in sev_rows}
|
||||
|
||||
threats_detected = sev_map.get("HIGH", 0) + sev_map.get("MEDIUM", 0)
|
||||
|
||||
ip_q = (
|
||||
select(ThreatEvent.source_ip, func.count().label("cnt")) # type: ignore[call-overload]
|
||||
.where(ThreatEvent.created_at >= cutoff)
|
||||
.group_by(ThreatEvent.source_ip)
|
||||
.order_by(func.count().desc())
|
||||
.limit(10)
|
||||
)
|
||||
select(ThreatEvent.source_ip,
|
||||
func.count().label("cnt")) # type: ignore[call-overload]
|
||||
.where(ThreatEvent.created_at >= cutoff).group_by(
|
||||
ThreatEvent.source_ip).order_by(func.count().desc()).limit(10))
|
||||
ip_rows = (await session.execute(ip_q)).all()
|
||||
|
||||
path_q = (
|
||||
select(ThreatEvent.request_path, func.count().label("cnt")) # type: ignore[call-overload]
|
||||
.where(ThreatEvent.created_at >= cutoff)
|
||||
.group_by(ThreatEvent.request_path)
|
||||
.order_by(func.count().desc())
|
||||
.limit(10)
|
||||
)
|
||||
select(ThreatEvent.request_path,
|
||||
func.count().label("cnt")) # type: ignore[call-overload]
|
||||
.where(ThreatEvent.created_at >= cutoff).group_by(
|
||||
ThreatEvent.request_path).order_by(func.count().desc()).limit(10))
|
||||
path_rows = (await session.execute(path_q)).all()
|
||||
|
||||
return StatsResponse(
|
||||
|
|
@ -77,6 +74,10 @@ async def get_stats(
|
|||
medium=sev_map.get("MEDIUM", 0),
|
||||
low=sev_map.get("LOW", 0),
|
||||
),
|
||||
top_source_ips=[IPStatEntry(source_ip=row[0], count=row[1]) for row in ip_rows],
|
||||
top_attacked_paths=[PathStatEntry(path=row[0], count=row[1]) for row in path_rows],
|
||||
top_source_ips=[
|
||||
IPStatEntry(source_ip=row[0], count=row[1]) for row in ip_rows
|
||||
],
|
||||
top_attacked_paths=[
|
||||
PathStatEntry(path=row[0], count=row[1]) for row in path_rows
|
||||
],
|
||||
)
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ from datetime import datetime
|
|||
from sqlalchemy import func, select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.core.detection.ensemble import classify_severity
|
||||
from app.core.ingestion.pipeline import ScoredRequest
|
||||
from app.models.threat_event import ThreatEvent
|
||||
from app.schemas.threats import (
|
||||
|
|
@ -63,19 +64,28 @@ async def get_threats(
|
|||
count_query = select(func.count()).select_from(ThreatEvent)
|
||||
|
||||
if severity:
|
||||
query = query.where(ThreatEvent.severity == severity.upper()) # type: ignore[arg-type]
|
||||
count_query = count_query.where(ThreatEvent.severity == severity.upper()) # type: ignore[arg-type]
|
||||
query = query.where(
|
||||
ThreatEvent.severity == severity.upper()) # type: ignore[arg-type]
|
||||
count_query = count_query.where(
|
||||
ThreatEvent.severity == severity.upper()) # type: ignore[arg-type]
|
||||
if source_ip:
|
||||
query = query.where(ThreatEvent.source_ip == source_ip) # type: ignore[arg-type]
|
||||
count_query = count_query.where(ThreatEvent.source_ip == source_ip) # type: ignore[arg-type]
|
||||
query = query.where(
|
||||
ThreatEvent.source_ip == source_ip) # type: ignore[arg-type]
|
||||
count_query = count_query.where(
|
||||
ThreatEvent.source_ip == source_ip) # type: ignore[arg-type]
|
||||
if since:
|
||||
query = query.where(ThreatEvent.created_at >= since) # type: ignore[arg-type]
|
||||
count_query = count_query.where(ThreatEvent.created_at >= since) # type: ignore[arg-type]
|
||||
query = query.where(ThreatEvent.created_at
|
||||
>= since) # type: ignore[arg-type]
|
||||
count_query = count_query.where(ThreatEvent.created_at
|
||||
>= since) # type: ignore[arg-type]
|
||||
if until:
|
||||
query = query.where(ThreatEvent.created_at <= until) # type: ignore[arg-type]
|
||||
count_query = count_query.where(ThreatEvent.created_at <= until) # type: ignore[arg-type]
|
||||
query = query.where(ThreatEvent.created_at
|
||||
<= until) # type: ignore[arg-type]
|
||||
count_query = count_query.where(ThreatEvent.created_at
|
||||
<= until) # type: ignore[arg-type]
|
||||
|
||||
query = query.order_by(ThreatEvent.created_at.desc()) # type: ignore[attr-defined]
|
||||
query = query.order_by(
|
||||
ThreatEvent.created_at.desc()) # type: ignore[attr-defined]
|
||||
query = query.offset(offset).limit(limit)
|
||||
|
||||
total = (await session.execute(count_query)).scalar_one()
|
||||
|
|
@ -116,16 +126,16 @@ async def create_threat_event(
|
|||
status_code=scored.entry.status_code,
|
||||
response_size=scored.entry.response_size,
|
||||
user_agent=scored.entry.user_agent,
|
||||
threat_score=scored.rule_result.threat_score,
|
||||
severity=scored.rule_result.severity,
|
||||
threat_score=scored.final_score,
|
||||
severity=classify_severity(scored.final_score),
|
||||
component_scores=scored.rule_result.component_scores,
|
||||
geo_country=scored.geo.country if scored.geo else None,
|
||||
geo_city=scored.geo.city if scored.geo else None,
|
||||
geo_lat=scored.geo.lat if scored.geo else None,
|
||||
geo_lon=scored.geo.lon if scored.geo else None,
|
||||
geo_country=(scored.geo.country if scored.geo else None),
|
||||
geo_city=(scored.geo.city if scored.geo else None),
|
||||
geo_lat=(scored.geo.lat if scored.geo else None),
|
||||
geo_lon=(scored.geo.lon if scored.geo else None),
|
||||
feature_vector=scored.feature_vector,
|
||||
matched_rules=scored.rule_result.matched_rules or None,
|
||||
model_version="rules-v1",
|
||||
matched_rules=(scored.rule_result.matched_rules or None),
|
||||
model_version=scored.detection_mode,
|
||||
)
|
||||
session.add(event)
|
||||
await session.flush()
|
||||
|
|
|
|||
|
|
@ -3,6 +3,8 @@
|
|||
main.py
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import typer
|
||||
|
||||
app = typer.Typer(
|
||||
|
|
@ -11,15 +13,21 @@ app = typer.Typer(
|
|||
no_args_is_help=True,
|
||||
)
|
||||
|
||||
DEFAULT_MODEL_DIR = "data/models"
|
||||
DEFAULT_EPOCHS = 100
|
||||
DEFAULT_BATCH_SIZE = 256
|
||||
DEFAULT_SERVER_URL = "http://localhost:8000"
|
||||
|
||||
|
||||
@app.command()
|
||||
def serve(
|
||||
host: str = typer.Option("0.0.0.0", help="Bind address"),
|
||||
port: int = typer.Option(8000, help="Bind port"),
|
||||
reload: bool = typer.Option(False, help="Enable auto-reload for development"),
|
||||
reload: bool = typer.Option(False,
|
||||
help="Enable auto-reload for development"),
|
||||
) -> None:
|
||||
"""
|
||||
Start the AngelusVigil API server.
|
||||
Start the AngelusVigil API server
|
||||
"""
|
||||
import uvicorn
|
||||
|
||||
|
|
@ -31,16 +39,148 @@ def serve(
|
|||
)
|
||||
|
||||
|
||||
@app.command()
|
||||
def train(
|
||||
dataset: Path = typer.Option(..., help="Path to CSIC 2010 CSV dataset"),
|
||||
output_dir: Path = typer.Option(
|
||||
DEFAULT_MODEL_DIR,
|
||||
help="Directory to save ONNX models",
|
||||
),
|
||||
epochs: int = typer.Option(
|
||||
DEFAULT_EPOCHS,
|
||||
help="Autoencoder training epochs",
|
||||
),
|
||||
batch_size: int = typer.Option(
|
||||
DEFAULT_BATCH_SIZE,
|
||||
help="Training batch size",
|
||||
),
|
||||
) -> None:
|
||||
"""
|
||||
Train all ML models and export to ONNX
|
||||
"""
|
||||
import json
|
||||
|
||||
import numpy as np
|
||||
from sklearn.model_selection import train_test_split
|
||||
|
||||
from ml.export_onnx import (
|
||||
export_autoencoder,
|
||||
export_isolation_forest,
|
||||
export_random_forest,
|
||||
)
|
||||
from ml.train_autoencoder import train_autoencoder
|
||||
from ml.train_classifiers import (
|
||||
train_isolation_forest,
|
||||
train_random_forest,
|
||||
)
|
||||
|
||||
if not dataset.exists():
|
||||
typer.echo(
|
||||
f"Error: dataset not found at {dataset}",
|
||||
err=True,
|
||||
)
|
||||
raise typer.Exit(code=1)
|
||||
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
typer.echo(f"Loading dataset from {dataset}")
|
||||
data = np.loadtxt(str(dataset), delimiter=",", dtype=np.float32)
|
||||
X = data[:, :-1]
|
||||
y = data[:, -1].astype(int)
|
||||
|
||||
X_normal = X[y == 0]
|
||||
X_train, X_test, y_train, y_test = train_test_split(X,
|
||||
y,
|
||||
test_size=0.2,
|
||||
stratify=y,
|
||||
random_state=42)
|
||||
|
||||
typer.echo(f"Training autoencoder for {epochs} epochs")
|
||||
ae_result = train_autoencoder(
|
||||
X_normal,
|
||||
epochs=epochs,
|
||||
batch_size=batch_size,
|
||||
)
|
||||
export_autoencoder(ae_result["model"], output_dir / "ae.onnx")
|
||||
ae_result["scaler"].save_json(output_dir / "scaler.json")
|
||||
threshold_data = {"threshold": float(ae_result["threshold"])}
|
||||
(output_dir / "threshold.json").write_text(json.dumps(threshold_data))
|
||||
typer.echo(f" AE threshold: {ae_result['threshold']:.6f}")
|
||||
|
||||
typer.echo("Training random forest")
|
||||
rf_result = train_random_forest(X_train, y_train)
|
||||
export_random_forest(rf_result["model"], X.shape[1],
|
||||
output_dir / "rf.onnx")
|
||||
typer.echo(f" RF F1: {rf_result['metrics']['f1']:.4f}")
|
||||
|
||||
typer.echo("Training isolation forest")
|
||||
if_result = train_isolation_forest(X_normal)
|
||||
export_isolation_forest(if_result["model"], X.shape[1],
|
||||
output_dir / "if.onnx")
|
||||
typer.echo(f" IF samples: {if_result['metrics']['n_samples']}")
|
||||
|
||||
typer.echo(f"Models exported to {output_dir}")
|
||||
|
||||
|
||||
@app.command()
|
||||
def replay(
|
||||
log_file: Path = typer.Option(...,
|
||||
help="Path to nginx access log file"),
|
||||
url: str = typer.Option(
|
||||
DEFAULT_SERVER_URL,
|
||||
help="Running server URL to send logs to",
|
||||
),
|
||||
batch_size: int = typer.Option(100, help="Lines per batch"),
|
||||
) -> None:
|
||||
"""
|
||||
Replay historical log lines through the pipeline
|
||||
"""
|
||||
import httpx
|
||||
|
||||
if not log_file.exists():
|
||||
typer.echo(
|
||||
f"Error: log file not found at {log_file}",
|
||||
err=True,
|
||||
)
|
||||
raise typer.Exit(code=1)
|
||||
|
||||
lines = log_file.read_text().strip().splitlines()
|
||||
typer.echo(f"Replaying {len(lines)} lines to {url}")
|
||||
|
||||
sent = 0
|
||||
with httpx.Client(timeout=30.0) as client:
|
||||
for i in range(0, len(lines), batch_size):
|
||||
batch = lines[i:i + batch_size]
|
||||
response = client.post(
|
||||
f"{url}/ingest/batch",
|
||||
json={"lines": batch},
|
||||
)
|
||||
if response.status_code == 200:
|
||||
sent += len(batch)
|
||||
else:
|
||||
typer.echo(
|
||||
f" Batch {i} failed: {response.status_code}",
|
||||
err=True,
|
||||
)
|
||||
|
||||
typer.echo(f"Replayed {sent}/{len(lines)} lines")
|
||||
|
||||
|
||||
@app.command()
|
||||
def config() -> None:
|
||||
"""
|
||||
Print the current configuration (secrets redacted).
|
||||
Print the current configuration (secrets redacted)
|
||||
"""
|
||||
from app.config import settings
|
||||
|
||||
safe_fields = {}
|
||||
for key, value in settings.model_dump().items():
|
||||
if any(secret in key for secret in ("key", "password", "secret", "token")):
|
||||
if any(secret in key for secret in (
|
||||
"key",
|
||||
"password",
|
||||
"secret",
|
||||
"token",
|
||||
)):
|
||||
safe_fields[key] = "***REDACTED***"
|
||||
elif "url" in key and "@" in str(value):
|
||||
safe_fields[key] = _redact_url(str(value))
|
||||
|
|
@ -54,12 +194,12 @@ def config() -> None:
|
|||
@app.command()
|
||||
def health(
|
||||
url: str = typer.Option(
|
||||
"http://localhost:8000",
|
||||
DEFAULT_SERVER_URL,
|
||||
help="Base URL of the running server",
|
||||
),
|
||||
) -> None:
|
||||
"""
|
||||
Ping the running server's /health endpoint.
|
||||
Ping the running server's /health endpoint
|
||||
"""
|
||||
import httpx
|
||||
|
||||
|
|
@ -69,18 +209,24 @@ def health(
|
|||
data = response.json()
|
||||
typer.echo(f" status: {data.get('status', 'unknown')}")
|
||||
typer.echo(f" uptime: {data.get('uptime_seconds', 0):.0f}s")
|
||||
typer.echo(f" pipeline: {'running' if data.get('pipeline_running') else 'stopped'}")
|
||||
typer.echo(
|
||||
f" pipeline: {'running' if data.get('pipeline_running') else 'stopped'}"
|
||||
)
|
||||
except httpx.ConnectError:
|
||||
typer.echo("Error: cannot connect to server", err=True)
|
||||
raise typer.Exit(code=1) from None
|
||||
except httpx.HTTPStatusError as exc:
|
||||
typer.echo(f"Error: server returned {exc.response.status_code}", err=True)
|
||||
typer.echo(
|
||||
f"Error: server returned {exc.response.status_code}",
|
||||
err=True,
|
||||
)
|
||||
raise typer.Exit(code=1) from None
|
||||
|
||||
|
||||
def _redact_url(url: str) -> str:
|
||||
"""
|
||||
Replace the user:password portion of a database URL with ***:***.
|
||||
Replace the user:password portion of a database
|
||||
URL with ***:***
|
||||
"""
|
||||
if "://" not in url or "@" not in url:
|
||||
return url
|
||||
|
|
|
|||
|
|
@ -71,4 +71,4 @@ class ThreatAutoencoder(nn.Module):
|
|||
Per-sample mean squared error between input and reconstruction.
|
||||
"""
|
||||
reconstructed = self.forward(x)
|
||||
return torch.mean((x - reconstructed) ** 2, dim=1)
|
||||
return torch.mean((x - reconstructed)**2, dim=1)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,104 @@
|
|||
"""
|
||||
©AngelaMos | 2026
|
||||
experiment.py
|
||||
"""
|
||||
|
||||
import platform
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from types import TracebackType
|
||||
|
||||
import mlflow
|
||||
|
||||
|
||||
class VigilExperiment:
|
||||
"""
|
||||
Context manager wrapping MLflow experiment runs
|
||||
with automatic system metadata logging
|
||||
"""
|
||||
|
||||
def __init__(self, experiment_name: str) -> None:
|
||||
self._experiment_name = experiment_name
|
||||
self._run: mlflow.ActiveRun | None = None
|
||||
self._run_id: str | None = None
|
||||
|
||||
@property
|
||||
def run_id(self) -> str | None:
|
||||
"""
|
||||
The MLflow run ID, set after entering context
|
||||
"""
|
||||
return self._run_id
|
||||
|
||||
def __enter__(self) -> VigilExperiment:
|
||||
mlflow.set_experiment(self._experiment_name)
|
||||
self._run = mlflow.start_run()
|
||||
self._run_id = self._run.info.run_id
|
||||
self._log_system_metadata()
|
||||
return self
|
||||
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: TracebackType | None,
|
||||
) -> None:
|
||||
if exc_type is not None:
|
||||
mlflow.set_tag("status", "failed")
|
||||
mlflow.set_tag("error", str(exc_val)[:500])
|
||||
else:
|
||||
mlflow.set_tag("status", "completed")
|
||||
mlflow.end_run()
|
||||
self._run = None
|
||||
|
||||
def log_params(self, params: dict[str, object]) -> None:
|
||||
"""
|
||||
Log a dictionary of parameters to the active run
|
||||
"""
|
||||
mlflow.log_params(params)
|
||||
|
||||
def log_metrics(
|
||||
self,
|
||||
metrics: dict[str, float],
|
||||
step: int | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Log a dictionary of metrics to the active run
|
||||
"""
|
||||
mlflow.log_metrics(metrics, step=step)
|
||||
|
||||
def log_artifact(self, path: Path | str) -> None:
|
||||
"""
|
||||
Log a local file as an artifact
|
||||
"""
|
||||
mlflow.log_artifact(str(path))
|
||||
|
||||
def _log_system_metadata(self) -> None:
|
||||
"""
|
||||
Record Python version and git commit hash
|
||||
"""
|
||||
mlflow.set_tag("python_version", sys.version.split()[0])
|
||||
mlflow.set_tag("platform", platform.system())
|
||||
|
||||
git_hash = _get_git_hash()
|
||||
if git_hash is not None:
|
||||
mlflow.set_tag("git_commit", git_hash)
|
||||
|
||||
|
||||
def _get_git_hash() -> str | None:
|
||||
"""
|
||||
Return the short git commit hash or None
|
||||
"""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["git", "rev-parse", "--short", "HEAD"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5,
|
||||
check=False,
|
||||
)
|
||||
if result.returncode == 0:
|
||||
return result.stdout.strip()
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
return None
|
||||
|
|
@ -0,0 +1,90 @@
|
|||
"""
|
||||
©AngelaMos | 2026
|
||||
export_onnx.py
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from skl2onnx import convert_sklearn
|
||||
from skl2onnx.common.data_types import FloatTensorType
|
||||
from sklearn.base import BaseEstimator
|
||||
|
||||
from ml.autoencoder import ThreatAutoencoder
|
||||
|
||||
ONNX_OPSET = 17
|
||||
SKL_TARGET_OPSET = {"": 17, "ai.onnx.ml": 3}
|
||||
|
||||
|
||||
def export_autoencoder(
|
||||
model: ThreatAutoencoder,
|
||||
path: Path | str,
|
||||
opset: int = ONNX_OPSET,
|
||||
) -> Path:
|
||||
"""
|
||||
Export a PyTorch autoencoder to ONNX with dynamic batch dimension
|
||||
"""
|
||||
path = Path(path)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
model.eval()
|
||||
dummy = torch.randn(1, model.input_dim)
|
||||
|
||||
batch_dim = torch.export.Dim("batch_size", min=1)
|
||||
|
||||
torch.onnx.export(
|
||||
model,
|
||||
dummy,
|
||||
str(path),
|
||||
opset_version=opset,
|
||||
export_params=True,
|
||||
do_constant_folding=True,
|
||||
input_names=["features"],
|
||||
output_names=["reconstructed"],
|
||||
dynamic_shapes={"x": {
|
||||
0: batch_dim
|
||||
}},
|
||||
)
|
||||
return path
|
||||
|
||||
|
||||
def export_random_forest(
|
||||
model: BaseEstimator,
|
||||
n_features: int,
|
||||
path: Path | str,
|
||||
) -> Path:
|
||||
"""
|
||||
Export a sklearn random forest (or calibrated wrapper) to ONNX
|
||||
"""
|
||||
path = Path(path)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
initial_type = [("features", FloatTensorType([None, n_features]))]
|
||||
onnx_model = convert_sklearn(
|
||||
model,
|
||||
initial_types=initial_type,
|
||||
target_opset=SKL_TARGET_OPSET,
|
||||
)
|
||||
path.write_bytes(onnx_model.SerializeToString())
|
||||
return path
|
||||
|
||||
|
||||
def export_isolation_forest(
|
||||
model: BaseEstimator,
|
||||
n_features: int,
|
||||
path: Path | str,
|
||||
) -> Path:
|
||||
"""
|
||||
Export a sklearn isolation forest to ONNX
|
||||
"""
|
||||
path = Path(path)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
initial_type = [("features", FloatTensorType([None, n_features]))]
|
||||
onnx_model = convert_sklearn(
|
||||
model,
|
||||
initial_types=initial_type,
|
||||
target_opset=SKL_TARGET_OPSET,
|
||||
)
|
||||
path.write_bytes(onnx_model.SerializeToString())
|
||||
return path
|
||||
|
|
@ -0,0 +1,115 @@
|
|||
"""
|
||||
©AngelaMos | 2026
|
||||
train_autoencoder.py
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.utils.data import DataLoader, TensorDataset
|
||||
|
||||
from ml.autoencoder import ThreatAutoencoder
|
||||
from ml.scaler import FeatureScaler
|
||||
|
||||
|
||||
def train_autoencoder(
|
||||
X_normal: np.ndarray,
|
||||
epochs: int = 100,
|
||||
batch_size: int = 256,
|
||||
lr: float = 1e-3,
|
||||
percentile: float = 99.5,
|
||||
val_split: float = 0.15,
|
||||
patience: int = 10,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Train the autoencoder on normal-only traffic and calibrate the
|
||||
anomaly detection threshold.
|
||||
"""
|
||||
input_dim = X_normal.shape[1]
|
||||
|
||||
split_idx = int(len(X_normal) * (1 - val_split))
|
||||
X_train_raw = X_normal[:split_idx]
|
||||
X_val_raw = X_normal[split_idx:]
|
||||
|
||||
scaler = FeatureScaler()
|
||||
X_train_scaled = scaler.fit_transform(X_train_raw)
|
||||
X_val_scaled = scaler.transform(X_val_raw)
|
||||
|
||||
train_tensor = torch.from_numpy(X_train_scaled)
|
||||
val_tensor = torch.from_numpy(X_val_scaled)
|
||||
|
||||
train_loader = DataLoader(
|
||||
TensorDataset(train_tensor),
|
||||
batch_size=batch_size,
|
||||
shuffle=True,
|
||||
drop_last=len(train_tensor) > batch_size,
|
||||
)
|
||||
|
||||
model = ThreatAutoencoder(input_dim=input_dim)
|
||||
optimizer = torch.optim.AdamW(model.parameters(),
|
||||
lr=lr,
|
||||
weight_decay=1e-5,
|
||||
betas=(0.9, 0.999))
|
||||
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,
|
||||
mode="min",
|
||||
factor=0.5,
|
||||
patience=5,
|
||||
min_lr=1e-6)
|
||||
|
||||
history: dict[str, list[float]] = {"train_loss": [], "val_loss": []}
|
||||
best_val_loss = float("inf")
|
||||
best_state = None
|
||||
epochs_without_improvement = 0
|
||||
|
||||
for _epoch in range(epochs):
|
||||
model.train()
|
||||
epoch_loss = 0.0
|
||||
n_batches = 0
|
||||
|
||||
for (batch, ) in train_loader:
|
||||
reconstructed = model(batch)
|
||||
loss = torch.nn.functional.mse_loss(reconstructed, batch)
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
|
||||
optimizer.step()
|
||||
epoch_loss += loss.item()
|
||||
n_batches += 1
|
||||
|
||||
avg_train_loss = epoch_loss / max(n_batches, 1)
|
||||
history["train_loss"].append(avg_train_loss)
|
||||
|
||||
model.eval()
|
||||
with torch.no_grad():
|
||||
val_reconstructed = model(val_tensor)
|
||||
val_loss = torch.nn.functional.mse_loss(val_reconstructed,
|
||||
val_tensor).item()
|
||||
history["val_loss"].append(val_loss)
|
||||
|
||||
scheduler.step(val_loss)
|
||||
|
||||
if val_loss < best_val_loss:
|
||||
best_val_loss = val_loss
|
||||
best_state = {k: v.clone() for k, v in model.state_dict().items()}
|
||||
epochs_without_improvement = 0
|
||||
else:
|
||||
epochs_without_improvement += 1
|
||||
|
||||
if epochs_without_improvement >= patience:
|
||||
break
|
||||
|
||||
if best_state is not None:
|
||||
model.load_state_dict(best_state)
|
||||
|
||||
model.eval()
|
||||
with torch.no_grad():
|
||||
val_errors = model.compute_reconstruction_error(val_tensor)
|
||||
threshold = float(np.percentile(val_errors.numpy(), percentile))
|
||||
|
||||
return {
|
||||
"model": model,
|
||||
"scaler": scaler,
|
||||
"threshold": threshold,
|
||||
"history": history,
|
||||
}
|
||||
|
|
@ -0,0 +1,79 @@
|
|||
"""
|
||||
©AngelaMos | 2026
|
||||
train_classifiers.py
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
from sklearn.calibration import CalibratedClassifierCV
|
||||
from sklearn.ensemble import IsolationForest, RandomForestClassifier
|
||||
from sklearn.metrics import (
|
||||
accuracy_score,
|
||||
average_precision_score,
|
||||
f1_score,
|
||||
precision_score,
|
||||
recall_score,
|
||||
)
|
||||
from sklearn.model_selection import train_test_split
|
||||
|
||||
|
||||
def train_random_forest(
|
||||
X: np.ndarray,
|
||||
y: np.ndarray,
|
||||
n_estimators: int = 200,
|
||||
max_depth: int = 20,
|
||||
calibration_split: float = 0.2,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Train a random forest with isotonic probability calibration
|
||||
"""
|
||||
X_train, X_eval, y_train, y_eval = train_test_split(
|
||||
X, y, test_size=calibration_split, stratify=y, random_state=42)
|
||||
|
||||
base_rf = RandomForestClassifier(
|
||||
n_estimators=n_estimators,
|
||||
max_depth=max_depth,
|
||||
class_weight="balanced",
|
||||
random_state=42,
|
||||
n_jobs=-1,
|
||||
)
|
||||
|
||||
calibrated = CalibratedClassifierCV(base_rf, method="isotonic", cv=3)
|
||||
calibrated.fit(X_train, y_train)
|
||||
|
||||
y_pred = calibrated.predict(X_eval)
|
||||
y_proba = calibrated.predict_proba(X_eval)[:, 1]
|
||||
|
||||
metrics = {
|
||||
"accuracy": float(accuracy_score(y_eval, y_pred)),
|
||||
"precision": float(precision_score(y_eval, y_pred, zero_division=0)),
|
||||
"recall": float(recall_score(y_eval, y_pred, zero_division=0)),
|
||||
"f1": float(f1_score(y_eval, y_pred, zero_division=0)),
|
||||
"pr_auc": float(average_precision_score(y_eval, y_proba)),
|
||||
}
|
||||
|
||||
return {"model": calibrated, "metrics": metrics}
|
||||
|
||||
|
||||
def train_isolation_forest(
|
||||
X_normal: np.ndarray,
|
||||
n_estimators: int = 200,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Train an isolation forest on normal-only traffic
|
||||
"""
|
||||
iso = IsolationForest(
|
||||
n_estimators=n_estimators,
|
||||
contamination="auto",
|
||||
random_state=42,
|
||||
n_jobs=-1,
|
||||
)
|
||||
iso.fit(X_normal)
|
||||
|
||||
return {
|
||||
"model": iso,
|
||||
"metrics": {
|
||||
"n_samples": len(X_normal)
|
||||
},
|
||||
}
|
||||
|
|
@ -50,6 +50,7 @@ ml = [
|
|||
"numpy>=2.4.2",
|
||||
"pandas>=2.2.0,<3",
|
||||
"imbalanced-learn>=0.14.1",
|
||||
"onnxscript>=0.6.2",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
"""
|
||||
©AngelaMos | 2026
|
||||
conftest.py
|
||||
|
||||
Shared pytest fixtures for in-memory SQLite database and HTTPX test client setup.
|
||||
"""
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
|
|
@ -84,5 +86,6 @@ async def db_client(db_engine) -> AsyncIterator[AsyncClient]:
|
|||
test_app.dependency_overrides[get_session] = override_get_session
|
||||
|
||||
transport = ASGITransport(app=test_app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
async with AsyncClient(transport=transport,
|
||||
base_url="http://test") as client:
|
||||
yield client
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
"""
|
||||
©AngelaMos | 2026
|
||||
test_api.py
|
||||
|
||||
Tests the FastAPI REST endpoints for threats, stats, health, readiness, and model management.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
|
|
@ -19,7 +21,8 @@ async def test_health_returns_200() -> None:
|
|||
Health endpoint returns 200 with status and uptime.
|
||||
"""
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
async with AsyncClient(transport=transport,
|
||||
base_url="http://test") as client:
|
||||
response = await client.get("/health")
|
||||
|
||||
assert response.status_code == 200
|
||||
|
|
@ -36,7 +39,8 @@ async def test_health_returns_pipeline_status() -> None:
|
|||
Health response includes pipeline_running boolean.
|
||||
"""
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
async with AsyncClient(transport=transport,
|
||||
base_url="http://test") as client:
|
||||
response = await client.get("/health")
|
||||
|
||||
data = response.json()
|
||||
|
|
@ -49,7 +53,8 @@ async def test_ready_returns_check_structure() -> None:
|
|||
Readiness endpoint returns structured component checks.
|
||||
"""
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
async with AsyncClient(transport=transport,
|
||||
base_url="http://test") as client:
|
||||
response = await client.get("/ready")
|
||||
|
||||
assert response.status_code in (200, 503)
|
||||
|
|
@ -178,15 +183,16 @@ async def test_stats_empty_window(db_client) -> None:
|
|||
@pytest.mark.asyncio
|
||||
async def test_model_status() -> None:
|
||||
"""
|
||||
GET /models/status returns rules-only detection mode.
|
||||
GET /models/status returns rules detection mode
|
||||
"""
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
async with AsyncClient(transport=transport,
|
||||
base_url="http://test") as client:
|
||||
response = await client.get("/models/status")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["detection_mode"] == "rules-only"
|
||||
assert data["detection_mode"] == "rules"
|
||||
assert data["active_models"] == []
|
||||
|
||||
|
||||
|
|
@ -196,7 +202,8 @@ async def test_retrain_returns_202() -> None:
|
|||
POST /models/retrain returns 202 Accepted with a job ID.
|
||||
"""
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
async with AsyncClient(transport=transport,
|
||||
base_url="http://test") as client:
|
||||
response = await client.post("/models/retrain")
|
||||
|
||||
assert response.status_code == 202
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
"""
|
||||
©AngelaMos | 2026
|
||||
test_autoencoder.py
|
||||
|
||||
Tests the ThreatAutoencoder architecture: shapes, output range, reconstruction error, and training behavior.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
|
@ -12,18 +14,27 @@ from ml.autoencoder import ThreatAutoencoder
|
|||
class TestAutoencoderArchitecture:
|
||||
|
||||
def test_output_shape_matches_input(self) -> None:
|
||||
"""
|
||||
Forward pass on a batch of 16 produces output matching input shape.
|
||||
"""
|
||||
model = ThreatAutoencoder(input_dim=35)
|
||||
x = torch.randn(16, 35)
|
||||
out = model(x)
|
||||
assert out.shape == (16, 35)
|
||||
|
||||
def test_bottleneck_dim_is_six(self) -> None:
|
||||
"""
|
||||
Encoder bottleneck compresses 35 features to a 6-dimensional latent vector.
|
||||
"""
|
||||
model = ThreatAutoencoder(input_dim=35)
|
||||
x = torch.randn(4, 35)
|
||||
encoded = model.encode(x)
|
||||
assert encoded.shape == (4, 6)
|
||||
|
||||
def test_single_sample_forward(self) -> None:
|
||||
"""
|
||||
Single-sample forward pass completes without error in eval mode.
|
||||
"""
|
||||
model = ThreatAutoencoder(input_dim=35)
|
||||
model.eval()
|
||||
x = torch.randn(1, 35)
|
||||
|
|
@ -32,6 +43,9 @@ class TestAutoencoderArchitecture:
|
|||
assert out.shape == (1, 35)
|
||||
|
||||
def test_output_values_in_zero_one_range(self) -> None:
|
||||
"""
|
||||
Decoder output is bounded to [0, 1] by the final activation.
|
||||
"""
|
||||
model = ThreatAutoencoder(input_dim=35)
|
||||
model.eval()
|
||||
x = torch.randn(32, 35)
|
||||
|
|
@ -41,14 +55,20 @@ class TestAutoencoderArchitecture:
|
|||
assert out.max() <= 1.0
|
||||
|
||||
def test_reconstruction_error_shape(self) -> None:
|
||||
"""
|
||||
compute_reconstruction_error returns one scalar per sample in the batch.
|
||||
"""
|
||||
model = ThreatAutoencoder(input_dim=35)
|
||||
model.eval()
|
||||
x = torch.randn(8, 35)
|
||||
with torch.no_grad():
|
||||
errors = model.compute_reconstruction_error(x)
|
||||
assert errors.shape == (8,)
|
||||
assert errors.shape == (8, )
|
||||
|
||||
def test_reconstruction_error_positive(self) -> None:
|
||||
"""
|
||||
Reconstruction error is non-negative for all samples.
|
||||
"""
|
||||
model = ThreatAutoencoder(input_dim=35)
|
||||
model.eval()
|
||||
x = torch.randn(8, 35)
|
||||
|
|
@ -56,7 +76,11 @@ class TestAutoencoderArchitecture:
|
|||
errors = model.compute_reconstruction_error(x)
|
||||
assert (errors >= 0.0).all()
|
||||
|
||||
def test_trained_model_reconstructs_normal_better_than_anomaly(self) -> None:
|
||||
def test_trained_model_reconstructs_normal_better_than_anomaly(
|
||||
self) -> None:
|
||||
"""
|
||||
After training on normal data, reconstruction error is lower for normals than anomalies.
|
||||
"""
|
||||
torch.manual_seed(42)
|
||||
model = ThreatAutoencoder(input_dim=35)
|
||||
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
|
||||
|
|
@ -74,13 +98,17 @@ class TestAutoencoderArchitecture:
|
|||
|
||||
model.eval()
|
||||
with torch.no_grad():
|
||||
normal_errors = model.compute_reconstruction_error(normal_data[:50])
|
||||
normal_errors = model.compute_reconstruction_error(
|
||||
normal_data[:50])
|
||||
anomaly_data = torch.rand(50, 35) * 3.0 - 1.0
|
||||
anomaly_errors = model.compute_reconstruction_error(anomaly_data)
|
||||
|
||||
assert anomaly_errors.mean() > normal_errors.mean()
|
||||
|
||||
def test_eval_mode_disables_dropout(self) -> None:
|
||||
"""
|
||||
Identical inputs produce identical outputs in eval mode (dropout is off).
|
||||
"""
|
||||
model = ThreatAutoencoder(input_dim=35)
|
||||
model.eval()
|
||||
x = torch.randn(4, 35)
|
||||
|
|
@ -91,6 +119,9 @@ class TestAutoencoderArchitecture:
|
|||
|
||||
@pytest.mark.parametrize("batch_size", [1, 8, 32, 128])
|
||||
def test_variable_batch_sizes(self, batch_size: int) -> None:
|
||||
"""
|
||||
Output shape matches input for batch sizes 1, 8, 32, and 128.
|
||||
"""
|
||||
model = ThreatAutoencoder(input_dim=35)
|
||||
model.eval()
|
||||
x = torch.randn(batch_size, 35)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,81 @@
|
|||
"""
|
||||
©AngelaMos | 2026
|
||||
test_cli.py
|
||||
|
||||
Tests the Typer CLI commands: help output, argument validation, and error handling for missing paths.
|
||||
"""
|
||||
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from cli.main import app
|
||||
|
||||
runner = CliRunner()
|
||||
|
||||
|
||||
class TestCLICommands:
|
||||
|
||||
def test_train_help(self) -> None:
|
||||
"""
|
||||
train --help exits cleanly and mentions 'dataset' in its output.
|
||||
"""
|
||||
result = runner.invoke(app, ["train", "--help"])
|
||||
assert result.exit_code == 0
|
||||
assert "dataset" in result.output.lower()
|
||||
|
||||
def test_replay_help(self) -> None:
|
||||
"""
|
||||
replay --help exits cleanly and mentions 'log' in its output.
|
||||
"""
|
||||
result = runner.invoke(app, ["replay", "--help"])
|
||||
assert result.exit_code == 0
|
||||
assert "log" in result.output.lower()
|
||||
|
||||
def test_serve_help(self) -> None:
|
||||
"""
|
||||
serve --help exits cleanly and mentions 'host' in its output.
|
||||
"""
|
||||
result = runner.invoke(app, ["serve", "--help"])
|
||||
assert result.exit_code == 0
|
||||
assert "host" in result.output.lower()
|
||||
|
||||
def test_config_help(self) -> None:
|
||||
"""
|
||||
config --help exits cleanly.
|
||||
"""
|
||||
result = runner.invoke(app, ["config", "--help"])
|
||||
assert result.exit_code == 0
|
||||
|
||||
def test_health_help(self) -> None:
|
||||
"""
|
||||
health --help exits cleanly.
|
||||
"""
|
||||
result = runner.invoke(app, ["health", "--help"])
|
||||
assert result.exit_code == 0
|
||||
|
||||
def test_train_missing_dataset_fails(self) -> None:
|
||||
"""
|
||||
train with a non-existent dataset path exits with a non-zero code.
|
||||
"""
|
||||
result = runner.invoke(
|
||||
app,
|
||||
[
|
||||
"train",
|
||||
"--dataset",
|
||||
"/nonexistent/data.csv",
|
||||
],
|
||||
)
|
||||
assert result.exit_code != 0
|
||||
|
||||
def test_replay_missing_log_fails(self) -> None:
|
||||
"""
|
||||
replay with a non-existent log file exits with a non-zero code.
|
||||
"""
|
||||
result = runner.invoke(
|
||||
app,
|
||||
[
|
||||
"replay",
|
||||
"--log-file",
|
||||
"/nonexistent/access.log",
|
||||
],
|
||||
)
|
||||
assert result.exit_code != 0
|
||||
|
|
@ -1,31 +1,45 @@
|
|||
"""
|
||||
©AngelaMos | 2026
|
||||
test_config_ml.py
|
||||
|
||||
Tests ML-related settings defaults: detection mode, ensemble weights, model paths, and MLflow URI.
|
||||
"""
|
||||
|
||||
from app.config import settings
|
||||
|
||||
|
||||
def test_default_detection_mode_is_rules() -> None:
|
||||
"""
|
||||
Default detection_mode is 'rules' before any ML models are loaded.
|
||||
"""
|
||||
assert settings.detection_mode == "rules"
|
||||
|
||||
|
||||
def test_default_ensemble_weights_sum_to_one() -> None:
|
||||
total = (
|
||||
settings.ensemble_weight_ae
|
||||
+ settings.ensemble_weight_rf
|
||||
+ settings.ensemble_weight_if
|
||||
)
|
||||
"""
|
||||
AE + RF + IF ensemble weights sum to exactly 1.0.
|
||||
"""
|
||||
total = (settings.ensemble_weight_ae + settings.ensemble_weight_rf +
|
||||
settings.ensemble_weight_if)
|
||||
assert abs(total - 1.0) < 1e-6
|
||||
|
||||
|
||||
def test_default_model_dir() -> None:
|
||||
"""
|
||||
Default model artifact directory is data/models.
|
||||
"""
|
||||
assert settings.model_dir == "data/models"
|
||||
|
||||
|
||||
def test_default_ae_threshold_percentile() -> None:
|
||||
"""
|
||||
Default autoencoder threshold percentile is 99.5.
|
||||
"""
|
||||
assert settings.ae_threshold_percentile == 99.5
|
||||
|
||||
|
||||
def test_default_mlflow_tracking_uri() -> None:
|
||||
"""
|
||||
Default MLflow tracking URI uses local file storage.
|
||||
"""
|
||||
assert settings.mlflow_tracking_uri == "file:./mlruns"
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
"""
|
||||
©AngelaMos | 2026
|
||||
test_detection.py
|
||||
|
||||
Tests the rule engine's threat scoring, severity classification, and attack pattern matching.
|
||||
"""
|
||||
|
||||
from datetime import datetime, UTC
|
||||
|
|
|
|||
|
|
@ -0,0 +1,156 @@
|
|||
"""
|
||||
©AngelaMos | 2026
|
||||
test_ensemble.py
|
||||
|
||||
Tests ensemble score normalization, weighted fusion, ML/rule blending, and severity classification.
|
||||
"""
|
||||
|
||||
from app.core.detection.ensemble import (
|
||||
blend_scores,
|
||||
classify_severity,
|
||||
fuse_scores,
|
||||
normalize_ae_score,
|
||||
normalize_if_score,
|
||||
)
|
||||
|
||||
|
||||
class TestScoreNormalization:
|
||||
|
||||
def test_ae_score_below_threshold(self) -> None:
|
||||
"""
|
||||
AE error below threshold maps to a score below 0.5.
|
||||
"""
|
||||
result = normalize_ae_score(0.05, threshold=0.10)
|
||||
assert abs(result - 0.25) < 1e-6
|
||||
|
||||
def test_ae_score_above_double_threshold_caps_at_one(self) -> None:
|
||||
"""
|
||||
AE error at 3x threshold is capped at 1.0.
|
||||
"""
|
||||
result = normalize_ae_score(0.30, threshold=0.10)
|
||||
assert result == 1.0
|
||||
|
||||
def test_ae_score_zero_error(self) -> None:
|
||||
"""
|
||||
Zero reconstruction error maps to a score of 0.0.
|
||||
"""
|
||||
result = normalize_ae_score(0.0, threshold=0.10)
|
||||
assert result == 0.0
|
||||
|
||||
def test_if_score_negative(self) -> None:
|
||||
"""
|
||||
Negative IF score (anomalous region) maps above 0.5.
|
||||
"""
|
||||
result = normalize_if_score(-0.5)
|
||||
assert abs(result - 0.75) < 1e-6
|
||||
|
||||
def test_if_score_positive(self) -> None:
|
||||
"""
|
||||
Positive IF score (normal region) maps below 0.5.
|
||||
"""
|
||||
result = normalize_if_score(0.5)
|
||||
assert abs(result - 0.25) < 1e-6
|
||||
|
||||
def test_if_score_zero(self) -> None:
|
||||
"""
|
||||
Zero IF score maps to exactly 0.5.
|
||||
"""
|
||||
result = normalize_if_score(0.0)
|
||||
assert abs(result - 0.5) < 1e-6
|
||||
|
||||
|
||||
class TestEnsembleFusion:
|
||||
|
||||
def test_weighted_average(self) -> None:
|
||||
"""
|
||||
Fused score is the weighted average of AE, RF, and IF scores.
|
||||
"""
|
||||
scores = {"ae": 0.8, "rf": 0.6, "if": 0.5}
|
||||
weights = {"ae": 0.4, "rf": 0.4, "if": 0.2}
|
||||
result = fuse_scores(scores, weights)
|
||||
expected = 0.8 * 0.4 + 0.6 * 0.4 + 0.5 * 0.2
|
||||
assert abs(result - expected) < 1e-6
|
||||
|
||||
def test_all_zero_scores(self) -> None:
|
||||
"""
|
||||
All-zero model scores fuse to 0.0.
|
||||
"""
|
||||
scores = {"ae": 0.0, "rf": 0.0, "if": 0.0}
|
||||
weights = {"ae": 0.4, "rf": 0.4, "if": 0.2}
|
||||
assert fuse_scores(scores, weights) == 0.0
|
||||
|
||||
def test_all_max_scores(self) -> None:
|
||||
"""
|
||||
All-one model scores fuse to 1.0.
|
||||
"""
|
||||
scores = {"ae": 1.0, "rf": 1.0, "if": 1.0}
|
||||
weights = {"ae": 0.4, "rf": 0.4, "if": 0.2}
|
||||
assert abs(fuse_scores(scores, weights) - 1.0) < 1e-6
|
||||
|
||||
def test_partial_models(self) -> None:
|
||||
"""
|
||||
Fusion works correctly with only two models present.
|
||||
"""
|
||||
scores = {"ae": 0.9, "rf": 0.7}
|
||||
weights = {"ae": 0.5, "rf": 0.5}
|
||||
expected = 0.9 * 0.5 + 0.7 * 0.5
|
||||
assert abs(fuse_scores(scores, weights) - expected) < 1e-6
|
||||
|
||||
|
||||
class TestBlendScores:
|
||||
|
||||
def test_blend_with_rule_score(self) -> None:
|
||||
"""
|
||||
ML and rule scores blend according to the specified ml_weight.
|
||||
"""
|
||||
result = blend_scores(ml_score=0.7, rule_score=0.9, ml_weight=0.7)
|
||||
expected = 0.7 * 0.7 + 0.9 * 0.3
|
||||
assert abs(result - expected) < 1e-6
|
||||
|
||||
def test_blend_full_ml_weight(self) -> None:
|
||||
"""
|
||||
ml_weight of 1.0 returns the pure ML score.
|
||||
"""
|
||||
result = blend_scores(ml_score=0.8, rule_score=0.2, ml_weight=1.0)
|
||||
assert abs(result - 0.8) < 1e-6
|
||||
|
||||
def test_blend_full_rule_weight(self) -> None:
|
||||
"""
|
||||
ml_weight of 0.0 returns the pure rule score.
|
||||
"""
|
||||
result = blend_scores(ml_score=0.8, rule_score=0.2, ml_weight=0.0)
|
||||
assert abs(result - 0.2) < 1e-6
|
||||
|
||||
def test_blend_clamped_to_one(self) -> None:
|
||||
"""
|
||||
Blended score is clamped to 1.0 even when both inputs are 1.0.
|
||||
"""
|
||||
result = blend_scores(ml_score=1.0, rule_score=1.0, ml_weight=0.5)
|
||||
assert result <= 1.0
|
||||
|
||||
|
||||
class TestClassifySeverity:
|
||||
|
||||
def test_high(self) -> None:
|
||||
"""
|
||||
Scores >= 0.7 classify as HIGH.
|
||||
"""
|
||||
assert classify_severity(0.8) == "HIGH"
|
||||
assert classify_severity(0.7) == "HIGH"
|
||||
assert classify_severity(1.0) == "HIGH"
|
||||
|
||||
def test_medium(self) -> None:
|
||||
"""
|
||||
Scores in [0.5, 0.7) classify as MEDIUM.
|
||||
"""
|
||||
assert classify_severity(0.55) == "MEDIUM"
|
||||
assert classify_severity(0.5) == "MEDIUM"
|
||||
assert classify_severity(0.69) == "MEDIUM"
|
||||
|
||||
def test_low(self) -> None:
|
||||
"""
|
||||
Scores below 0.5 classify as LOW.
|
||||
"""
|
||||
assert classify_severity(0.3) == "LOW"
|
||||
assert classify_severity(0.0) == "LOW"
|
||||
assert classify_severity(0.49) == "LOW"
|
||||
|
|
@ -0,0 +1,101 @@
|
|||
"""
|
||||
©AngelaMos | 2026
|
||||
test_experiment.py
|
||||
|
||||
Tests the VigilExperiment MLflow wrapper: run lifecycle, param/metric logging, and status tagging.
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import mlflow
|
||||
import pytest
|
||||
|
||||
from ml.experiment import VigilExperiment
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _mlflow_tmp(tmp_path: Path) -> None:
|
||||
"""
|
||||
Point MLflow at a temp directory for isolation
|
||||
"""
|
||||
mlflow.set_tracking_uri(f"file:{tmp_path}/mlruns")
|
||||
|
||||
|
||||
class TestVigilExperiment:
|
||||
|
||||
def test_creates_run_with_id(self) -> None:
|
||||
"""
|
||||
Entering the context manager creates an MLflow run with a non-None run ID.
|
||||
"""
|
||||
with VigilExperiment("test-exp") as exp:
|
||||
assert exp.run_id is not None
|
||||
|
||||
def test_run_id_is_none_before_enter(self) -> None:
|
||||
"""
|
||||
run_id is None until the context manager is entered.
|
||||
"""
|
||||
exp = VigilExperiment("test-exp")
|
||||
assert exp.run_id is None
|
||||
|
||||
def test_log_params(self) -> None:
|
||||
"""
|
||||
log_params writes key-value pairs to the MLflow run as strings.
|
||||
"""
|
||||
with VigilExperiment("test-exp") as exp:
|
||||
exp.log_params({"lr": 0.001, "epochs": 10})
|
||||
run = mlflow.get_run(exp.run_id)
|
||||
assert run.data.params["lr"] == "0.001"
|
||||
assert run.data.params["epochs"] == "10"
|
||||
|
||||
def test_log_metrics(self) -> None:
|
||||
"""
|
||||
log_metrics stores numeric values on the MLflow run.
|
||||
"""
|
||||
with VigilExperiment("test-exp") as exp:
|
||||
exp.log_metrics({"f1": 0.95, "loss": 0.02})
|
||||
run = mlflow.get_run(exp.run_id)
|
||||
assert run.data.metrics["f1"] == 0.95
|
||||
|
||||
def test_log_artifact(self, tmp_path: Path) -> None:
|
||||
"""
|
||||
log_artifact uploads a file so it appears in the run's artifact list.
|
||||
"""
|
||||
artifact = tmp_path / "dummy.txt"
|
||||
artifact.write_text("test content")
|
||||
with VigilExperiment("test-exp") as exp:
|
||||
exp.log_artifact(artifact)
|
||||
run_id = exp.run_id
|
||||
client = mlflow.MlflowClient()
|
||||
artifacts = client.list_artifacts(run_id)
|
||||
names = [a.path for a in artifacts]
|
||||
assert "dummy.txt" in names
|
||||
|
||||
def test_system_metadata_logged(self) -> None:
|
||||
"""
|
||||
python_version and platform tags are added automatically on run start.
|
||||
"""
|
||||
with VigilExperiment("test-exp") as exp:
|
||||
run = mlflow.get_run(exp.run_id)
|
||||
assert "python_version" in run.data.tags
|
||||
assert "platform" in run.data.tags
|
||||
|
||||
def test_completed_status_on_success(self) -> None:
|
||||
"""
|
||||
Run tagged with status='completed' when the context exits cleanly.
|
||||
"""
|
||||
with VigilExperiment("test-exp") as exp:
|
||||
run_id = exp.run_id
|
||||
run = mlflow.get_run(run_id)
|
||||
assert run.data.tags["status"] == "completed"
|
||||
|
||||
def test_failed_status_on_exception(self) -> None:
|
||||
"""
|
||||
Run tagged with status='failed' and error message when an exception is raised.
|
||||
"""
|
||||
run_id = None
|
||||
with pytest.raises(ValueError), VigilExperiment("test-exp") as exp:
|
||||
run_id = exp.run_id
|
||||
raise ValueError("boom")
|
||||
run = mlflow.get_run(run_id)
|
||||
assert run.data.tags["status"] == "failed"
|
||||
assert "boom" in run.data.tags["error"]
|
||||
|
|
@ -0,0 +1,131 @@
|
|||
"""
|
||||
©AngelaMos | 2026
|
||||
test_export_onnx.py
|
||||
|
||||
Tests ONNX export for the autoencoder, random forest, and isolation forest models.
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
import torch
|
||||
from sklearn.ensemble import IsolationForest, RandomForestClassifier
|
||||
|
||||
from ml.autoencoder import ThreatAutoencoder
|
||||
from ml.export_onnx import (
|
||||
export_autoencoder,
|
||||
export_isolation_forest,
|
||||
export_random_forest,
|
||||
)
|
||||
|
||||
|
||||
class TestAutoencoderExport:
|
||||
|
||||
def test_creates_onnx_file(self, tmp_path: Path) -> None:
|
||||
"""
|
||||
Exporting creates a non-empty .onnx file at the given path.
|
||||
"""
|
||||
model = ThreatAutoencoder(input_dim=35)
|
||||
path = export_autoencoder(model, tmp_path / "ae.onnx")
|
||||
assert path.exists()
|
||||
assert path.stat().st_size > 0
|
||||
|
||||
def test_onnx_output_matches_pytorch(self, tmp_path: Path) -> None:
|
||||
"""
|
||||
ONNX inference output matches PyTorch forward pass within 1e-5 tolerance.
|
||||
"""
|
||||
torch.manual_seed(42)
|
||||
model = ThreatAutoencoder(input_dim=35)
|
||||
model.eval()
|
||||
onnx_path = export_autoencoder(model, tmp_path / "ae.onnx")
|
||||
|
||||
rng = np.random.default_rng(42)
|
||||
x = rng.standard_normal((8, 35)).astype(np.float32)
|
||||
|
||||
with torch.no_grad():
|
||||
pt_out = model(torch.from_numpy(x)).numpy()
|
||||
|
||||
session = ort.InferenceSession(str(onnx_path))
|
||||
ort_out = session.run(None, {"features": x})[0]
|
||||
|
||||
np.testing.assert_allclose(pt_out, ort_out, atol=1e-5)
|
||||
|
||||
def test_dynamic_batch_dimension(self, tmp_path: Path) -> None:
|
||||
"""
|
||||
Exported model accepts variable batch sizes (1, 16, 64).
|
||||
"""
|
||||
model = ThreatAutoencoder(input_dim=35)
|
||||
model.eval()
|
||||
onnx_path = export_autoencoder(model, tmp_path / "ae.onnx")
|
||||
session = ort.InferenceSession(str(onnx_path))
|
||||
|
||||
for batch_size in (1, 16, 64):
|
||||
rng = np.random.default_rng(batch_size)
|
||||
x = rng.standard_normal((batch_size, 35)).astype(np.float32)
|
||||
out = session.run(None, {"features": x})[0]
|
||||
assert out.shape == (batch_size, 35)
|
||||
|
||||
|
||||
class TestRandomForestExport:
|
||||
|
||||
def test_creates_onnx_file(self, tmp_path: Path) -> None:
|
||||
"""
|
||||
Exporting a fitted RandomForest creates a non-empty .onnx file.
|
||||
"""
|
||||
rng = np.random.default_rng(42)
|
||||
rf = RandomForestClassifier(n_estimators=10, random_state=42)
|
||||
X = rng.standard_normal((100, 35)).astype(np.float32)
|
||||
y = np.concatenate([np.zeros(70, dtype=int), np.ones(30, dtype=int)])
|
||||
rf.fit(X, y)
|
||||
path = export_random_forest(rf, 35, tmp_path / "rf.onnx")
|
||||
assert path.exists()
|
||||
assert path.stat().st_size > 0
|
||||
|
||||
def test_onnx_produces_valid_output(self, tmp_path: Path) -> None:
|
||||
"""
|
||||
ONNX inference returns class predictions and probabilities for each sample.
|
||||
"""
|
||||
rng = np.random.default_rng(42)
|
||||
rf = RandomForestClassifier(n_estimators=10, random_state=42)
|
||||
X = rng.standard_normal((100, 35)).astype(np.float32)
|
||||
y = np.concatenate([np.zeros(70, dtype=int), np.ones(30, dtype=int)])
|
||||
rf.fit(X, y)
|
||||
onnx_path = export_random_forest(rf, 35, tmp_path / "rf.onnx")
|
||||
session = ort.InferenceSession(str(onnx_path))
|
||||
|
||||
x_test = rng.standard_normal((5, 35)).astype(np.float32)
|
||||
result = session.run(None, {"features": x_test})
|
||||
assert len(result) == 2
|
||||
assert len(result[0]) == 5
|
||||
|
||||
|
||||
class TestIsolationForestExport:
|
||||
|
||||
def test_creates_onnx_file(self, tmp_path: Path) -> None:
|
||||
"""
|
||||
Exporting a fitted IsolationForest creates a non-empty .onnx file.
|
||||
"""
|
||||
rng = np.random.default_rng(42)
|
||||
iso = IsolationForest(n_estimators=10, random_state=42)
|
||||
iso.fit(rng.standard_normal((100, 35)).astype(np.float32))
|
||||
path = export_isolation_forest(iso, 35, tmp_path / "if.onnx")
|
||||
assert path.exists()
|
||||
assert path.stat().st_size > 0
|
||||
|
||||
def test_onnx_scores_match_decision_function(self, tmp_path: Path) -> None:
|
||||
"""
|
||||
ONNX anomaly scores match sklearn decision_function within 1e-4 tolerance.
|
||||
"""
|
||||
rng = np.random.default_rng(42)
|
||||
iso = IsolationForest(n_estimators=10, random_state=42)
|
||||
X = rng.standard_normal((100, 35)).astype(np.float32)
|
||||
iso.fit(X)
|
||||
onnx_path = export_isolation_forest(iso, 35, tmp_path / "if.onnx")
|
||||
session = ort.InferenceSession(str(onnx_path))
|
||||
|
||||
x_test = rng.standard_normal((10, 35)).astype(np.float32)
|
||||
sk_decision = iso.decision_function(x_test)
|
||||
ort_scores = session.run(None, {"features": x_test})[1].flatten()
|
||||
|
||||
np.testing.assert_allclose(sk_decision, ort_scores, atol=1e-4)
|
||||
|
|
@ -1,6 +1,8 @@
|
|||
"""
|
||||
©AngelaMos | 2026
|
||||
test_features.py
|
||||
|
||||
Tests per-request feature extraction, Redis sliding-window aggregation, and feature encoding.
|
||||
"""
|
||||
|
||||
from datetime import datetime, UTC
|
||||
|
|
@ -79,16 +81,20 @@ def test_path_depth() -> None:
|
|||
Path depth counts non-empty segments between slashes.
|
||||
"""
|
||||
assert extract_request_features(_make_entry(path="/"))["path_depth"] == 0
|
||||
assert extract_request_features(_make_entry(path="/api"))["path_depth"] == 1
|
||||
assert extract_request_features(_make_entry(path="/api/v1/users"))["path_depth"] == 3
|
||||
assert extract_request_features(
|
||||
_make_entry(path="/api"))["path_depth"] == 1
|
||||
assert extract_request_features(
|
||||
_make_entry(path="/api/v1/users"))["path_depth"] == 3
|
||||
|
||||
|
||||
def test_path_entropy_high_vs_low() -> None:
|
||||
"""
|
||||
Random-character paths have higher entropy than simple paths.
|
||||
"""
|
||||
low = extract_request_features(_make_entry(path="/index.html"))["path_entropy"]
|
||||
high = extract_request_features(_make_entry(path="/x8Kp2mQz7wR4vL1n"))["path_entropy"]
|
||||
low = extract_request_features(
|
||||
_make_entry(path="/index.html"))["path_entropy"]
|
||||
high = extract_request_features(
|
||||
_make_entry(path="/x8Kp2mQz7wR4vL1n"))["path_entropy"]
|
||||
assert high > low
|
||||
|
||||
|
||||
|
|
@ -96,7 +102,8 @@ def test_query_string_features() -> None:
|
|||
"""
|
||||
Query param count and length are extracted correctly.
|
||||
"""
|
||||
features = extract_request_features(_make_entry(query_string="page=1&sort=name&limit=50"))
|
||||
features = extract_request_features(
|
||||
_make_entry(query_string="page=1&sort=name&limit=50"))
|
||||
assert features["query_param_count"] == 3
|
||||
assert features["query_string_length"] == len("page=1&sort=name&limit=50")
|
||||
|
||||
|
|
@ -110,11 +117,11 @@ def test_url_encoding_detection() -> None:
|
|||
Percent-encoded sequences are detected in path and query.
|
||||
"""
|
||||
encoded = extract_request_features(
|
||||
_make_entry(path="/search", query_string="q=%27OR+1%3D1")
|
||||
)
|
||||
_make_entry(path="/search", query_string="q=%27OR+1%3D1"))
|
||||
assert encoded["has_encoded_chars"] is True
|
||||
|
||||
clean = extract_request_features(_make_entry(path="/index.html", query_string=""))
|
||||
clean = extract_request_features(
|
||||
_make_entry(path="/index.html", query_string=""))
|
||||
assert clean["has_encoded_chars"] is False
|
||||
|
||||
|
||||
|
|
@ -133,9 +140,12 @@ def test_status_class() -> None:
|
|||
"""
|
||||
Status class groups status codes into Nxx buckets.
|
||||
"""
|
||||
assert extract_request_features(_make_entry(status_code=200))["status_class"] == "2xx"
|
||||
assert extract_request_features(_make_entry(status_code=404))["status_class"] == "4xx"
|
||||
assert extract_request_features(_make_entry(status_code=503))["status_class"] == "5xx"
|
||||
assert extract_request_features(
|
||||
_make_entry(status_code=200))["status_class"] == "2xx"
|
||||
assert extract_request_features(
|
||||
_make_entry(status_code=404))["status_class"] == "4xx"
|
||||
assert extract_request_features(
|
||||
_make_entry(status_code=503))["status_class"] == "5xx"
|
||||
|
||||
|
||||
def test_temporal_features() -> None:
|
||||
|
|
@ -161,14 +171,13 @@ def test_ua_bot_detection() -> None:
|
|||
"""
|
||||
bot = extract_request_features(
|
||||
_make_entry(
|
||||
user_agent="Mozilla/5.0 (compatible; Googlebot/2.1; +http://www.google.com/bot.html)"
|
||||
)
|
||||
)
|
||||
user_agent=
|
||||
"Mozilla/5.0 (compatible; Googlebot/2.1; +http://www.google.com/bot.html)"
|
||||
))
|
||||
assert bot["is_known_bot"] is True
|
||||
|
||||
normal = extract_request_features(
|
||||
_make_entry(user_agent="Mozilla/5.0 (Windows NT 10.0; Win64; x64)")
|
||||
)
|
||||
_make_entry(user_agent="Mozilla/5.0 (Windows NT 10.0; Win64; x64)"))
|
||||
assert normal["is_known_bot"] is False
|
||||
|
||||
|
||||
|
|
@ -176,7 +185,8 @@ def test_ua_scanner_detection() -> None:
|
|||
"""
|
||||
Known vulnerability scanner user agents are flagged.
|
||||
"""
|
||||
nikto = extract_request_features(_make_entry(user_agent="Mozilla/5.00 (Nikto/2.1.6)"))
|
||||
nikto = extract_request_features(
|
||||
_make_entry(user_agent="Mozilla/5.00 (Nikto/2.1.6)"))
|
||||
assert nikto["is_known_scanner"] is True
|
||||
|
||||
sqlmap = extract_request_features(_make_entry(user_agent="sqlmap/1.8"))
|
||||
|
|
@ -187,15 +197,17 @@ def test_attack_pattern_detection() -> None:
|
|||
"""
|
||||
SQLi, XSS, and path traversal patterns in paths are detected.
|
||||
"""
|
||||
sqli = extract_request_features(_make_entry(path="/users", query_string="id=1' OR 1=1--"))
|
||||
sqli = extract_request_features(
|
||||
_make_entry(path="/users", query_string="id=1' OR 1=1--"))
|
||||
assert sqli["has_attack_pattern"] is True
|
||||
|
||||
xss = extract_request_features(
|
||||
_make_entry(path="/comment", query_string="body=<script>alert(1)</script>")
|
||||
)
|
||||
_make_entry(path="/comment",
|
||||
query_string="body=<script>alert(1)</script>"))
|
||||
assert xss["has_attack_pattern"] is True
|
||||
|
||||
traversal = extract_request_features(_make_entry(path="/static/../../etc/passwd"))
|
||||
traversal = extract_request_features(
|
||||
_make_entry(path="/static/../../etc/passwd"))
|
||||
assert traversal["has_attack_pattern"] is True
|
||||
|
||||
clean = extract_request_features(_make_entry(path="/api/v1/health"))
|
||||
|
|
@ -206,11 +218,12 @@ def test_special_char_ratio() -> None:
|
|||
"""
|
||||
Paths with many non-alphanumeric characters have higher ratios.
|
||||
"""
|
||||
clean = extract_request_features(_make_entry(path="/api/users"))["special_char_ratio"]
|
||||
clean = extract_request_features(
|
||||
_make_entry(path="/api/users"))["special_char_ratio"]
|
||||
|
||||
noisy = extract_request_features(_make_entry(path="/<script>alert('xss')</script>"))[
|
||||
"special_char_ratio"
|
||||
]
|
||||
noisy = extract_request_features(
|
||||
_make_entry(
|
||||
path="/<script>alert('xss')</script>"))["special_char_ratio"]
|
||||
|
||||
assert noisy > clean
|
||||
|
||||
|
|
@ -219,20 +232,25 @@ def test_private_ip_detection() -> None:
|
|||
"""
|
||||
RFC 1918 and loopback addresses are classified as private.
|
||||
"""
|
||||
assert extract_request_features(_make_entry(ip="192.168.1.1"))["is_private_ip"] is True
|
||||
assert extract_request_features(
|
||||
_make_entry(ip="192.168.1.1"))["is_private_ip"] is True
|
||||
|
||||
assert extract_request_features(_make_entry(ip="127.0.0.1"))["is_private_ip"] is True
|
||||
assert extract_request_features(
|
||||
_make_entry(ip="127.0.0.1"))["is_private_ip"] is True
|
||||
|
||||
assert extract_request_features(_make_entry(ip="8.8.8.8"))["is_private_ip"] is False
|
||||
assert extract_request_features(
|
||||
_make_entry(ip="8.8.8.8"))["is_private_ip"] is False
|
||||
|
||||
|
||||
def test_file_extension() -> None:
|
||||
"""
|
||||
File extension is extracted from the path.
|
||||
"""
|
||||
assert extract_request_features(_make_entry(path="/style.css"))["file_extension"] == ".css"
|
||||
assert extract_request_features(
|
||||
_make_entry(path="/style.css"))["file_extension"] == ".css"
|
||||
|
||||
assert extract_request_features(_make_entry(path="/api/users"))["file_extension"] == ""
|
||||
assert extract_request_features(
|
||||
_make_entry(path="/api/users"))["file_extension"] == ""
|
||||
|
||||
|
||||
def test_country_code_passthrough() -> None:
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
"""
|
||||
©AngelaMos | 2026
|
||||
test_geoip.py
|
||||
|
||||
Tests the GeoIP lookup service including private IP handling and missing database fallback.
|
||||
"""
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
|
@ -108,7 +110,9 @@ async def test_lookup_missing_city_name(mock_reader) -> None:
|
|||
"""
|
||||
Responses with None city name are handled gracefully.
|
||||
"""
|
||||
mock_reader.city.return_value = _mock_city_response(city_name=None, lat=0.0, lon=0.0)
|
||||
mock_reader.city.return_value = _mock_city_response(city_name=None,
|
||||
lat=0.0,
|
||||
lon=0.0)
|
||||
service = GeoIPService.__new__(GeoIPService)
|
||||
service._reader = mock_reader
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,142 @@
|
|||
"""
|
||||
©AngelaMos | 2026
|
||||
test_inference.py
|
||||
|
||||
Tests the InferenceEngine: model loading, predict output shapes, score ranges, and missing-model handling.
|
||||
"""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from ml.autoencoder import ThreatAutoencoder
|
||||
from ml.export_onnx import (
|
||||
export_autoencoder,
|
||||
export_isolation_forest,
|
||||
export_random_forest,
|
||||
)
|
||||
from ml.scaler import FeatureScaler
|
||||
from sklearn.ensemble import IsolationForest, RandomForestClassifier
|
||||
|
||||
from app.core.detection.inference import InferenceEngine
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def model_dir(tmp_path: Path) -> Path:
|
||||
"""
|
||||
Create a temp directory with all 3 ONNX models + scaler + threshold
|
||||
"""
|
||||
rng = np.random.default_rng(42)
|
||||
X = rng.standard_normal((200, 35)).astype(np.float32)
|
||||
y = np.concatenate([np.zeros(140, dtype=int), np.ones(60, dtype=int)])
|
||||
|
||||
ae = ThreatAutoencoder(input_dim=35)
|
||||
export_autoencoder(ae, tmp_path / "ae.onnx")
|
||||
|
||||
rf = RandomForestClassifier(n_estimators=10, random_state=42)
|
||||
rf.fit(X, y)
|
||||
export_random_forest(rf, 35, tmp_path / "rf.onnx")
|
||||
|
||||
iso = IsolationForest(n_estimators=10, random_state=42)
|
||||
iso.fit(X[:140])
|
||||
export_isolation_forest(iso, 35, tmp_path / "if.onnx")
|
||||
|
||||
scaler = FeatureScaler()
|
||||
scaler.fit(X[:140])
|
||||
scaler.save_json(tmp_path / "scaler.json")
|
||||
|
||||
threshold_data = {"threshold": 0.05}
|
||||
(tmp_path / "threshold.json").write_text(json.dumps(threshold_data))
|
||||
|
||||
return tmp_path
|
||||
|
||||
|
||||
class TestInferenceEngine:
|
||||
|
||||
def test_loads_all_models(self, model_dir: Path) -> None:
|
||||
"""
|
||||
Engine reports is_loaded=True when all three ONNX models are present.
|
||||
"""
|
||||
engine = InferenceEngine(model_dir=str(model_dir))
|
||||
assert engine.is_loaded
|
||||
|
||||
def test_returns_none_when_no_models(self) -> None:
|
||||
"""
|
||||
Engine reports is_loaded=False when the model directory does not exist.
|
||||
"""
|
||||
engine = InferenceEngine(model_dir="/nonexistent/path")
|
||||
assert not engine.is_loaded
|
||||
|
||||
def test_predict_returns_none_when_not_loaded(self) -> None:
|
||||
"""
|
||||
predict returns None when the engine has no models loaded.
|
||||
"""
|
||||
engine = InferenceEngine(model_dir="/nonexistent/path")
|
||||
result = engine.predict(np.zeros((1, 35), dtype=np.float32))
|
||||
assert result is None
|
||||
|
||||
def test_predict_returns_scores(self, model_dir: Path) -> None:
|
||||
"""
|
||||
predict returns a dict with ae, rf, and if score arrays.
|
||||
"""
|
||||
engine = InferenceEngine(model_dir=str(model_dir))
|
||||
rng = np.random.default_rng(99)
|
||||
x = rng.standard_normal((4, 35)).astype(np.float32)
|
||||
result = engine.predict(x)
|
||||
assert result is not None
|
||||
assert "ae" in result
|
||||
assert "rf" in result
|
||||
assert "if" in result
|
||||
|
||||
def test_predict_ae_scores_are_positive(self, model_dir: Path) -> None:
|
||||
"""
|
||||
AE reconstruction error scores are non-negative for all samples.
|
||||
"""
|
||||
engine = InferenceEngine(model_dir=str(model_dir))
|
||||
rng = np.random.default_rng(99)
|
||||
x = rng.standard_normal((4, 35)).astype(np.float32)
|
||||
result = engine.predict(x)
|
||||
assert result is not None
|
||||
assert all(s >= 0.0 for s in result["ae"])
|
||||
|
||||
def test_predict_rf_probabilities_in_range(self, model_dir: Path) -> None:
|
||||
"""
|
||||
RF malicious-class probabilities are within [0, 1].
|
||||
"""
|
||||
engine = InferenceEngine(model_dir=str(model_dir))
|
||||
rng = np.random.default_rng(99)
|
||||
x = rng.standard_normal((4, 35)).astype(np.float32)
|
||||
result = engine.predict(x)
|
||||
assert result is not None
|
||||
assert all(0.0 <= p <= 1.0 for p in result["rf"])
|
||||
|
||||
def test_predict_single_sample(self, model_dir: Path) -> None:
|
||||
"""
|
||||
predict works on a single sample and returns one score per model.
|
||||
"""
|
||||
engine = InferenceEngine(model_dir=str(model_dir))
|
||||
rng = np.random.default_rng(99)
|
||||
x = rng.standard_normal((1, 35)).astype(np.float32)
|
||||
result = engine.predict(x)
|
||||
assert result is not None
|
||||
assert len(result["ae"]) == 1
|
||||
assert len(result["rf"]) == 1
|
||||
assert len(result["if"]) == 1
|
||||
|
||||
def test_threshold_loaded(self, model_dir: Path) -> None:
|
||||
"""
|
||||
Autoencoder threshold is read from threshold.json on initialization.
|
||||
"""
|
||||
engine = InferenceEngine(model_dir=str(model_dir))
|
||||
assert engine.threshold == 0.05
|
||||
|
||||
def test_partial_models_not_loaded(self, tmp_path: Path) -> None:
|
||||
"""
|
||||
Engine with only the AE model present reports is_loaded=False.
|
||||
"""
|
||||
ae = ThreatAutoencoder(input_dim=35)
|
||||
export_autoencoder(ae, tmp_path / "ae.onnx")
|
||||
engine = InferenceEngine(model_dir=str(tmp_path))
|
||||
assert not engine.is_loaded
|
||||
|
|
@ -1,6 +1,8 @@
|
|||
"""
|
||||
©AngelaMos | 2026
|
||||
test_integration.py
|
||||
|
||||
End-to-end tests covering the full path from log file write through tailer, pipeline, and database storage.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
|
@ -22,29 +24,22 @@ from app.core.ingestion.pipeline import Pipeline
|
|||
from app.core.ingestion.tailer import LogTailer
|
||||
from app.models.threat_event import ThreatEvent
|
||||
|
||||
NORMAL_LINE = (
|
||||
"192.168.1.100 - - [11/Feb/2026:10:00:00 +0000] "
|
||||
'"GET /index.html HTTP/1.1" 200 4523 "-" '
|
||||
'"Mozilla/5.0 (Windows NT 10.0; Win64; x64)"'
|
||||
)
|
||||
NORMAL_LINE = ("192.168.1.100 - - [11/Feb/2026:10:00:00 +0000] "
|
||||
'"GET /index.html HTTP/1.1" 200 4523 "-" '
|
||||
'"Mozilla/5.0 (Windows NT 10.0; Win64; x64)"')
|
||||
|
||||
SQLI_LINE = (
|
||||
"198.51.100.10 - - [11/Feb/2026:10:00:01 +0000] "
|
||||
'"GET /search?q=1%27+OR+1=1-- HTTP/1.1" 200 5678 "-" '
|
||||
'"Mozilla/5.0 (Windows NT 10.0; Win64; x64)"'
|
||||
)
|
||||
SQLI_LINE = ("198.51.100.10 - - [11/Feb/2026:10:00:01 +0000] "
|
||||
'"GET /search?q=1%27+OR+1=1-- HTTP/1.1" 200 5678 "-" '
|
||||
'"Mozilla/5.0 (Windows NT 10.0; Win64; x64)"')
|
||||
|
||||
XSS_LINE = (
|
||||
"198.51.100.11 - - [11/Feb/2026:10:00:02 +0000] "
|
||||
'"GET /comment?text=<script>alert(1)</script> HTTP/1.1" 200 3210 "-" '
|
||||
'"Mozilla/5.0 (Windows NT 10.0; Win64; x64)"'
|
||||
)
|
||||
'"Mozilla/5.0 (Windows NT 10.0; Win64; x64)"')
|
||||
|
||||
PATH_TRAVERSAL_LINE = (
|
||||
"198.51.100.12 - - [11/Feb/2026:10:00:03 +0000] "
|
||||
'"GET /../../etc/passwd HTTP/1.1" 400 230 "-" '
|
||||
'"Mozilla/5.0 (Windows NT 10.0; Win64; x64)"'
|
||||
)
|
||||
PATH_TRAVERSAL_LINE = ("198.51.100.12 - - [11/Feb/2026:10:00:03 +0000] "
|
||||
'"GET /../../etc/passwd HTTP/1.1" 400 230 "-" '
|
||||
'"Mozilla/5.0 (Windows NT 10.0; Win64; x64)"')
|
||||
|
||||
|
||||
def _write_lines(log_path: str, *lines: str) -> None:
|
||||
|
|
@ -127,7 +122,8 @@ async def _poll_threat_count(
|
|||
for _ in range(int(timeout / 0.1)):
|
||||
await asyncio.sleep(0.1)
|
||||
async with session_factory() as session:
|
||||
result = await session.execute(select(func.count()).select_from(ThreatEvent))
|
||||
result = await session.execute(
|
||||
select(func.count()).select_from(ThreatEvent))
|
||||
count = result.scalar_one()
|
||||
if count >= expected:
|
||||
return count
|
||||
|
|
@ -163,15 +159,15 @@ async def test_only_medium_plus_stored(integration_env) -> None:
|
|||
lines = [
|
||||
f"192.168.1.{i + 1} - - [11/Feb/2026:10:00:0{i} +0000] "
|
||||
f'"GET /page/{i} HTTP/1.1" 200 1234 "-" '
|
||||
f'"Mozilla/5.0 (Windows NT 10.0; Win64; x64)"'
|
||||
for i in range(5)
|
||||
f'"Mozilla/5.0 (Windows NT 10.0; Win64; x64)"' for i in range(5)
|
||||
]
|
||||
_write_lines(env["log_path"], *lines)
|
||||
|
||||
await asyncio.sleep(2.0)
|
||||
|
||||
async with env["session_factory"]() as session:
|
||||
result = await session.execute(select(func.count()).select_from(ThreatEvent))
|
||||
result = await session.execute(
|
||||
select(func.count()).select_from(ThreatEvent))
|
||||
count = result.scalar_one()
|
||||
|
||||
assert count == 0
|
||||
|
|
|
|||
|
|
@ -0,0 +1,179 @@
|
|||
"""
|
||||
©AngelaMos | 2026
|
||||
test_ml_integration.py
|
||||
"""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import fakeredis.aioredis
|
||||
import numpy as np
|
||||
import pytest
|
||||
from sklearn.ensemble import (
|
||||
IsolationForest,
|
||||
RandomForestClassifier,
|
||||
)
|
||||
|
||||
from app.core.detection.inference import InferenceEngine
|
||||
from app.core.detection.rules import RuleEngine
|
||||
from app.core.ingestion.pipeline import (
|
||||
Pipeline,
|
||||
ScoredRequest,
|
||||
)
|
||||
from ml.autoencoder import ThreatAutoencoder
|
||||
from ml.export_onnx import (
|
||||
export_autoencoder,
|
||||
export_isolation_forest,
|
||||
export_random_forest,
|
||||
)
|
||||
from ml.scaler import FeatureScaler
|
||||
|
||||
VALID_LINE = ("93.184.216.34 - - [11/Feb/2026:14:30:00 +0000] "
|
||||
'"GET /api/v1/users HTTP/1.1" 200 1234 '
|
||||
'"https://example.com" '
|
||||
'"Mozilla/5.0 (Windows NT 10.0; Win64; x64)"')
|
||||
|
||||
SQLI_LINE = ("93.184.216.34 - - [11/Feb/2026:14:30:01 +0000] "
|
||||
'"GET /users?id=1%27OR%201=1-- HTTP/1.1" 200 512 '
|
||||
'"-" "Mozilla/5.0"')
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def trained_model_dir(tmp_path: Path) -> Path:
|
||||
"""
|
||||
Create a temp directory with trained ONNX models,
|
||||
scaler, and threshold
|
||||
"""
|
||||
rng = np.random.default_rng(42)
|
||||
X = rng.standard_normal((200, 35)).astype(np.float32)
|
||||
y = np.concatenate([np.zeros(140, dtype=int), np.ones(60, dtype=int)])
|
||||
|
||||
ae = ThreatAutoencoder(input_dim=35)
|
||||
export_autoencoder(ae, tmp_path / "ae.onnx")
|
||||
|
||||
rf = RandomForestClassifier(n_estimators=10, random_state=42)
|
||||
rf.fit(X, y)
|
||||
export_random_forest(rf, 35, tmp_path / "rf.onnx")
|
||||
|
||||
iso = IsolationForest(n_estimators=10, random_state=42)
|
||||
iso.fit(X[:140])
|
||||
export_isolation_forest(iso, 35, tmp_path / "if.onnx")
|
||||
|
||||
scaler = FeatureScaler()
|
||||
scaler.fit(X[:140])
|
||||
scaler.save_json(tmp_path / "scaler.json")
|
||||
|
||||
threshold_data = {"threshold": 0.05}
|
||||
(tmp_path / "threshold.json").write_text(json.dumps(threshold_data))
|
||||
|
||||
return tmp_path
|
||||
|
||||
|
||||
async def _make_pipeline_with_ml(
|
||||
results: list[ScoredRequest],
|
||||
model_dir: Path,
|
||||
) -> Pipeline:
|
||||
"""
|
||||
Build a Pipeline with ML inference engine wired in
|
||||
"""
|
||||
redis = fakeredis.aioredis.FakeRedis(decode_responses=True)
|
||||
engine = InferenceEngine(model_dir=str(model_dir))
|
||||
|
||||
async def collect(sr: ScoredRequest) -> None:
|
||||
results.append(sr)
|
||||
|
||||
pipeline = Pipeline(
|
||||
redis_client=redis,
|
||||
rule_engine=RuleEngine(),
|
||||
on_result=collect,
|
||||
inference_engine=engine,
|
||||
)
|
||||
await pipeline.start()
|
||||
return pipeline
|
||||
|
||||
|
||||
async def _make_pipeline_rules_only(
|
||||
results: list[ScoredRequest], ) -> Pipeline:
|
||||
"""
|
||||
Build a Pipeline without ML inference engine
|
||||
"""
|
||||
redis = fakeredis.aioredis.FakeRedis(decode_responses=True)
|
||||
|
||||
async def collect(sr: ScoredRequest) -> None:
|
||||
results.append(sr)
|
||||
|
||||
pipeline = Pipeline(
|
||||
redis_client=redis,
|
||||
rule_engine=RuleEngine(),
|
||||
on_result=collect,
|
||||
)
|
||||
await pipeline.start()
|
||||
return pipeline
|
||||
|
||||
|
||||
class TestMLIntegration:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pipeline_with_ml_sets_hybrid_mode(
|
||||
self, trained_model_dir: Path) -> None:
|
||||
results: list[ScoredRequest] = []
|
||||
pipeline = await _make_pipeline_with_ml(results, trained_model_dir)
|
||||
|
||||
await pipeline.raw_queue.put(VALID_LINE)
|
||||
await pipeline.stop()
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].detection_mode == "hybrid"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pipeline_with_ml_produces_final_score(
|
||||
self, trained_model_dir: Path) -> None:
|
||||
results: list[ScoredRequest] = []
|
||||
pipeline = await _make_pipeline_with_ml(results, trained_model_dir)
|
||||
|
||||
await pipeline.raw_queue.put(VALID_LINE)
|
||||
await pipeline.stop()
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].final_score >= 0.0
|
||||
assert results[0].final_score <= 1.0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rules_only_uses_rule_score(self, ) -> None:
|
||||
results: list[ScoredRequest] = []
|
||||
pipeline = await _make_pipeline_rules_only(results)
|
||||
|
||||
await pipeline.raw_queue.put(VALID_LINE)
|
||||
await pipeline.stop()
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].detection_mode == "rules"
|
||||
assert results[0].final_score == results[0].rule_result.threat_score
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_attack_scores_higher_than_benign(
|
||||
self, trained_model_dir: Path) -> None:
|
||||
results: list[ScoredRequest] = []
|
||||
pipeline = await _make_pipeline_with_ml(results, trained_model_dir)
|
||||
|
||||
await pipeline.raw_queue.put(VALID_LINE)
|
||||
await pipeline.raw_queue.put(SQLI_LINE)
|
||||
await pipeline.stop()
|
||||
|
||||
assert len(results) == 2
|
||||
benign_score = results[0].final_score
|
||||
attack_score = results[1].final_score
|
||||
assert attack_score > benign_score
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rule_result_preserved_in_hybrid_mode(
|
||||
self, trained_model_dir: Path) -> None:
|
||||
results: list[ScoredRequest] = []
|
||||
pipeline = await _make_pipeline_with_ml(results, trained_model_dir)
|
||||
|
||||
await pipeline.raw_queue.put(SQLI_LINE)
|
||||
await pipeline.stop()
|
||||
|
||||
assert len(results) == 1
|
||||
assert "SQL_INJECTION" in results[0].rule_result.matched_rules
|
||||
assert results[0].rule_result.threat_score > 0
|
||||
|
|
@ -1,6 +1,8 @@
|
|||
"""
|
||||
©AngelaMos | 2026
|
||||
test_parsers.py
|
||||
|
||||
Tests nginx combined log line parsing via parse_combined.
|
||||
"""
|
||||
|
||||
from datetime import datetime, UTC
|
||||
|
|
@ -12,12 +14,10 @@ def test_parse_standard_combined_line() -> None:
|
|||
"""
|
||||
Parse a standard nginx combined log line into all fields.
|
||||
"""
|
||||
line = (
|
||||
"93.184.216.34 - - [11/Feb/2026:14:30:00 +0000] "
|
||||
'"GET /api/users?page=1 HTTP/1.1" 200 1234 '
|
||||
'"https://example.com" '
|
||||
'"Mozilla/5.0 (Windows NT 10.0; Win64; x64)"'
|
||||
)
|
||||
line = ("93.184.216.34 - - [11/Feb/2026:14:30:00 +0000] "
|
||||
'"GET /api/users?page=1 HTTP/1.1" 200 1234 '
|
||||
'"https://example.com" '
|
||||
'"Mozilla/5.0 (Windows NT 10.0; Win64; x64)"')
|
||||
result = parse_combined(line)
|
||||
|
||||
assert result is not None
|
||||
|
|
@ -52,10 +52,8 @@ def test_parse_missing_referer() -> None:
|
|||
"""
|
||||
A dash referer is normalized to an empty string.
|
||||
"""
|
||||
line = (
|
||||
"10.0.0.1 - - [11/Feb/2026:08:15:42 +0000] "
|
||||
'"POST /login HTTP/1.1" 302 0 "-" "Mozilla/5.0"'
|
||||
)
|
||||
line = ("10.0.0.1 - - [11/Feb/2026:08:15:42 +0000] "
|
||||
'"POST /login HTTP/1.1" 302 0 "-" "Mozilla/5.0"')
|
||||
result = parse_combined(line)
|
||||
|
||||
assert result is not None
|
||||
|
|
@ -71,8 +69,7 @@ def test_parse_complex_query_string() -> None:
|
|||
line = (
|
||||
"93.184.216.34 - - [11/Feb/2026:14:30:00 +0000] "
|
||||
'"GET /search?q=hello+world&lang=en&page=2&sort=relevance HTTP/1.1" '
|
||||
'200 5678 "https://example.com/search" "Mozilla/5.0"'
|
||||
)
|
||||
'200 5678 "https://example.com/search" "Mozilla/5.0"')
|
||||
result = parse_combined(line)
|
||||
|
||||
assert result is not None
|
||||
|
|
@ -110,11 +107,9 @@ def test_parse_full_ipv6_address() -> None:
|
|||
"""
|
||||
Parse a line with a full-length IPv6 address.
|
||||
"""
|
||||
line = (
|
||||
"2001:0db8:85a3:0000:0000:8a2e:0370:7334 - - "
|
||||
'[11/Feb/2026:14:30:00 +0000] "GET /api/v1/health HTTP/2.0" '
|
||||
'200 256 "-" "python-httpx/0.28"'
|
||||
)
|
||||
line = ("2001:0db8:85a3:0000:0000:8a2e:0370:7334 - - "
|
||||
'[11/Feb/2026:14:30:00 +0000] "GET /api/v1/health HTTP/2.0" '
|
||||
'200 256 "-" "python-httpx/0.28"')
|
||||
result = parse_combined(line)
|
||||
|
||||
assert result is not None
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
"""
|
||||
©AngelaMos | 2026
|
||||
test_pipeline.py
|
||||
|
||||
Tests the async ingestion pipeline: parsing, feature extraction, rule scoring, and shutdown.
|
||||
"""
|
||||
|
||||
import fakeredis.aioredis
|
||||
|
|
@ -9,18 +11,14 @@ import pytest
|
|||
from app.core.detection.rules import RuleEngine
|
||||
from app.core.ingestion.pipeline import Pipeline, ScoredRequest
|
||||
|
||||
VALID_LINE = (
|
||||
"93.184.216.34 - - [11/Feb/2026:14:30:00 +0000] "
|
||||
'"GET /api/v1/users HTTP/1.1" 200 1234 '
|
||||
'"https://example.com" '
|
||||
'"Mozilla/5.0 (Windows NT 10.0; Win64; x64)"'
|
||||
)
|
||||
VALID_LINE = ("93.184.216.34 - - [11/Feb/2026:14:30:00 +0000] "
|
||||
'"GET /api/v1/users HTTP/1.1" 200 1234 '
|
||||
'"https://example.com" '
|
||||
'"Mozilla/5.0 (Windows NT 10.0; Win64; x64)"')
|
||||
|
||||
SQLI_LINE = (
|
||||
"93.184.216.34 - - [11/Feb/2026:14:30:01 +0000] "
|
||||
'"GET /users?id=1%27OR%201=1-- HTTP/1.1" 200 512 '
|
||||
'"-" "Mozilla/5.0"'
|
||||
)
|
||||
SQLI_LINE = ("93.184.216.34 - - [11/Feb/2026:14:30:01 +0000] "
|
||||
'"GET /users?id=1%27OR%201=1-- HTTP/1.1" 200 512 '
|
||||
'"-" "Mozilla/5.0"')
|
||||
|
||||
|
||||
async def _make_pipeline(
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
"""
|
||||
©AngelaMos | 2026
|
||||
test_scaler.py
|
||||
|
||||
Tests the FeatureScaler: fitting, transform correctness, JSON serialization, and round-trip loading.
|
||||
"""
|
||||
|
||||
import json
|
||||
|
|
@ -14,6 +16,9 @@ from ml.scaler import FeatureScaler
|
|||
|
||||
@pytest.fixture
|
||||
def sample_data() -> np.ndarray:
|
||||
"""
|
||||
Two-hundred samples of 35-dimensional random float32 data.
|
||||
"""
|
||||
rng = np.random.default_rng(42)
|
||||
return rng.standard_normal((200, 35)).astype(np.float32)
|
||||
|
||||
|
|
@ -21,23 +26,36 @@ def sample_data() -> np.ndarray:
|
|||
class TestFeatureScaler:
|
||||
|
||||
def test_fit_sets_n_features(self, sample_data: np.ndarray) -> None:
|
||||
"""
|
||||
Fitting stores the number of input features.
|
||||
"""
|
||||
scaler = FeatureScaler()
|
||||
scaler.fit(sample_data)
|
||||
assert scaler.n_features == 35
|
||||
|
||||
def test_transform_preserves_shape(self, sample_data: np.ndarray) -> None:
|
||||
"""
|
||||
Transform output has the same shape as the input array.
|
||||
"""
|
||||
scaler = FeatureScaler()
|
||||
scaler.fit(sample_data)
|
||||
transformed = scaler.transform(sample_data)
|
||||
assert transformed.shape == sample_data.shape
|
||||
|
||||
def test_transform_dtype_float32(self, sample_data: np.ndarray) -> None:
|
||||
"""
|
||||
Transformed array dtype remains float32.
|
||||
"""
|
||||
scaler = FeatureScaler()
|
||||
scaler.fit(sample_data)
|
||||
transformed = scaler.transform(sample_data)
|
||||
assert transformed.dtype == np.float32
|
||||
|
||||
def test_transformed_median_near_zero(self, sample_data: np.ndarray) -> None:
|
||||
def test_transformed_median_near_zero(self,
|
||||
sample_data: np.ndarray) -> None:
|
||||
"""
|
||||
Median of each feature column is approximately zero after scaling.
|
||||
"""
|
||||
scaler = FeatureScaler()
|
||||
scaler.fit(sample_data)
|
||||
transformed = scaler.transform(sample_data)
|
||||
|
|
@ -45,17 +63,21 @@ class TestFeatureScaler:
|
|||
assert np.allclose(medians, 0.0, atol=0.15)
|
||||
|
||||
def test_inverse_transform_recovers_original(
|
||||
self, sample_data: np.ndarray
|
||||
) -> None:
|
||||
self, sample_data: np.ndarray) -> None:
|
||||
"""
|
||||
Inverse transform recovers the original values within floating-point tolerance.
|
||||
"""
|
||||
scaler = FeatureScaler()
|
||||
scaler.fit(sample_data)
|
||||
transformed = scaler.transform(sample_data)
|
||||
recovered = scaler.inverse_transform(transformed)
|
||||
np.testing.assert_allclose(recovered, sample_data, atol=1e-5)
|
||||
|
||||
def test_save_json_creates_file(
|
||||
self, sample_data: np.ndarray, tmp_path: Path
|
||||
) -> None:
|
||||
def test_save_json_creates_file(self, sample_data: np.ndarray,
|
||||
tmp_path: Path) -> None:
|
||||
"""
|
||||
save_json writes a non-empty file to the given path.
|
||||
"""
|
||||
scaler = FeatureScaler()
|
||||
scaler.fit(sample_data)
|
||||
path = tmp_path / "scaler.json"
|
||||
|
|
@ -63,9 +85,11 @@ class TestFeatureScaler:
|
|||
assert path.exists()
|
||||
assert path.stat().st_size > 0
|
||||
|
||||
def test_save_json_is_valid_json(
|
||||
self, sample_data: np.ndarray, tmp_path: Path
|
||||
) -> None:
|
||||
def test_save_json_is_valid_json(self, sample_data: np.ndarray,
|
||||
tmp_path: Path) -> None:
|
||||
"""
|
||||
Saved JSON contains center, scale, and n_features keys.
|
||||
"""
|
||||
scaler = FeatureScaler()
|
||||
scaler.fit(sample_data)
|
||||
path = tmp_path / "scaler.json"
|
||||
|
|
@ -75,9 +99,11 @@ class TestFeatureScaler:
|
|||
assert "scale" in data
|
||||
assert "n_features" in data
|
||||
|
||||
def test_load_json_round_trip(
|
||||
self, sample_data: np.ndarray, tmp_path: Path
|
||||
) -> None:
|
||||
def test_load_json_round_trip(self, sample_data: np.ndarray,
|
||||
tmp_path: Path) -> None:
|
||||
"""
|
||||
Loading from JSON produces a scaler with identical transform output.
|
||||
"""
|
||||
scaler = FeatureScaler()
|
||||
scaler.fit(sample_data)
|
||||
path = tmp_path / "scaler.json"
|
||||
|
|
@ -91,11 +117,17 @@ class TestFeatureScaler:
|
|||
np.testing.assert_allclose(original_out, loaded_out, atol=1e-6)
|
||||
|
||||
def test_transform_before_fit_raises(self) -> None:
|
||||
"""
|
||||
Calling transform before fit raises RuntimeError.
|
||||
"""
|
||||
scaler = FeatureScaler()
|
||||
with pytest.raises(RuntimeError):
|
||||
scaler.transform(np.zeros((5, 35), dtype=np.float32))
|
||||
|
||||
def test_fit_transform_convenience(self, sample_data: np.ndarray) -> None:
|
||||
"""
|
||||
fit_transform fits and transforms in one call.
|
||||
"""
|
||||
scaler = FeatureScaler()
|
||||
result = scaler.fit_transform(sample_data)
|
||||
assert result.shape == sample_data.shape
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
"""
|
||||
©AngelaMos | 2026
|
||||
test_training.py
|
||||
|
||||
Tests training pipelines for the autoencoder, random forest, and isolation forest models.
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
|
|
@ -14,10 +16,18 @@ class TestAutoencoderTraining:
|
|||
|
||||
@pytest.fixture
|
||||
def normal_data(self) -> np.ndarray:
|
||||
"""
|
||||
Three-hundred samples of clipped 35-dimensional float32 data representing normal traffic.
|
||||
"""
|
||||
rng = np.random.default_rng(42)
|
||||
return (rng.standard_normal((300, 35)) * 0.3 + 0.5).astype(np.float32).clip(0, 1)
|
||||
return (rng.standard_normal(
|
||||
(300, 35)) * 0.3 + 0.5).astype(np.float32).clip(0, 1)
|
||||
|
||||
def test_returns_model_and_threshold(self, normal_data: np.ndarray) -> None:
|
||||
def test_returns_model_and_threshold(self,
|
||||
normal_data: np.ndarray) -> None:
|
||||
"""
|
||||
Training returns model, threshold, scaler, and history keys.
|
||||
"""
|
||||
result = train_autoencoder(normal_data, epochs=5, batch_size=32)
|
||||
assert "model" in result
|
||||
assert "threshold" in result
|
||||
|
|
@ -25,24 +35,38 @@ class TestAutoencoderTraining:
|
|||
assert "history" in result
|
||||
|
||||
def test_threshold_is_positive(self, normal_data: np.ndarray) -> None:
|
||||
"""
|
||||
Computed reconstruction error threshold is a positive value.
|
||||
"""
|
||||
result = train_autoencoder(normal_data, epochs=5, batch_size=32)
|
||||
assert result["threshold"] > 0.0
|
||||
|
||||
def test_history_has_train_loss(self, normal_data: np.ndarray) -> None:
|
||||
"""
|
||||
Training history includes one train_loss entry per epoch.
|
||||
"""
|
||||
result = train_autoencoder(normal_data, epochs=5, batch_size=32)
|
||||
assert "train_loss" in result["history"]
|
||||
assert len(result["history"]["train_loss"]) == 5
|
||||
|
||||
def test_custom_percentile(self, normal_data: np.ndarray) -> None:
|
||||
result_95 = train_autoencoder(
|
||||
normal_data, epochs=3, batch_size=32, percentile=95.0
|
||||
)
|
||||
result_99 = train_autoencoder(
|
||||
normal_data, epochs=3, batch_size=32, percentile=99.0
|
||||
)
|
||||
"""
|
||||
Higher percentile produces a higher or equal reconstruction threshold.
|
||||
"""
|
||||
result_95 = train_autoencoder(normal_data,
|
||||
epochs=3,
|
||||
batch_size=32,
|
||||
percentile=95.0)
|
||||
result_99 = train_autoencoder(normal_data,
|
||||
epochs=3,
|
||||
batch_size=32,
|
||||
percentile=99.0)
|
||||
assert result_99["threshold"] >= result_95["threshold"]
|
||||
|
||||
def test_model_is_in_eval_mode(self, normal_data: np.ndarray) -> None:
|
||||
"""
|
||||
Returned model is in eval mode after training completes.
|
||||
"""
|
||||
result = train_autoencoder(normal_data, epochs=3, batch_size=32)
|
||||
assert not result["model"].training
|
||||
|
||||
|
|
@ -51,37 +75,50 @@ class TestRandomForestTraining:
|
|||
|
||||
@pytest.fixture
|
||||
def labeled_data(self) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""
|
||||
Four-hundred samples with 300 benign and 100 attack labels.
|
||||
"""
|
||||
rng = np.random.default_rng(42)
|
||||
X = rng.standard_normal((400, 35)).astype(np.float32)
|
||||
y = np.concatenate([np.zeros(300, dtype=np.int64), np.ones(100, dtype=np.int64)])
|
||||
y = np.concatenate(
|
||||
[np.zeros(300, dtype=np.int64),
|
||||
np.ones(100, dtype=np.int64)])
|
||||
return X, y
|
||||
|
||||
def test_returns_model_and_metrics(
|
||||
self, labeled_data: tuple[np.ndarray, np.ndarray]
|
||||
) -> None:
|
||||
self, labeled_data: tuple[np.ndarray, np.ndarray]) -> None:
|
||||
"""
|
||||
Training returns model and metrics dict.
|
||||
"""
|
||||
X, y = labeled_data
|
||||
result = train_random_forest(X, y)
|
||||
assert "model" in result
|
||||
assert "metrics" in result
|
||||
|
||||
def test_model_has_predict_proba(
|
||||
self, labeled_data: tuple[np.ndarray, np.ndarray]
|
||||
) -> None:
|
||||
self, labeled_data: tuple[np.ndarray, np.ndarray]) -> None:
|
||||
"""
|
||||
Trained model exposes predict_proba for probability scoring.
|
||||
"""
|
||||
X, y = labeled_data
|
||||
result = train_random_forest(X, y)
|
||||
assert hasattr(result["model"], "predict_proba")
|
||||
|
||||
def test_metrics_contain_required_keys(
|
||||
self, labeled_data: tuple[np.ndarray, np.ndarray]
|
||||
) -> None:
|
||||
self, labeled_data: tuple[np.ndarray, np.ndarray]) -> None:
|
||||
"""
|
||||
Metrics dict contains f1, pr_auc, accuracy, precision, and recall.
|
||||
"""
|
||||
X, y = labeled_data
|
||||
result = train_random_forest(X, y)
|
||||
for key in ("f1", "pr_auc", "accuracy", "precision", "recall"):
|
||||
assert key in result["metrics"]
|
||||
|
||||
def test_probabilities_in_valid_range(
|
||||
self, labeled_data: tuple[np.ndarray, np.ndarray]
|
||||
) -> None:
|
||||
self, labeled_data: tuple[np.ndarray, np.ndarray]) -> None:
|
||||
"""
|
||||
Predicted probabilities are within the [0, 1] range.
|
||||
"""
|
||||
X, y = labeled_data
|
||||
result = train_random_forest(X, y)
|
||||
proba = result["model"].predict_proba(X[:10])
|
||||
|
|
@ -89,8 +126,10 @@ class TestRandomForestTraining:
|
|||
assert proba.max() <= 1.0
|
||||
|
||||
def test_metrics_values_in_valid_range(
|
||||
self, labeled_data: tuple[np.ndarray, np.ndarray]
|
||||
) -> None:
|
||||
self, labeled_data: tuple[np.ndarray, np.ndarray]) -> None:
|
||||
"""
|
||||
All metric values fall within [0, 1].
|
||||
"""
|
||||
X, y = labeled_data
|
||||
result = train_random_forest(X, y)
|
||||
for value in result["metrics"].values():
|
||||
|
|
@ -101,24 +140,39 @@ class TestIsolationForestTraining:
|
|||
|
||||
@pytest.fixture
|
||||
def normal_data(self) -> np.ndarray:
|
||||
"""
|
||||
Two-hundred samples of 35-dimensional standard normal float32 data.
|
||||
"""
|
||||
rng = np.random.default_rng(42)
|
||||
return rng.standard_normal((200, 35)).astype(np.float32)
|
||||
|
||||
def test_returns_model(self, normal_data: np.ndarray) -> None:
|
||||
"""
|
||||
Training returns a model key in the result dict.
|
||||
"""
|
||||
result = train_isolation_forest(normal_data)
|
||||
assert "model" in result
|
||||
|
||||
def test_model_has_score_samples(self, normal_data: np.ndarray) -> None:
|
||||
"""
|
||||
Trained model exposes score_samples for anomaly scoring.
|
||||
"""
|
||||
result = train_isolation_forest(normal_data)
|
||||
assert hasattr(result["model"], "score_samples")
|
||||
|
||||
def test_returns_metrics_with_n_samples(self, normal_data: np.ndarray) -> None:
|
||||
def test_returns_metrics_with_n_samples(self,
|
||||
normal_data: np.ndarray) -> None:
|
||||
"""
|
||||
Metrics include n_samples matching the training set size.
|
||||
"""
|
||||
result = train_isolation_forest(normal_data)
|
||||
assert result["metrics"]["n_samples"] == 200
|
||||
|
||||
def test_anomaly_scores_distinguish_normal_and_outlier(
|
||||
self, normal_data: np.ndarray
|
||||
) -> None:
|
||||
self, normal_data: np.ndarray) -> None:
|
||||
"""
|
||||
Normal training data scores higher than extreme outliers.
|
||||
"""
|
||||
result = train_isolation_forest(normal_data)
|
||||
model = result["model"]
|
||||
normal_scores = model.score_samples(normal_data[:50])
|
||||
|
|
|
|||
|
|
@ -136,6 +136,7 @@ ml = [
|
|||
{ name = "mlflow" },
|
||||
{ name = "numpy" },
|
||||
{ name = "onnxruntime" },
|
||||
{ name = "onnxscript" },
|
||||
{ name = "pandas" },
|
||||
{ name = "scikit-learn" },
|
||||
{ name = "skl2onnx" },
|
||||
|
|
@ -157,6 +158,7 @@ requires-dist = [
|
|||
{ name = "mypy", marker = "extra == 'dev'", specifier = ">=1.19.1" },
|
||||
{ name = "numpy", marker = "extra == 'ml'", specifier = ">=2.4.2" },
|
||||
{ name = "onnxruntime", marker = "extra == 'ml'", specifier = ">=1.24.1" },
|
||||
{ name = "onnxscript", marker = "extra == 'ml'", specifier = ">=0.6.2" },
|
||||
{ name = "orjson", specifier = ">=3.11.7" },
|
||||
{ name = "pandas", marker = "extra == 'ml'", specifier = ">=2.2.0,<3" },
|
||||
{ name = "pydantic", specifier = ">=2.12.5" },
|
||||
|
|
@ -1589,6 +1591,22 @@ wheels = [
|
|||
{ url = "https://files.pythonhosted.org/packages/01/a0/4fb0e6d36eaf079af366b2c1f68bafe92df6db963e2295da84388af64abc/onnx-1.20.1-cp312-abi3-win_arm64.whl", hash = "sha256:21d747348b1c8207406fa2f3e12b82f53e0d5bb3958bcd0288bd27d3cb6ebb00", size = 16344155, upload-time = "2026-01-10T01:39:45.536Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "onnx-ir"
|
||||
version = "0.2.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "ml-dtypes" },
|
||||
{ name = "numpy" },
|
||||
{ name = "onnx" },
|
||||
{ name = "sympy" },
|
||||
{ name = "typing-extensions" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/b2/a5/acc43c8fa6edbc584d127fb6bbd13ae9ebfc01b9675c74e0da2de15fa4a6/onnx_ir-0.2.0.tar.gz", hash = "sha256:8bad3906691987290789b26d05e0dbff467029a0b1e411e12e4cae02e43503e4", size = 141693, upload-time = "2026-02-24T02:31:10.998Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/4a/df/a99736bcca6b16e36c687ce4996abcf4ce73c514fddd9e730cfcb6a334f2/onnx_ir-0.2.0-py3-none-any.whl", hash = "sha256:eb14d1399c2442bd1ff702719e70074e9cedfa3af5729416a32752c9e0f82591", size = 164100, upload-time = "2026-02-24T02:31:09.454Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "onnxruntime"
|
||||
version = "1.24.1"
|
||||
|
|
@ -1609,6 +1627,23 @@ wheels = [
|
|||
{ url = "https://files.pythonhosted.org/packages/ab/a2/cfcf009eb38d90cc628c087b6506b3dfe1263387f3cbbf8d272af4fef957/onnxruntime-1.24.1-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:34488aa760fb5c2e6d06a7ca9241124eb914a6a06f70936a14c669d1b3df9598", size = 17099815, upload-time = "2026-02-05T17:31:43.092Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "onnxscript"
|
||||
version = "0.6.2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "ml-dtypes" },
|
||||
{ name = "numpy" },
|
||||
{ name = "onnx" },
|
||||
{ name = "onnx-ir" },
|
||||
{ name = "packaging" },
|
||||
{ name = "typing-extensions" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/e7/2b/538fdeb0e25bed5d7e0f954af5710543e2629499fb74381afc3333f8a8ae/onnxscript-0.6.2.tar.gz", hash = "sha256:abb2e6f464db40c9b8c7fbb3e64cca04cf3f4495e67c4eda5eac17b784191ce3", size = 590865, upload-time = "2026-02-10T22:53:39.638Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/66/56/e6b179397497ab93266b6eb00743403a6a699a29063a423c4a14595d3db9/onnxscript-0.6.2-py3-none-any.whl", hash = "sha256:20e3c3fd1da19b3655549d5455a2df719db47374fe430e01e865ae69127c37b9", size = 689064, upload-time = "2026-02-10T22:53:41.663Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "opentelemetry-api"
|
||||
version = "1.39.1"
|
||||
|
|
|
|||
|
|
@ -24,20 +24,19 @@ default:
|
|||
|
||||
[group('lint')]
|
||||
ruff *ARGS:
|
||||
cd backend && uv run ruff check app/ cli/ {{ARGS}}
|
||||
cd backend && uv run ruff check app/ cli/ ml/ tests/ {{ARGS}}
|
||||
|
||||
[group('lint')]
|
||||
ruff-fix:
|
||||
cd backend && uv run ruff check app/ cli/ --fix
|
||||
cd backend && uv run ruff format app/ cli/
|
||||
cd backend && uv run ruff check app/ cli/ ml/ tests/ --fix
|
||||
|
||||
[group('lint')]
|
||||
ruff-format:
|
||||
cd backend && uv run ruff format app/ cli/
|
||||
yapf *ARGS:
|
||||
cd backend && uv run yapf -rip app/ cli/ ml/ tests/ {{ARGS}}
|
||||
|
||||
[group('lint')]
|
||||
pylint *ARGS:
|
||||
cd backend && uv run pylint app {{ARGS}}
|
||||
cd backend && uv run pylint app cli ml {{ARGS}}
|
||||
|
||||
[group('lint')]
|
||||
lint: ruff pylint
|
||||
|
|
@ -48,7 +47,7 @@ lint: ruff pylint
|
|||
|
||||
[group('types')]
|
||||
mypy *ARGS:
|
||||
cd backend && uv run mypy app/ cli/ {{ARGS}}
|
||||
cd backend && uv run mypy app/ cli/ ml/ {{ARGS}}
|
||||
|
||||
[group('types')]
|
||||
typecheck: mypy
|
||||
|
|
@ -64,9 +63,13 @@ pytest *ARGS:
|
|||
[group('test')]
|
||||
test: pytest
|
||||
|
||||
[group('test')]
|
||||
test-v *ARGS:
|
||||
cd backend && uv run pytest tests/ -v {{ARGS}}
|
||||
|
||||
[group('test')]
|
||||
test-cov:
|
||||
cd backend && uv run pytest tests/ --cov=app --cov-report=term-missing --cov-report=html
|
||||
cd backend && uv run pytest tests/ --cov=app --cov=cli --cov=ml --cov-report=term-missing --cov-report=html
|
||||
|
||||
# =============================================================================
|
||||
# CI / Quality
|
||||
|
|
@ -222,6 +225,18 @@ vigil-config:
|
|||
vigil-health:
|
||||
curl -sf http://localhost:8000/health | python -m json.tool
|
||||
|
||||
[group('vigil')]
|
||||
vigil-train *ARGS:
|
||||
cd backend && uv run vigil train {{ARGS}}
|
||||
|
||||
[group('vigil')]
|
||||
vigil-replay *ARGS:
|
||||
cd backend && uv run vigil replay {{ARGS}}
|
||||
|
||||
[group('vigil')]
|
||||
vigil-status:
|
||||
curl -sf http://localhost:8000/models/status | python -m json.tool
|
||||
|
||||
# =============================================================================
|
||||
# Setup
|
||||
# =============================================================================
|
||||
|
|
@ -233,8 +248,8 @@ setup: setup-backend env
|
|||
[group('setup')]
|
||||
setup-backend:
|
||||
@echo "Setting up backend..."
|
||||
cd backend && uv sync
|
||||
@echo "Backend setup complete!"
|
||||
cd backend && uv sync --all-groups
|
||||
@echo "Backend setup complete (with dev + ml deps)!"
|
||||
|
||||
[group('setup')]
|
||||
env:
|
||||
|
|
|
|||
Loading…
Reference in New Issue