calvana/ayn-antivirus/ayn_antivirus/dashboard/collector.py

"""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(),
        }