app/ingest.py

"""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