"""FastAPI application: realtime security-log RAG, CSV ingest, mock emitter. Run (self-host): uvicorn app.main:app --host 0.0.0.0 --port 8000 Realtime needs a single worker so all WebSocket clients share one process. """ from __future__ import annotations from contextlib import asynccontextmanager from pathlib import Path from typing import Any from fastapi import FastAPI, File, Query, UploadFile, WebSocket, WebSocketDisconnect from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import FileResponse, JSONResponse, Response from fastapi.staticfiles import StaticFiles from pydantic import BaseModel from . import __version__ from .config import get_settings from .db import fetch_alerts, fetch_events, fetch_stats, init_db from .generator import generate, generate_live, to_csv from .ingest import ingest_rows, parse_csv from .rag import ask as rag_ask from .realtime import manager, stream ROOT = Path(__file__).resolve().parents[1] WEB_DIR = ROOT / "web" HTML_DIR = ROOT / "html" settings = get_settings() @asynccontextmanager async def lifespan(app: FastAPI): await init_db() yield await stream.stop() app = FastAPI(title="ROGLAG Realtime", version=__version__, lifespan=lifespan) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], allow_credentials=False, ) # ---------------------------------------------------------------- broadcasting async def broadcast_events(events: list[dict[str, Any]], summary: dict[str, Any], kind: str) -> None: """Push events (capped) + a summary frame to all WebSocket clients.""" cap = 300 to_push = events if len(events) <= cap else ( [e for e in events if e.get("severity") == "critical"][:cap] or events[:cap] ) for ev in to_push: await manager.broadcast({"type": "event", "event": _slim(ev)}) await manager.broadcast({"type": "summary", "kind": kind, "summary": summary}) def _slim(ev: dict[str, Any]) -> dict[str, Any]: return { "ts": ev.get("timestamp") or ev.get("ts"), "source": ev.get("source"), "category": ev.get("category"), "label": ev.get("label"), "severity": ev.get("severity"), "mitre": ev.get("mitre"), "risk_score": ev.get("risk_score"), "src_ip": ev.get("src_ip"), "dst_ip": ev.get("dst_ip"), "user": ev.get("user") or ev.get("username"), "message": ev.get("message"), "reasons": ev.get("reasons"), "expected_category": ev.get("expected_category"), "matches_label": ev.get("matches_label"), } async def emit_mock(count: int, profile: str) -> dict[str, Any]: rows = generate_live(count, profile=profile) events, summary = await ingest_rows(rows) await broadcast_events(events, summary, kind="mock") return summary # ---------------------------------------------------------------- API: health/stats @app.get("/health") async def health() -> dict[str, Any]: return { "ok": True, "version": __version__, "storage": "postgres" if settings.is_postgres else "sqlite", "llm_ready": settings.llm_ready, "ws_clients": manager.count, "stream_running": stream.is_running(), } @app.get("/api/stats") async def api_stats() -> Any: return await fetch_stats() @app.get("/api/logs") async def api_logs(limit: int = 100, offset: int = 0, category: str = "", severity: str = "", source: str = "", q: str = "") -> Any: return await fetch_events(limit, offset, category, severity, source, q) @app.get("/api/alerts") async def api_alerts(limit: int = 20) -> Any: return await fetch_alerts(limit) class AskBody(BaseModel): question: str limit: int = 8 use_llm: bool = False model: str = "" @app.post("/api/ask") async def api_ask(body: AskBody) -> Any: if not body.question.strip(): return JSONResponse({"error": "question is required"}, status_code=400) return await rag_ask(body.question.strip(), min(body.limit, 20), body.use_llm, body.model) # ---------------------------------------------------------------- CSV upload (realtime detect) @app.post("/api/ingest/csv") async def api_ingest_csv(file: UploadFile = File(...)) -> Any: data = await file.read() rows = parse_csv(data) if not rows: return JSONResponse({"error": "no parseable rows (need a 'message' or 'source' column)"}, status_code=400) if len(rows) > settings.max_csv_rows: rows = rows[: settings.max_csv_rows] events, summary = await ingest_rows(rows) summary["filename"] = file.filename await broadcast_events(events, summary, kind="csv") return summary # ---------------------------------------------------------------- CSV generator (unique each call) @app.get("/api/generate/csv") async def api_generate_csv(count: int = Query(200, ge=1, le=20000), profile: str = Query("mixed")) -> Response: records = generate(count, profile=profile, seed=None) # unseeded -> unique each time csv_text = to_csv(records) fname = f"roglag_logs_{profile}_{len(records)}.csv" return Response( content=csv_text, media_type="text/csv", headers={"Content-Disposition": f'attachment; filename="{fname}"'}, ) # ---------------------------------------------------------------- mock emitter (realtime) @app.post("/api/mock/emit") async def api_mock_emit(count: int = Query(20, ge=1), profile: str = Query("mixed")) -> Any: count = min(count, settings.max_mock_burst) return await emit_mock(count, profile) @app.post("/api/mock/stream/start") async def api_stream_start(rate: int = Query(3, ge=1, le=50), interval: float = Query(1.0, ge=0.2, le=10.0), profile: str = Query("mixed")) -> Any: await stream.start(emit_mock, rate=min(rate, settings.max_mock_burst), interval=interval, profile=profile) return {"running": stream.is_running(), "rate": rate, "interval": interval, "profile": profile} @app.post("/api/mock/stream/stop") async def api_stream_stop() -> Any: await stream.stop() return {"running": stream.is_running()} # ---------------------------------------------------------------- WebSocket @app.websocket("/ws") async def ws_feed(ws: WebSocket) -> None: await manager.connect(ws) try: await ws.send_json({"type": "hello", "clients": manager.count}) while True: await ws.receive_text() # keep-alive; we ignore client messages except WebSocketDisconnect: await manager.disconnect(ws) except Exception: await manager.disconnect(ws) # ---------------------------------------------------------------- static / pages @app.get("/") async def index() -> Response: live = WEB_DIR / "live.html" return FileResponse(live) if live.exists() else JSONResponse({"ok": True, "see": "/docs-html"}) @app.get("/docs-html") async def docs_html() -> Response: return FileResponse(HTML_DIR / "docs.html") @app.get("/plan") async def plan_html() -> Response: return FileResponse(HTML_DIR / "plan.html") if WEB_DIR.exists(): app.mount("/web", StaticFiles(directory=str(WEB_DIR)), name="web") if HTML_DIR.exists(): app.mount("/html", StaticFiles(directory=str(HTML_DIR)), name="html")