152 lines
4.3 KiB
Python
152 lines
4.3 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
sysmonstm hub - Local aggregator that receives from collectors and forwards to edge.
|
|
|
|
Runs on the local network, receives metrics from collectors via WebSocket,
|
|
and forwards them to the cloud edge.
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import os
|
|
|
|
from fastapi import FastAPI, Query, WebSocket, WebSocketDisconnect
|
|
|
|
# Configuration
|
|
API_KEY = os.environ.get("API_KEY", "")
|
|
EDGE_URL = os.environ.get("EDGE_URL", "") # e.g., wss://sysmonstm.mcrn.ar/ws
|
|
EDGE_API_KEY = os.environ.get("EDGE_API_KEY", "")
|
|
LOG_LEVEL = os.environ.get("LOG_LEVEL", "INFO")
|
|
|
|
# Logging setup
|
|
logging.basicConfig(
|
|
level=getattr(logging, LOG_LEVEL.upper(), logging.INFO),
|
|
format="%(asctime)s [%(levelname)s] %(message)s",
|
|
datefmt="%Y-%m-%d %H:%M:%S",
|
|
)
|
|
log = logging.getLogger("hub")
|
|
|
|
app = FastAPI(title="sysmonstm-hub")
|
|
|
|
# State
|
|
collector_connections: list[WebSocket] = []
|
|
machines: dict = {}
|
|
edge_ws = None
|
|
|
|
|
|
async def connect_to_edge():
|
|
"""Maintain persistent connection to edge and forward metrics."""
|
|
global edge_ws
|
|
|
|
if not EDGE_URL:
|
|
log.info("No EDGE_URL configured, running in local-only mode")
|
|
return
|
|
|
|
import websockets
|
|
|
|
url = EDGE_URL
|
|
if EDGE_API_KEY:
|
|
separator = "&" if "?" in url else "?"
|
|
url = f"{url}{separator}key={EDGE_API_KEY}"
|
|
|
|
while True:
|
|
try:
|
|
log.info(f"Connecting to edge: {EDGE_URL}")
|
|
async with websockets.connect(url) as ws:
|
|
edge_ws = ws
|
|
log.info("Connected to edge")
|
|
|
|
while True:
|
|
try:
|
|
msg = await asyncio.wait_for(ws.recv(), timeout=30)
|
|
# Ignore messages from edge (pings, etc)
|
|
except asyncio.TimeoutError:
|
|
await ws.ping()
|
|
|
|
except asyncio.CancelledError:
|
|
break
|
|
except Exception as e:
|
|
edge_ws = None
|
|
log.warning(f"Edge connection error: {e}. Reconnecting in 5s...")
|
|
await asyncio.sleep(5)
|
|
|
|
|
|
async def forward_to_edge(data: dict):
|
|
"""Forward metrics to edge if connected."""
|
|
global edge_ws
|
|
if edge_ws:
|
|
try:
|
|
await edge_ws.send(json.dumps(data))
|
|
log.debug(f"Forwarded to edge: {data.get('machine_id')}")
|
|
except Exception as e:
|
|
log.warning(f"Failed to forward to edge: {e}")
|
|
|
|
|
|
@app.on_event("startup")
|
|
async def startup():
|
|
asyncio.create_task(connect_to_edge())
|
|
|
|
|
|
@app.get("/health")
|
|
async def health():
|
|
return {
|
|
"status": "ok",
|
|
"machines": len(machines),
|
|
"collectors": len(collector_connections),
|
|
"edge_connected": edge_ws is not None,
|
|
}
|
|
|
|
|
|
@app.get("/api/machines")
|
|
async def get_machines():
|
|
return machines
|
|
|
|
|
|
@app.websocket("/ws")
|
|
async def websocket_endpoint(websocket: WebSocket, key: str = Query(default="")):
|
|
# Validate API key
|
|
if API_KEY and key != API_KEY:
|
|
log.warning(f"Invalid API key from {websocket.client}")
|
|
await websocket.close(code=4001, reason="Invalid API key")
|
|
return
|
|
|
|
await websocket.accept()
|
|
collector_connections.append(websocket)
|
|
client = websocket.client.host if websocket.client else "unknown"
|
|
log.info(f"Collector connected: {client}")
|
|
|
|
try:
|
|
while True:
|
|
try:
|
|
msg = await asyncio.wait_for(websocket.receive_text(), timeout=30)
|
|
data = json.loads(msg)
|
|
|
|
if data.get("type") == "metrics":
|
|
machine_id = data.get("machine_id", "unknown")
|
|
machines[machine_id] = data
|
|
log.debug(f"Metrics from {machine_id}: cpu={data.get('cpu')}%")
|
|
|
|
# Forward to edge
|
|
await forward_to_edge(data)
|
|
|
|
except asyncio.TimeoutError:
|
|
await websocket.send_json({"type": "ping"})
|
|
|
|
except WebSocketDisconnect:
|
|
log.info(f"Collector disconnected: {client}")
|
|
except Exception as e:
|
|
log.error(f"WebSocket error: {e}")
|
|
finally:
|
|
if websocket in collector_connections:
|
|
collector_connections.remove(websocket)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
|
|
log.info("Starting sysmonstm hub")
|
|
log.info(f" API key: {'configured' if API_KEY else 'not set (open)'}")
|
|
log.info(f" Edge URL: {EDGE_URL or 'not configured (local only)'}")
|
|
uvicorn.run(app, host="0.0.0.0", port=8080)
|