monitor.py - master - Befun/FamalyControlle

monitor.py
282 строки11 KB
BeFun
Init project
07 май 2026, 10:43
07 май 2026, 10:437e96900
"""
Десктоп-агент.
Мониторит процессы, блокирует запрещённые приложения,
контролирует загрузки, выполняет команды питания.
Общается с сервером через REST API.
"""

import json
import logging
import os
import time
import platform
from datetime import datetime
from pathlib import Path
from typing import Optional

import httpx
import psutil

from core.config import Config
from utils.helpers import (
    safe_kill_process, shutdown_pc, restart_pc,
    cancel_shutdown, is_time_in_schedule, now_iso,
)

logger = logging.getLogger("agent")


class DesktopAgent:
    """Агент, работающий на контролируемом ПК."""

    def __init__(self):
        self.config = Config
        self.server_url = self.config.agent_server_url
        self.poll_interval = self.config.agent_poll_interval
        self.block_categories = self.config.block_categories
        self.download_dirs = self.config.download_dirs
        self._running = False
        self._killed_pids: set[int] = set()      # PIDs, убитые в текущем цикле (кратков. кэш)
        self._last_blocked_extensions: list[str] = self.block_categories.get(
            "downloads.extensions", [".exe", ".msi", ".bat", ".torrent", ".zip", ".rar", ".7z"]
        )

    # ── Основной цикл ─────────────────────────────────────────

    def start(self) -> None:
        """Запуск агента (блокирующий)."""
        logger.info("Десктоп-агент запущен")
        logger.info("Категории блокировки: %s", list(self.block_categories.keys()))
        logger.info("Интервал проверки: %d сек", self.poll_interval)
        self._running = True

        while self._running:
            try:
                self._tick()
            except Exception as e:
                logger.error("Ошибка в цикле агента: %s", e, exc_info=True)

            time.sleep(self.poll_interval)

    def stop(self) -> None:
        """Остановить агента."""
        self._running = False
        logger.info("Десктоп-агент остановлен")

    def _tick(self) -> None:
        """Один такт мониторинга."""
        # 1. Получить активные ограничения с сервера
        restrictions = self._fetch_restrictions()
        if restrictions is None:
            logger.warning("Не удалось получить ограничения от сервера")
            return

        # 2. Получить расписания и проверить, активно ли сейчас
        schedules = self._fetch_schedules()
        schedule_active = self._check_schedules(schedules) if schedules else False

        if schedule_active:
            # Во время расписания сна/отдыха — блокируем всё
            self._block_all_running()
            self._monitor_downloads(block=True)
            return

        # 3. Собрать целевые процессы из ограничений
        targets = self._build_targets(restrictions)

        if not targets:
            self._monitor_downloads(block="downloads" in self._get_active_categories(restrictions))
            return

        # 4. Мониторинг процессов
        self._enforce_blocks(targets)

        # 5. Мониторинг загрузок
        block_downloads = any(r["category"] == "downloads" for r in restrictions)
        self._monitor_downloads(block=block_downloads)

        # 6. Очистить кэш убитых PID
        self._killed_pids.clear()

    # ── Сетевые запросы к серверу ─────────────────────────────

    def _fetch_restrictions(self) -> Optional[list[dict]]:
        try:
            resp = httpx.get(
                f"{self.server_url}/api/restrictions/active",
                timeout=5,
            )
            if resp.status_code == 200:
                return resp.json()
            return None
        except Exception as e:
            logger.error("Ошибка получения ограничений: %s", e)
            return None

    def _fetch_schedules(self) -> Optional[list[dict]]:
        try:
            resp = httpx.get(
                f"{self.server_url}/api/schedules",
                timeout=5,
            )
            if resp.status_code == 200:
                return resp.json()
            return None
        except Exception as e:
            logger.error("Ошибка получения расписаний: %s", e)
            return None

    def _report_event(self, event: str, details: str = "") -> None:
        """Отправить событие на WebSocket сервера (если доступен)."""
        # Используем HTTP как fallback для логирования событий
        try:
            httpx.post(
                f"{self.server_url}/api/agent/event",
                json={"event": event, "details": details, "timestamp": now_iso()},
                timeout=3,
            )
        except Exception:
            pass

    # ── Обработка ограничений ─────────────────────────────────

    def _get_active_categories(self, restrictions: list[dict]) -> list[str]:
        cats = set()
        for r in restrictions:
            if r.get("is_active"):
                cats.add(r["category"])
        return list(cats)

    def _build_targets(self, restrictions: list[dict]) -> list[str]:
        """
        Построить список имён процессов для блокировки.
        Учитывает категории из config и конкретные приложения.
        """
        targets: list[str] = []
        active_categories: set[str] = set()

        for r in restrictions:
            if not r.get("is_active"):
                continue

            category = r["category"]
            target = r.get("target", "*")

            if category == "app" and target != "*":
                # Конкретное приложение
                targets.append(target.lower())
            elif category in ("games", "browsers"):
                active_categories.add(category)

        # Добавить все процессы из активных категорий
        for cat in active_categories:
            cat_key = cat if cat in self.block_categories else None
            if cat_key:
                targets.extend(self.block_categories[cat_key])

        return list(set(targets))

    # ── Блокировка процессов ──────────────────────────────────

    def _enforce_blocks(self, targets: list[str]) -> None:
        """Проверить все процессы и убить запрещённые."""
        for proc in psutil.process_iter(["pid", "name", "exe"]):
            try:
                proc_info = proc.info
                pid = proc_info["pid"]
                name = (proc_info["name"] or "").lower()

                if pid in self._killed_pids:
                    continue

                for target in targets:
                    if name == target or name.startswith(target):
                        logger.info(
                            "Блокировка: PID=%d процесс='%s' (цель='%s')",
                            pid, name, target,
                        )
                        if safe_kill_process(pid):
                            self._killed_pids.add(pid)
                            self._report_event("process_killed", f"{name} (PID {pid})")
                        break
            except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                continue

    def _block_all_running(self) -> None:
        """Заблокировать все процессы из всех категорий (во время расписания)."""
        all_targets: list[str] = []
        for cat, procs in self.block_categories.items():
            if isinstance(procs, list):
                all_targets.extend(procs)
        self._enforce_blocks(list(set(all_targets)))

    # ── Мониторинг загрузок ───────────────────────────────────

    def _monitor_downloads(self, block: bool) -> None:
        """Проверить папки загрузок на запрещённые файлы."""
        if not block or not self._last_blocked_extensions:
            return

        for dir_path in self.download_dirs:
            expanded = os.path.expanduser(dir_path)
            if not os.path.isdir(expanded):
                continue

            try:
                for entry in os.scandir(expanded):
                    if not entry.is_file():
                        continue
                    ext = Path(entry.name).suffix.lower()
                    if ext in self._last_blocked_extensions:
                        # Удаляем запрещённый файл
                        try:
                            os.remove(entry.path)
                            logger.info(
                                "Удалена загрузка: %s (расширение: %s)",
                                entry.path, ext,
                            )
                            self._report_event(
                                "download_blocked",
                                f"{entry.name} ({ext})",
                            )
                        except Exception as e:
                            logger.error("Не удалось удалить %s: %s", entry.path, e)
            except PermissionError:
                logger.warning("Нет доступа к папке: %s", expanded)

    # ── Проверка расписаний ───────────────────────────────────

    def _check_schedules(self, schedules: list[dict]) -> bool:
        """Проверить, есть ли активное расписание прямо сейчас."""
        for sch in schedules:
            if not sch.get("is_active"):
                continue
            days = sch.get("days", [])
            if is_time_in_schedule(
                sch.get("start_time", "00:00"),
                sch.get("end_time", "23:59"),
                days,
            ):
                logger.info(
                    "Расписание активно: %s (%s - %s)",
                    sch.get("name", "?"),
                    sch.get("start_time"),
                    sch.get("end_time"),
                )
                return True
        return False


# ── Экспорт для запуска ───────────────────────────────────────

def run_agent():
    """Точка входа для десктоп-агента."""
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(name)s] %(levelname)s: %(message)s",
    )
    agent = DesktopAgent()
    try:
        agent.start()
    except KeyboardInterrupt:
        agent.stop()