dipl-edr/agent/main.py

"""EDR monitoring agent.

Collects process and system metrics every 5 seconds using psutil,
monitors filesystem paths with watchdog, and POSTs batches to the backend.
"""

import logging
import os
import sys
import threading
import time
from datetime import datetime, timezone
from pathlib import Path

import psutil
import requests

# Normalise per-process cpu_percent to 0-100% of *total* machine capacity.
# psutil returns cumulative core-time (can exceed 100% on multi-core).
_CPU_COUNT: int = psutil.cpu_count(logical=True) or 1
from watchdog.events import FileSystemEventHandler
from watchdog.observers import Observer

sys.path.insert(0, str(Path(__file__).parent))
from model import AnomalyDetector

# ── configuration ─────────────────────────────────────────────────────────────

BACKEND_URL = os.getenv("BACKEND_URL", "http://localhost:8000")
COLLECT_INTERVAL = int(os.getenv("COLLECT_INTERVAL", "5"))
WATCH_PATHS = [p for p in ["/etc", "/tmp", "/var/log"] if Path(p).exists()]

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
logger = logging.getLogger(__name__)

# ── shared fs-event buffer ────────────────────────────────────────────────────

_fs_events: list[dict] = []
_fs_lock = threading.Lock()


class _FSHandler(FileSystemEventHandler):
    """Records create/modify/delete events into the shared buffer."""

    def on_created(self, event) -> None:
        if not event.is_directory:
            self._push(event.src_path, "create")

    def on_modified(self, event) -> None:
        if not event.is_directory:
            self._push(event.src_path, "modify")

    def on_deleted(self, event) -> None:
        self._push(event.src_path, "delete")

    @staticmethod
    def _push(path: str, event_type: str) -> None:
        with _fs_lock:
            _fs_events.append(
                {
                    "path": path,
                    "event_type": event_type,
                    "timestamp": datetime.now(timezone.utc).isoformat(),
                }
            )
            # Bound buffer to avoid unbounded growth between flushes
            if len(_fs_events) > 500:
                _fs_events.pop(0)


def _drain_fs_events() -> list[dict]:
    """Return accumulated events and clear the buffer."""
    with _fs_lock:
        out = list(_fs_events)
        _fs_events.clear()
    return out


# ── metric collection ─────────────────────────────────────────────────────────


def _proc_connections(proc: psutil.Process) -> int:
    """Return TCP connection count, compatible with all psutil versions."""
    for method in ("net_connections", "connections"):
        fn = getattr(proc, method, None)
        if fn is None:
            continue
        try:
            return len(fn())
        except (psutil.NoSuchProcess, psutil.AccessDenied, OSError):
            return 0
    return 0


def collect_processes() -> list[dict]:
    """Return per-process metrics for all running processes."""
    result = []
    # 'connections' was removed from process_iter attrs in psutil 6.x
    attrs = ["pid", "name", "cpu_percent", "memory_info", "open_files"]
    for proc in psutil.process_iter(attrs):
        try:
            info = proc.info
            mem_mb = (info["memory_info"].rss / 1024 / 1024
                      if info["memory_info"] else 0.0)
            open_files = len(info["open_files"]) if info["open_files"] else 0
            connections = _proc_connections(proc)
            result.append(
                {
                    "pid": info["pid"],
                    "name": info["name"] or "unknown",
                    "cpu_percent": round(float(info["cpu_percent"] or 0.0) / _CPU_COUNT, 2),
                    "memory_mb": round(mem_mb, 2),
                    "open_files": open_files,
                    "connections": connections,
                }
            )
        except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
            pass
    return result


def collect_system() -> dict:
    """Return system-wide CPU, RAM, and network connection totals."""
    return {
        "cpu_percent": psutil.cpu_percent(interval=None),
        "ram_percent": psutil.virtual_memory().percent,
        "network_connections": len(psutil.net_connections()),
    }


# ── backend communication ─────────────────────────────────────────────────────


def send_batch(detector: AnomalyDetector) -> None:
    """Collect metrics, annotate anomalies, and POST to backend."""
    processes = collect_processes()
    system = collect_system()
    fs_events = _drain_fs_events()

    processes = detector.predict(processes)

    anomaly_count = sum(1 for p in processes if p.get("is_anomaly"))

    payload = {
        "processes": processes,
        "system": system,
        "file_events": fs_events,
        "timestamp": datetime.now(timezone.utc).isoformat(),
    }

    try:
        resp = requests.post(
            f"{BACKEND_URL}/api/metrics",
            json=payload,
            timeout=10,
        )
        resp.raise_for_status()
        logger.info(
            "Sent %d procs | anomalies=%d | fs_events=%d",
            len(processes),
            anomaly_count,
            len(fs_events),
        )
    except requests.exceptions.ConnectionError:
        logger.warning("Backend unreachable — skipping batch")
    except requests.exceptions.RequestException as exc:
        logger.error("POST /api/metrics failed: %s", exc)


# ── entry point ───────────────────────────────────────────────────────────────


def main() -> None:
    """Start filesystem watcher and run the collection loop."""
    logger.info("EDR Agent starting (backend=%s, interval=%ds)", BACKEND_URL, COLLECT_INTERVAL)

    # Load or train the anomaly model
    detector = AnomalyDetector()
    detector.load_or_train()

    # Start filesystem observer
    handler = _FSHandler()
    observer = Observer()
    for path in WATCH_PATHS:
        observer.schedule(handler, path, recursive=True)
        logger.info("Watching path: %s", path)
    observer.start()

    # Warm up cpu_percent (first call always returns 0.0)
    psutil.cpu_percent(interval=None)
    for proc in psutil.process_iter(["cpu_percent"]):
        try:
            proc.cpu_percent(interval=None)
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            pass
    time.sleep(1)

    try:
        while True:
            send_batch(detector)
            time.sleep(COLLECT_INTERVAL)
    except KeyboardInterrupt:
        logger.info("Agent stopped by user")
    finally:
        observer.stop()
        observer.join()


if __name__ == "__main__":
    main()