"""Async database layer (SQLAlchemy Core). Portable across SQLite (default, zero-config self-host) and PostgreSQL (docker-compose / production) using the same code. Stores classified log events; alerts and stats are computed on read. """ from __future__ import annotations from collections import Counter, defaultdict from datetime import datetime, timezone from typing import Any from sqlalchemy import ( JSON, BigInteger, Boolean, Column, Integer, MetaData, String, Table, Text, func, select, ) from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine from .classifier import CATEGORY_META from .config import get_settings metadata = MetaData() # BigInteger on Postgres, but Integer on SQLite so it aliases rowid and # autoincrements (SQLite only autoincrements INTEGER PRIMARY KEY). PK_TYPE = BigInteger().with_variant(Integer, "sqlite") log_events = Table( "log_events", metadata, Column("id", PK_TYPE, primary_key=True, autoincrement=True), Column("ts", String(40), index=True), # event time (ISO8601) Column("source", String(64), index=True), Column("category", String(32), index=True), # detected category Column("severity", String(16), index=True), Column("mitre", String(32)), Column("risk_score", Integer), Column("src_ip", String(64), index=True), Column("dst_ip", String(64)), Column("username", String(64)), Column("message", Text), Column("reasons", Text), Column("expected_category", String(32)), # original label if CSV provided one Column("matches_label", Boolean), Column("raw", JSON), Column("batch_id", String(64), index=True), Column("ingested_at", String(40), index=True), ) _engine: AsyncEngine | None = None def get_engine() -> AsyncEngine: global _engine if _engine is None: settings = get_settings() _engine = create_async_engine(settings.database_url, future=True, pool_pre_ping=True) return _engine async def init_db() -> None: engine = get_engine() async with engine.begin() as conn: await conn.run_sync(metadata.create_all) async def insert_events(events: list[dict[str, Any]]) -> int: """Insert enriched event rows. Returns count inserted.""" if not events: return 0 rows = [_to_row(e) for e in events] engine = get_engine() async with engine.begin() as conn: await conn.execute(log_events.insert(), rows) return len(rows) def _to_row(event: dict[str, Any]) -> dict[str, Any]: return { "ts": str(event.get("timestamp") or event.get("ts") or ""), "source": event.get("source"), "category": event.get("category"), "severity": event.get("severity"), "mitre": event.get("mitre"), "risk_score": event.get("risk_score"), "src_ip": event.get("src_ip"), "dst_ip": event.get("dst_ip"), "username": event.get("user") or event.get("username"), "message": event.get("message"), "reasons": "; ".join(event.get("reasons", [])) if isinstance(event.get("reasons"), list) else event.get("reasons"), "expected_category": event.get("expected_category"), "matches_label": event.get("matches_label"), "raw": event.get("raw") or event, "batch_id": event.get("batch_id"), "ingested_at": event.get("ingested_at") or datetime.now(timezone.utc).isoformat(timespec="seconds"), } def _row_to_dict(row: Any) -> dict[str, Any]: d = dict(row._mapping) if isinstance(d.get("reasons"), str) and d["reasons"]: d["reasons"] = d["reasons"].split("; ") return d async def fetch_events(limit: int = 100, offset: int = 0, category: str = "", severity: str = "", source: str = "", q: str = "") -> dict[str, Any]: engine = get_engine() stmt = select(log_events) if category: stmt = stmt.where(log_events.c.category == category) if severity: stmt = stmt.where(log_events.c.severity == severity) if source: stmt = stmt.where(log_events.c.source == source) if q: stmt = stmt.where(log_events.c.message.ilike(f"%{q}%")) count_stmt = select(func.count()).select_from(stmt.subquery()) page_stmt = stmt.order_by(log_events.c.id.desc()).limit(min(limit, 500)).offset(max(offset, 0)) async with engine.connect() as conn: total = (await conn.execute(count_stmt)).scalar() or 0 rows = (await conn.execute(page_stmt)).fetchall() return {"total": total, "limit": limit, "offset": offset, "rows": [_row_to_dict(r) for r in rows]} async def fetch_stats() -> dict[str, Any]: engine = get_engine() async with engine.connect() as conn: total = (await conn.execute(select(func.count()).select_from(log_events))).scalar() or 0 cat_rows = (await conn.execute( select(log_events.c.category, func.count()).group_by(log_events.c.category))).fetchall() sev_rows = (await conn.execute( select(log_events.c.severity, func.count()).group_by(log_events.c.severity))).fetchall() src_rows = (await conn.execute( select(log_events.c.source, func.count()).group_by(log_events.c.source))).fetchall() ip_rows = (await conn.execute( select(log_events.c.src_ip, func.count()).where(log_events.c.src_ip.isnot(None)) .group_by(log_events.c.src_ip).order_by(func.count().desc()).limit(10))).fetchall() by_category = {k or "unknown": v for k, v in cat_rows} risky = sum(v for k, v in by_category.items() if k != "benign") return { "total": total, "risky": risky, "risk_rate": round((risky / total) * 100, 2) if total else 0.0, "by_category": by_category, "by_severity": {k or "unknown": v for k, v in sev_rows}, "by_source": {k or "unknown": v for k, v in src_rows}, "top_ips": [[k, v] for k, v in ip_rows], "categories": CATEGORY_META, } async def fetch_alerts(limit: int = 20) -> dict[str, Any]: """Group critical events into alerts (computed on read).""" engine = get_engine() stmt = select(log_events).where(log_events.c.severity == "critical").order_by(log_events.c.id.desc()).limit(2000) async with engine.connect() as conn: rows = [_row_to_dict(r) for r in (await conn.execute(stmt)).fetchall()] groups: dict[str, list[dict[str, Any]]] = defaultdict(list) for row in rows: if row.get("category") == "priv_esc": key = "|".join([str(row.get("category")), str(row.get("username")), str(row.get("src_ip"))]) else: key = "|".join([str(row.get("category")), str(row.get("src_ip")), str(row.get("dst_ip")), str(row.get("username"))]) groups[key].append(row) alerts = [] for key, items in groups.items(): items.sort(key=lambda r: str(r.get("ts", ""))) first, last = items[0], items[-1] meta = CATEGORY_META.get(first.get("category"), CATEGORY_META["benign"]) alerts.append({ "title": meta["label"], "category": first.get("category"), "severity": "critical", "mitre": meta["mitre"], "count": len(items), "first_seen": first.get("ts"), "last_seen": last.get("ts"), "src_ip": first.get("src_ip") or "-", "dst_ip": first.get("dst_ip") or "-", "user": first.get("username") or "-", "message": last.get("message"), "reasons": last.get("reasons"), "recommended_action": meta["runbook"], }) alerts.sort(key=lambda a: str(a["last_seen"]), reverse=True) return {"total": len(alerts), "limit": limit, "alerts": alerts[:limit]} async def search_rows(q: str, limit: int = 400) -> list[dict[str, Any]]: """Fetch candidate rows for keyword RAG (message LIKE any token).""" engine = get_engine() stmt = select(log_events).order_by(log_events.c.id.desc()).limit(limit) if q: stmt = select(log_events).where(log_events.c.message.ilike(f"%{q}%")).order_by(log_events.c.id.desc()).limit(limit) async with engine.connect() as conn: rows = (await conn.execute(stmt)).fetchall() return [_row_to_dict(r) for r in rows]