add baseline and sqlite

src/fgai/baseline.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+import hashlib
+import sqlite3
+import time
+from collections import defaultdict
+from pathlib import Path
+from statistics import mean, pstdev
+
+from .logs import THREAT_ACTIONS, is_utm_event
+from .models import LogEvent
+
+
+def _number(value: str | None) -> int:
+    try:
+        return int(float(value or 0))
+    except ValueError:
+        return 0
+
+
+class BaselineStore:
+    """Persistent five-minute behavior baseline, implemented with stdlib SQLite."""
+
+    def __init__(self, path: str, *, bucket_seconds: int = 300) -> None:
+        self.path = Path(path)
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        self.bucket_seconds = bucket_seconds
+        with self._connect() as connection:
+            connection.executescript(
+                """
+                create table if not exists seen_events (fingerprint text primary key);
+                create table if not exists source_buckets (
+                  source_ip text not null, bucket_start integer not null,
+                  events integer not null, bytes integer not null, hitcount integer not null,
+                  denies integer not null, utm integer not null,
+                  primary key (source_ip, bucket_start)
+                );
+                """
+            )
+
+    def _connect(self) -> sqlite3.Connection:
+        return sqlite3.connect(self.path)
+
+    def ingest(self, events: list[LogEvent], *, observed_at: int | None = None) -> int:
+        observed_at = observed_at or int(time.time())
+        pending: dict[tuple[str, int], list[int]] = defaultdict(lambda: [0, 0, 0, 0, 0])
+        inserted = 0
+        with self._connect() as connection:
+            for event in events:
+                if not event.src_ip:
+                    continue
+                fingerprint = hashlib.sha256(event.raw.encode("utf-8", errors="replace")).hexdigest()
+                if connection.execute("insert or ignore into seen_events values (?)", (fingerprint,)).rowcount != 1:
+                    continue
+                bucket = observed_at - (observed_at % self.bucket_seconds)
+                values = pending[(event.src_ip, bucket)]
+                values[0] += 1
+                values[1] += _number(event.fields.get("sentbyte")) + _number(event.fields.get("rcvdbyte"))
+                values[2] += _number(event.fields.get("hitcount"))
+                values[3] += int(event.action in THREAT_ACTIONS)
+                values[4] += int(is_utm_event(event))
+                inserted += 1
+            for (source_ip, bucket), values in pending.items():
+                connection.execute(
+                    """insert into source_buckets values (?, ?, ?, ?, ?, ?, ?)
+                    on conflict(source_ip, bucket_start) do update set
+                      events=events+excluded.events, bytes=bytes+excluded.bytes,
+                      hitcount=hitcount+excluded.hitcount, denies=denies+excluded.denies, utm=utm+excluded.utm""",
+                    (source_ip, bucket, *values),
+                )
+        return inserted
+
+    def profiles(self, source_ips: set[str]) -> dict[str, dict[str, float | int]]:
+        profiles: dict[str, dict[str, float | int]] = {}
+        with self._connect() as connection:
+            for source_ip in source_ips:
+                rows = connection.execute(
+                    "select events, bytes, hitcount from source_buckets where source_ip=? order by bucket_start", (source_ip,)
+                ).fetchall()
+                # Do not compare the current (often incomplete) bucket to itself.
+                rows = rows[:-1]
+                if len(rows) < 12:
+                    continue
+                rates = [row[0] * 60 / self.bucket_seconds for row in rows]
+                hit_rates = [row[2] * 60 / self.bucket_seconds for row in rows]
+                profiles[source_ip] = {
+                    "samples": len(rows),
+                    "event_rate_mean": mean(rates), "event_rate_stddev": pstdev(rates) or 1.0,
+                    "hitcount_rate_mean": mean(hit_rates), "hitcount_rate_stddev": pstdev(hit_rates) or 1.0,
+                }
+        return profiles