"""Background metrics collector for the AYN Antivirus dashboard.""" from __future__ import annotations import asyncio import logging import os import random from datetime import datetime from typing import Any, Dict, Optional import psutil from ayn_antivirus.constants import DASHBOARD_COLLECTOR_INTERVAL logger = logging.getLogger("ayn_antivirus.dashboard.collector") class MetricsCollector: """Periodically sample system metrics and store them in the dashboard DB. Parameters ---------- store: A :class:`DashboardStore` instance to write metrics into. interval: Seconds between samples. """ def __init__(self, store: Any, interval: int = DASHBOARD_COLLECTOR_INTERVAL) -> None: self.store = store self.interval = interval self._task: Optional[asyncio.Task] = None self._running = False async def start(self) -> None: """Begin collecting metrics on a background asyncio task.""" self._running = True self._task = asyncio.create_task(self._collect_loop()) logger.info("Metrics collector started (interval=%ds)", self.interval) async def stop(self) -> None: """Cancel the background task and wait for it to finish.""" self._running = False if self._task: self._task.cancel() try: await self._task except asyncio.CancelledError: pass logger.info("Metrics collector stopped") # ------------------------------------------------------------------ # Internal loop # ------------------------------------------------------------------ async def _collect_loop(self) -> None: while self._running: try: await asyncio.to_thread(self._sample) except Exception as exc: logger.error("Collector error: %s", exc) await asyncio.sleep(self.interval) def _sample(self) -> None: """Take a single metric snapshot and persist it.""" cpu = psutil.cpu_percent(interval=1) mem = psutil.virtual_memory() disks = [] for part in psutil.disk_partitions(all=False): try: usage = psutil.disk_usage(part.mountpoint) disks.append({ "mount": part.mountpoint, "device": part.device, "total": usage.total, "used": usage.used, "free": usage.free, "percent": usage.percent, }) except (PermissionError, OSError): continue try: load = list(os.getloadavg()) except (OSError, AttributeError): load = [0.0, 0.0, 0.0] try: net_conns = len(psutil.net_connections(kind="inet")) except (psutil.AccessDenied, OSError): net_conns = 0 self.store.record_metric( cpu=cpu, mem_pct=mem.percent, mem_used=mem.used, mem_total=mem.total, disk_usage=disks, load_avg=load, net_conns=net_conns, ) # Periodic cleanup (~1 in 100 samples). if random.randint(1, 100) == 1: self.store.cleanup_old_metrics() # ------------------------------------------------------------------ # One-shot snapshot (no storage) # ------------------------------------------------------------------ @staticmethod def get_snapshot() -> Dict[str, Any]: """Return a live system snapshot without persisting it.""" cpu = psutil.cpu_percent(interval=0.1) cpu_per_core = psutil.cpu_percent(interval=0.1, percpu=True) cpu_freq = psutil.cpu_freq(percpu=False) mem = psutil.virtual_memory() swap = psutil.swap_memory() disks = [] for part in psutil.disk_partitions(all=False): try: usage = psutil.disk_usage(part.mountpoint) disks.append({ "mount": part.mountpoint, "device": part.device, "total": usage.total, "used": usage.used, "percent": usage.percent, }) except (PermissionError, OSError): continue try: load = list(os.getloadavg()) except (OSError, AttributeError): load = [0.0, 0.0, 0.0] try: net_conns = len(psutil.net_connections(kind="inet")) except (psutil.AccessDenied, OSError): net_conns = 0 # Top processes by CPU top_procs = [] try: for p in sorted(psutil.process_iter(['pid', 'name', 'cpu_percent', 'memory_percent']), key=lambda x: x.info.get('cpu_percent', 0) or 0, reverse=True)[:8]: info = p.info if (info.get('cpu_percent') or 0) > 0.1: top_procs.append({ "pid": info['pid'], "name": info['name'] or '?', "cpu": round(info.get('cpu_percent', 0) or 0, 1), "mem": round(info.get('memory_percent', 0) or 0, 1), }) except Exception: pass return { "cpu_percent": cpu, "cpu_per_core": cpu_per_core, "cpu_cores": psutil.cpu_count(logical=True), "cpu_freq_mhz": round(cpu_freq.current) if cpu_freq else 0, "mem_percent": mem.percent, "mem_used": mem.used, "mem_total": mem.total, "mem_available": mem.available, "mem_cached": getattr(mem, 'cached', 0), "mem_buffers": getattr(mem, 'buffers', 0), "swap_percent": swap.percent, "swap_used": swap.used, "swap_total": swap.total, "disk_usage": disks, "load_avg": load, "net_connections": net_conns, "top_processes": top_procs, "timestamp": datetime.utcnow().isoformat(), }