kilyabin
350 lines
13 KiB
Python
350 lines
13 KiB
Python
"""EDR Backend — FastAPI REST API + WebSocket broadcast.
|
|
|
|
Receives metric batches from the agent, stores them in SQLite,
|
|
and pushes real-time anomaly alerts to connected dashboard clients.
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
from contextlib import asynccontextmanager
|
|
from datetime import date, datetime, timezone
|
|
from typing import Any
|
|
|
|
import aiosqlite
|
|
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from pydantic import BaseModel
|
|
|
|
logger = logging.getLogger(__name__)
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="%(asctime)s [%(levelname)s] %(message)s",
|
|
)
|
|
|
|
DB_PATH = os.getenv("DB_PATH", "../data/edr.db")
|
|
|
|
|
|
# ── Database ──────────────────────────────────────────────────────────────────
|
|
|
|
|
|
async def _init_db() -> None:
|
|
"""Create all tables on first startup."""
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
await db.executescript("""
|
|
CREATE TABLE IF NOT EXISTS metrics (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
timestamp TEXT NOT NULL,
|
|
pid INTEGER,
|
|
name TEXT,
|
|
cpu_percent REAL,
|
|
memory_mb REAL,
|
|
open_files INTEGER,
|
|
connections INTEGER,
|
|
is_anomaly INTEGER DEFAULT 0
|
|
);
|
|
CREATE TABLE IF NOT EXISTS alerts (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
timestamp TEXT NOT NULL,
|
|
pid INTEGER,
|
|
process_name TEXT,
|
|
reason TEXT,
|
|
severity TEXT
|
|
);
|
|
CREATE TABLE IF NOT EXISTS file_events (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
timestamp TEXT NOT NULL,
|
|
path TEXT,
|
|
event_type TEXT
|
|
);
|
|
CREATE TABLE IF NOT EXISTS system_metrics (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
timestamp TEXT NOT NULL,
|
|
cpu_percent REAL,
|
|
ram_percent REAL,
|
|
net_connections INTEGER
|
|
);
|
|
""")
|
|
await db.commit()
|
|
|
|
|
|
# ── WebSocket manager ─────────────────────────────────────────────────────────
|
|
|
|
|
|
class _WsManager:
|
|
"""Thread-safe registry of connected WebSocket clients."""
|
|
|
|
def __init__(self) -> None:
|
|
self._clients: list[WebSocket] = []
|
|
|
|
async def connect(self, ws: WebSocket) -> None:
|
|
await ws.accept()
|
|
self._clients.append(ws)
|
|
logger.info("WS connected total=%d", len(self._clients))
|
|
|
|
def disconnect(self, ws: WebSocket) -> None:
|
|
self._clients.remove(ws)
|
|
logger.info("WS disconnected total=%d", len(self._clients))
|
|
|
|
async def broadcast(self, data: dict) -> None:
|
|
"""Send JSON payload to every connected client; drop stale ones."""
|
|
stale: list[WebSocket] = []
|
|
for ws in list(self._clients):
|
|
try:
|
|
await ws.send_json(data)
|
|
except Exception:
|
|
stale.append(ws)
|
|
for ws in stale:
|
|
if ws in self._clients:
|
|
self._clients.remove(ws)
|
|
|
|
|
|
_ws = _WsManager()
|
|
|
|
|
|
# ── App lifecycle ─────────────────────────────────────────────────────────────
|
|
|
|
|
|
@asynccontextmanager
|
|
async def _lifespan(app: FastAPI):
|
|
await _init_db()
|
|
logger.info("Database ready at %s", DB_PATH)
|
|
yield
|
|
|
|
|
|
app = FastAPI(title="EDR Backend", version="1.0.0", lifespan=_lifespan)
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
|
|
# ── Pydantic models ───────────────────────────────────────────────────────────
|
|
|
|
|
|
class ProcessMetric(BaseModel):
|
|
pid: int
|
|
name: str
|
|
cpu_percent: float
|
|
memory_mb: float
|
|
open_files: int
|
|
connections: int
|
|
is_anomaly: bool = False
|
|
|
|
|
|
class SystemMetric(BaseModel):
|
|
cpu_percent: float
|
|
ram_percent: float
|
|
network_connections: int
|
|
|
|
|
|
class FileEvent(BaseModel):
|
|
path: str
|
|
event_type: str
|
|
timestamp: str
|
|
|
|
|
|
class MetricsBatch(BaseModel):
|
|
processes: list[ProcessMetric]
|
|
system: SystemMetric
|
|
file_events: list[FileEvent] = []
|
|
timestamp: str
|
|
|
|
|
|
# ── Helpers ───────────────────────────────────────────────────────────────────
|
|
|
|
|
|
def _severity(proc: ProcessMetric) -> str:
|
|
"""Derive alert severity from process resource usage."""
|
|
if proc.cpu_percent > 80 or proc.memory_mb > 2000:
|
|
return "high"
|
|
if proc.cpu_percent > 50 or proc.memory_mb > 1000:
|
|
return "medium"
|
|
return "low"
|
|
|
|
|
|
def _reason(proc: ProcessMetric) -> str:
|
|
"""Build a human-readable anomaly reason string (Russian)."""
|
|
parts: list[str] = []
|
|
if proc.cpu_percent > 80:
|
|
parts.append(f"CPU {proc.cpu_percent:.1f}%")
|
|
if proc.memory_mb > 2000:
|
|
parts.append(f"ОЗУ {proc.memory_mb:.0f} МБ")
|
|
if proc.connections > 50:
|
|
parts.append(f"{proc.connections} TCP-соединений")
|
|
if proc.open_files > 200:
|
|
parts.append(f"{proc.open_files} открытых файлов")
|
|
detail = ", ".join(parts) if parts else "статистический выброс (ансамбль ML: IF + LOF + OC-SVM)"
|
|
return f"Аномальное поведение: {detail}"
|
|
|
|
|
|
# ── REST endpoints ────────────────────────────────────────────────────────────
|
|
|
|
|
|
@app.get("/api/metrics", summary="Recent process metrics")
|
|
async def get_metrics(limit: int = 200) -> list[dict]:
|
|
"""Return the most recent *limit* rows from the metrics table."""
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
db.row_factory = aiosqlite.Row
|
|
async with db.execute(
|
|
"SELECT * FROM metrics ORDER BY id DESC LIMIT ?", (limit,)
|
|
) as cur:
|
|
return [dict(r) for r in await cur.fetchall()]
|
|
|
|
|
|
@app.get("/api/alerts", summary="Recent alerts")
|
|
async def get_alerts(limit: int = 50) -> list[dict]:
|
|
"""Return the most recent *limit* alert rows."""
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
db.row_factory = aiosqlite.Row
|
|
async with db.execute(
|
|
"SELECT * FROM alerts ORDER BY id DESC LIMIT ?", (limit,)
|
|
) as cur:
|
|
return [dict(r) for r in await cur.fetchall()]
|
|
|
|
|
|
@app.get("/api/system-metrics", summary="System-level CPU/RAM timeline")
|
|
async def get_system_metrics(limit: int = 60) -> list[dict]:
|
|
"""Return last *limit* system snapshots in chronological order (for charts)."""
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
db.row_factory = aiosqlite.Row
|
|
async with db.execute(
|
|
"SELECT * FROM system_metrics ORDER BY id DESC LIMIT ?", (limit,)
|
|
) as cur:
|
|
rows = await cur.fetchall()
|
|
return list(reversed([dict(r) for r in rows]))
|
|
|
|
|
|
@app.get("/api/file-events", summary="Recent filesystem events")
|
|
async def get_file_events(limit: int = 100) -> list[dict]:
|
|
"""Return the most recent *limit* filesystem events."""
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
db.row_factory = aiosqlite.Row
|
|
async with db.execute(
|
|
"SELECT * FROM file_events ORDER BY id DESC LIMIT ?", (limit,)
|
|
) as cur:
|
|
return [dict(r) for r in await cur.fetchall()]
|
|
|
|
|
|
@app.get("/api/stats", summary="Dashboard statistics")
|
|
async def get_stats() -> dict[str, Any]:
|
|
"""Return aggregated statistics for the Overview panel."""
|
|
today = date.today().isoformat()
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
db.row_factory = aiosqlite.Row
|
|
|
|
async with db.execute(
|
|
"SELECT COUNT(DISTINCT pid) FROM metrics"
|
|
) as cur:
|
|
total_processes = (await cur.fetchone())[0]
|
|
|
|
async with db.execute(
|
|
"SELECT COUNT(*) FROM metrics WHERE is_anomaly=1 AND timestamp LIKE ?",
|
|
(f"{today}%",),
|
|
) as cur:
|
|
anomalies_today = (await cur.fetchone())[0]
|
|
|
|
async with db.execute(
|
|
"SELECT COUNT(*) FROM alerts WHERE timestamp LIKE ?", (f"{today}%",)
|
|
) as cur:
|
|
alerts_today = (await cur.fetchone())[0]
|
|
|
|
# Средняя нагрузка системы за сегодня (из таблицы system_metrics)
|
|
async with db.execute(
|
|
"SELECT AVG(cpu_percent), AVG(ram_percent) FROM system_metrics WHERE timestamp LIKE ?",
|
|
(f"{today}%",),
|
|
) as cur:
|
|
row = await cur.fetchone()
|
|
cpu_avg = round(float(row[0] or 0), 1)
|
|
ram_avg = round(float(row[1] or 0), 1)
|
|
|
|
return {
|
|
"total_processes": total_processes,
|
|
"anomalies_today": anomalies_today,
|
|
"alerts_today": alerts_today,
|
|
"cpu_avg": cpu_avg,
|
|
"ram_avg": ram_avg,
|
|
}
|
|
|
|
|
|
@app.post("/api/metrics", status_code=201, summary="Ingest metric batch from agent")
|
|
async def post_metrics(batch: MetricsBatch) -> dict[str, int]:
|
|
"""Store process metrics, create alerts for anomalies, broadcast via WS."""
|
|
ts = batch.timestamp
|
|
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
# Bulk-insert process metrics
|
|
await db.executemany(
|
|
"""INSERT INTO metrics
|
|
(timestamp, pid, name, cpu_percent, memory_mb, open_files, connections, is_anomaly)
|
|
VALUES (?,?,?,?,?,?,?,?)""",
|
|
[
|
|
(ts, p.pid, p.name, p.cpu_percent, p.memory_mb,
|
|
p.open_files, p.connections, int(p.is_anomaly))
|
|
for p in batch.processes
|
|
],
|
|
)
|
|
|
|
# Системные метрики — одна строка на батч (для графика)
|
|
await db.execute(
|
|
"INSERT INTO system_metrics (timestamp, cpu_percent, ram_percent, net_connections) VALUES (?,?,?,?)",
|
|
(ts, batch.system.cpu_percent, batch.system.ram_percent, batch.system.network_connections),
|
|
)
|
|
|
|
# Bulk-insert filesystem events
|
|
if batch.file_events:
|
|
await db.executemany(
|
|
"INSERT INTO file_events (timestamp, path, event_type) VALUES (?,?,?)",
|
|
[(e.timestamp, e.path, e.event_type) for e in batch.file_events],
|
|
)
|
|
|
|
# Пороговый алерт работает независимо от ML — для любых процессов
|
|
threshold_anomalies = [
|
|
p for p in batch.processes
|
|
if not p.is_anomaly and (p.cpu_percent > 80 or p.memory_mb > 3000 or p.connections > 100)
|
|
]
|
|
for proc in threshold_anomalies:
|
|
proc.is_anomaly = True # пометить для записи в метрики
|
|
|
|
# Все аномалии: ML + пороговые
|
|
anomalies = [p for p in batch.processes if p.is_anomaly]
|
|
for proc in anomalies:
|
|
sev = _severity(proc)
|
|
reason = _reason(proc)
|
|
await db.execute(
|
|
"INSERT INTO alerts (timestamp, pid, process_name, reason, severity) VALUES (?,?,?,?,?)",
|
|
(ts, proc.pid, proc.name, reason, sev),
|
|
)
|
|
await _ws.broadcast(
|
|
{
|
|
"type": "alert",
|
|
"pid": proc.pid,
|
|
"process_name": proc.name,
|
|
"reason": reason,
|
|
"severity": sev,
|
|
"timestamp": ts,
|
|
}
|
|
)
|
|
|
|
await db.commit()
|
|
|
|
logger.info("Batch stored: procs=%d anomalies=%d fs=%d",
|
|
len(batch.processes), len(anomalies), len(batch.file_events))
|
|
return {"received": len(batch.processes), "anomalies": len(anomalies)}
|
|
|
|
|
|
# ── WebSocket ─────────────────────────────────────────────────────────────────
|
|
|
|
|
|
@app.websocket("/ws")
|
|
async def ws_endpoint(ws: WebSocket) -> None:
|
|
"""WebSocket endpoint — each dashboard connects here for live alerts."""
|
|
await _ws.connect(ws)
|
|
try:
|
|
while True:
|
|
await ws.receive_text() # keep connection alive; messages ignored
|
|
except WebSocketDisconnect:
|
|
_ws.disconnect(ws)
|