feat: initial EDR system — agent + backend + frontend + ensemble ML
This commit is contained in:
kilyabin
@@ -0,0 +1,4 @@
|
||||
__pycache__
|
||||
*.py[cod]
|
||||
.venv
|
||||
*.log
|
||||
@@ -0,0 +1,13 @@
|
||||
FROM python:3.11-slim
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# Зависимости отдельным слоем для кеширования
|
||||
COPY requirements.txt .
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
|
||||
COPY . .
|
||||
|
||||
EXPOSE 8000
|
||||
|
||||
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]
|
||||
+349
@@ -0,0 +1,349 @@
|
||||
"""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)
|
||||
@@ -0,0 +1,4 @@
|
||||
fastapi>=0.111.0
|
||||
uvicorn[standard]>=0.29.0
|
||||
aiosqlite>=0.20.0
|
||||
python-multipart>=0.0.9
|
||||
Reference in New Issue
Block a user