File size: 3,202 Bytes
c4f5819 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | """Ingest pipeline: raw rows -> classify -> persist -> (caller broadcasts).
Shared by CSV upload and the mock emitter so detection behaves identically
no matter where the data comes from.
"""
from __future__ import annotations
import csv
import io
import uuid
from collections import Counter
from datetime import datetime, timezone
from typing import Any
from .classifier import detect_rule
from .db import insert_events
# Fields we lift from a raw row onto the enriched event (besides detection).
PASSTHROUGH = ["source", "src_ip", "dst_ip", "user", "username", "dst_port", "path", "action"]
def parse_csv(data: bytes | str) -> list[dict[str, Any]]:
text = data.decode("utf-8", errors="replace") if isinstance(data, bytes) else data
reader = csv.DictReader(io.StringIO(text))
rows: list[dict[str, Any]] = []
for raw in reader:
# Normalise empty strings to None and strip keys.
row = {(k or "").strip(): (v.strip() if isinstance(v, str) else v) for k, v in raw.items()}
row = {k: (v if v != "" else None) for k, v in row.items()}
if row.get("message") or row.get("source"):
rows.append(row)
return rows
def enrich_row(raw: dict[str, Any], batch_id: str, ingested_at: str) -> dict[str, Any]:
"""Run detection and build a complete event dict ready to store/broadcast."""
rule = detect_rule(raw)
event: dict[str, Any] = {
"timestamp": raw.get("timestamp") or datetime.now(timezone.utc).isoformat(timespec="seconds"),
"category": rule["category"],
"label": rule["label"],
"severity": rule["severity"],
"mitre": rule["mitre"],
"risk_score": rule["risk_score"],
"reasons": rule["reasons"],
"message": raw.get("message", ""),
"expected_category": raw.get("category"), # ground-truth label if the CSV had one
"matches_label": rule["matches_label"],
"batch_id": batch_id,
"ingested_at": ingested_at,
"raw": raw,
}
for field in PASSTHROUGH:
if raw.get(field) is not None:
event[field] = raw.get(field)
return event
async def ingest_rows(rows: list[dict[str, Any]], batch_id: str | None = None) -> tuple[list[dict[str, Any]], dict[str, Any]]:
"""Classify + persist a list of raw rows. Returns (enriched_events, summary)."""
batch_id = batch_id or uuid.uuid4().hex[:12]
ingested_at = datetime.now(timezone.utc).isoformat(timespec="seconds")
events = [enrich_row(r, batch_id, ingested_at) for r in rows]
await insert_events(events)
by_cat = Counter(e["category"] for e in events)
by_sev = Counter(e["severity"] for e in events)
correct = sum(1 for e in events if e.get("matches_label") is True)
labelled = sum(1 for e in events if e.get("expected_category"))
summary = {
"batch_id": batch_id,
"ingested": len(events),
"by_category": dict(by_cat),
"by_severity": dict(by_sev),
"risky": sum(v for k, v in by_cat.items() if k != "benign"),
"critical": by_sev.get("critical", 0),
"detection_accuracy": round(correct / labelled * 100, 1) if labelled else None,
}
return events, summary
|