process_monitor.py - master - Befun/FamileControll

process_monitor.py
641 строка26 KB
BeFun
Fixed bugs
02 мар 2026, 22:09
02 мар 2026, 22:09036d2e9
"""
Мониторинг и обнаружение процессов (игры, браузеры)
"""
import os
import sys
import subprocess
import psutil
import asyncio
import logging
from pathlib import Path
from typing import List, Dict, Optional, Set
from datetime import datetime

from config import KNOWN_GAMES, KNOWN_BROWSERS, SEARCH_PATHS
from database import Database

logger = logging.getLogger(__name__)

# Таймауты
TERMINATE_TIMEOUT = 3   # секунд ждать после terminate()
KILL_TIMEOUT = 3         # секунд ждать после kill()
TASKKILL_TIMEOUT = 5     # секунд ждать после taskkill


class ProcessMonitor:
    """Класс для мониторинга процессов на компьютере"""

    def __init__(self, db: Database):
        self.db = db
        self._running_processes: Dict[str, dict] = {}
        self._discovered_apps: Dict[str, dict] = {}
        self._stop_flag = False

    @staticmethod
    def get_running_processes() -> List[Dict]:
        """Получение списка всех запущенных процессов"""
        processes = []
        for proc in psutil.process_iter(['pid', 'name', 'exe', 'cmdline', 'username']):
            try:
                proc_info = proc.info
                if proc_info['name']:
                    processes.append({
                        'pid': proc_info['pid'],
                        'name': proc_info['name'],
                        'exe': proc_info.get('exe', ''),
                        'cmdline': ' '.join(proc_info.get('cmdline') or []),
                        'username': proc_info.get('username', '')
                    })
            except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                continue
        return processes

    def identify_app_type(self, process_name: str, exe_path: str = None) -> tuple:
        """
        Определение типа приложения (игра, браузер, другое)
        Возвращает (app_type, category)
        """
        name_lower = process_name.lower().replace('.exe', '')

        # Проверяем игры
        for game in KNOWN_GAMES:
            if game.lower() in name_lower:
                return ('game', 'game')

        # Проверяем браузеры
        for browser in KNOWN_BROWSERS:
            if browser.lower() in name_lower:
                return ('browser', 'browser')

        # Проверяем по пути, если есть
        if exe_path:
            path_lower = exe_path.lower()

            game_indicators = [
                '\\games\\', '\\game\\', 'steamapps', 'epic games',
                'origin games', 'ubisoft', 'battle.net', 'gog galaxy'
            ]
            for indicator in game_indicators:
                if indicator in path_lower:
                    return ('game', 'game')

            browser_indicators = [
                '\\browser\\', 'chrome\\application', 'firefox\\',
                'edge\\application', 'opera\\'
            ]
            for indicator in browser_indicators:
                if indicator in path_lower:
                    return ('browser', 'browser')

        return ('app', 'other')

    def scan_for_applications(self) -> Dict[str, List[Dict]]:
        """
        Сканирование системы на наличие игр и браузеров
        Возвращает словарь {'games': [...], 'browsers': [...], 'other': [...]}
        """
        result = {'games': [], 'browsers': [], 'other': []}
        found_paths = set()

        for search_path in SEARCH_PATHS:
            if not search_path or not os.path.exists(search_path):
                continue

            try:
                for root, dirs, files in os.walk(search_path):
                    dirs[:] = [
                        d for d in dirs if not d.startswith('.') and
                        d.lower() not in ['windows', 'system32', 'programdata']
                    ]

                    for file in files:
                        if file.endswith('.exe'):
                            full_path = os.path.join(root, file)
                            if full_path in found_paths:
                                continue
                            found_paths.add(full_path)

                            name = file.replace('.exe', '')
                            app_type, category = self.identify_app_type(name, full_path)

                            app_info = {
                                'name': name,
                                'path': full_path,
                                'type': app_type,
                                'category': category
                            }

                            if app_type == 'game':
                                result['games'].append(app_info)
                            elif app_type == 'browser':
                                result['browsers'].append(app_info)
                            else:
                                result['other'].append(app_info)

                            self.db.add_discovered_app(name, full_path, app_type, category)

            except (PermissionError, OSError) as e:
                logger.debug(f"Ошибка доступа к {search_path}: {e}")
                continue

        return result

    def get_monitored_processes(self) -> List[Dict]:
        """Получение списка процессов, которые нужно мониторить (игры и браузеры)"""
        processes = self.get_running_processes()
        monitored = []

        for proc in processes:
            app_type, category = self.identify_app_type(proc['name'], proc['exe'])
            if app_type in ('game', 'browser'):
                proc['app_type'] = app_type
                proc['category'] = category
                monitored.append(proc)

        return monitored

    def is_process_running(self, process_name: str) -> bool:
        """Проверка, запущен ли процесс"""
        for proc in psutil.process_iter(['name']):
            try:
                if proc.info['name'] and process_name.lower() in proc.info['name'].lower():
                    return True
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                continue
        return False

    # ------------------------------------------------------------------ #
    #  ГАРАНТИРОВАННОЕ ЗАВЕРШЕНИЕ ПРОЦЕССА — внутренние хелперы
    # ------------------------------------------------------------------ #

    @staticmethod
    def _collect_matching_procs(process_name: str) -> List[psutil.Process]:
        """Собрать все живые процессы, чьё имя содержит *process_name*."""
        name_lower = process_name.lower()
        matched: List[psutil.Process] = []
        for proc in psutil.process_iter(['name', 'pid']):
            try:
                pname = proc.info.get('name') or ''
                if name_lower in pname.lower():
                    matched.append(proc)
            except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                continue
        return matched

    @staticmethod
    def _proc_alive(proc: psutil.Process) -> bool:
        """Проверить, жив ли процесс."""
        try:
            return proc.is_running() and proc.status() != psutil.STATUS_ZOMBIE
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            return False

    @staticmethod
    def _force_kill_single(proc: psutil.Process, timeout: int = TERMINATE_TIMEOUT) -> bool:
        """
        Трёхступенчатое завершение *одного* psutil.Process:
        1. terminate()  → ждём timeout
        2. kill()       → ждём timeout
        3. (Windows) taskkill /F /PID  → ждём timeout
        Возвращает True, если процесс мёртв.
        """
        pid = proc.pid
        name = ''
        try:
            name = proc.name()
        except Exception:
            pass

        # --- Шаг 1: terminate (SIGTERM / TerminateProcess) ---
        try:
            proc.terminate()
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            pass
        try:
            proc.wait(timeout=timeout)
        except psutil.TimeoutExpired:
            pass
        except psutil.NoSuchProcess:
            logger.info(f"[terminate] Процесс мёртв: {name} PID={pid}")
            return True

        if not ProcessMonitor._proc_alive(proc):
            logger.info(f"[terminate] Процесс завершён: {name} PID={pid}")
            return True

        # --- Шаг 2: kill (SIGKILL / TerminateProcess с exit_code=1) ---
        try:
            proc.kill()
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            pass
        try:
            proc.wait(timeout=KILL_TIMEOUT)
        except psutil.TimeoutExpired:
            pass
        except psutil.NoSuchProcess:
            logger.info(f"[kill] Процесс мёртв: {name} PID={pid}")
            return True

        if not ProcessMonitor._proc_alive(proc):
            logger.info(f"[kill] Процесс завершён: {name} PID={pid}")
            return True

        # --- Шаг 3: OS-уровень ---
        if sys.platform == 'win32':
            try:
                subprocess.run(
                    ['taskkill', '/F', '/PID', str(pid)],
                    stdout=subprocess.DEVNULL,
                    stderr=subprocess.DEVNULL,
                    timeout=TASKKILL_TIMEOUT,
                )
            except Exception as exc:
                logger.warning(f"taskkill /F /PID {pid} провалился: {exc}")
        else:
            # Linux/macOS
            try:
                os.kill(pid, 9)
            except OSError:
                pass

        # Финальная проверка
        try:
            proc.wait(timeout=KILL_TIMEOUT)
        except Exception:
            pass

        alive = ProcessMonitor._proc_alive(proc)
        if alive:
            logger.error(f"НЕ УДАЛОСЬ завершить: {name} PID={pid}")
        else:
            logger.info(f"[os-level] Процесс завершён: {name} PID={pid}")
        return not alive

    # ------------------------------------------------------------------ #
    #  ПУБЛИЧНЫЙ ИНТЕРФЕЙС — тот же, что был
    # ------------------------------------------------------------------ #

    def kill_process(self, process_name: str) -> bool:
        """
        Гарантированное завершение **всех** процессов, содержащих
        *process_name* в имени.

        Стратегия (для каждого найденного процесса):
          1. terminate()  → ожидание
          2. kill()        → ожидание
          3. taskkill /F /PID (Windows) или SIGKILL (Linux)
          4. taskkill /F /IM <image_name> (Windows) — «контрольный»
          5. Повторный поиск: если кто-то возродился — убиваем снова

        Возвращает True, если хотя бы один процесс был убит
        (или если процессов с таким именем уже не существовало).
        """
        name_lower = process_name.lower()
        total_killed = 0
        any_found = False

        # === Раунд 1: собираем и убиваем по-одному ===
        procs = self._collect_matching_procs(process_name)
        if not procs:
            logger.info(f"Процесс '{process_name}' не найден среди запущенных")
            return False

        any_found = True
        # Сначала убиваем дочерние, потом родительские
        # Сортируем: дети → родители (больший PID обычно = потомок, но
        # надёжнее через children()).
        all_targets: List[psutil.Process] = []
        for p in procs:
            try:
                children = p.children(recursive=True)
                all_targets.extend(children)
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                pass
            all_targets.append(p)

        # Убираем дубли, сохраняя порядок
        seen_pids: Set[int] = set()
        unique_targets: List[psutil.Process] = []
        for p in all_targets:
            if p.pid not in seen_pids:
                seen_pids.add(p.pid)
                unique_targets.append(p)

        for proc in unique_targets:
            if self._force_kill_single(proc):
                total_killed += 1

        # === Раунд 2 (Windows): taskkill /F /IM — «ковровая» зачистка ===
        if sys.platform == 'win32':
            # Собираем уникальные image-имена, которые мы хотели убить
            image_names: Set[str] = set()
            for proc in procs:
                try:
                    img = proc.name()
                    if img:
                        image_names.add(img)
                except Exception:
                    pass

            # Если имя процесса выглядит как «chrome.exe» — используем как есть
            # Иначе добавляем .exe, если его нет
            candidate = process_name if '.' in process_name else process_name + '.exe'
            image_names.add(candidate)

            for img in image_names:
                try:
                    subprocess.run(
                        ['taskkill', '/F', '/IM', img],
                        stdout=subprocess.DEVNULL,
                        stderr=subprocess.DEVNULL,
                        timeout=TASKKILL_TIMEOUT,
                    )
                except Exception:
                    pass

        # === Раунд 3: финальная проверка, пере-поиск ===
        # Даём ОС время на зачистку
        import time
        time.sleep(0.5)

        survivors = self._collect_matching_procs(process_name)
        if survivors:
            logger.warning(
                f"После основной зачистки остались {len(survivors)} процесс(ов) "
                f"'{process_name}'. Повторная попытка…"
            )
            for proc in survivors:
                self._force_kill_single(proc, timeout=KILL_TIMEOUT)

            # Последняя проверка
            time.sleep(0.3)
            final_survivors = self._collect_matching_procs(process_name)
            if final_survivors:
                logger.error(
                    f"Не удалось завершить {len(final_survivors)} процесс(ов) "
                    f"'{process_name}': "
                    f"{[(p.pid, p.name()) for p in final_survivors]}"
                )
            else:
                total_killed += len(survivors)
                logger.info(f"Все процессы '{process_name}' успешно завершены (повторная попытка)")
        else:
            if any_found:
                logger.info(f"Все процессы '{process_name}' успешно завершены")

        return total_killed > 0 or (any_found and not self.is_process_running(process_name))

    def kill_process_by_pid(self, pid: int) -> bool:
        """Гарантированное завершение процесса по PID"""
        try:
            proc = psutil.Process(pid)
        except psutil.NoSuchProcess:
            logger.info(f"Процесс PID={pid} уже не существует")
            return True
        except psutil.AccessDenied as e:
            logger.warning(f"Нет доступа к процессу PID={pid}: {e}")
            # Попытка через taskkill
            if sys.platform == 'win32':
                try:
                    result = subprocess.run(
                        ['taskkill', '/F', '/PID', str(pid)],
                        capture_output=True, text=True, timeout=TASKKILL_TIMEOUT,
                    )
                    if result.returncode == 0:
                        logger.info(f"taskkill /F /PID {pid}: успех")
                        return True
                except Exception:
                    pass
            return False

        return self._force_kill_single(proc)

    async def monitor_loop(self, callback=None, interval: int = 5):
        """
        Основной цикл мониторинга
        callback - функция, вызываемая при обнаружении заблокированного процесса
        """
        self._stop_flag = False

        while not self._stop_flag:
            try:
                # Очистка истекших правил
                self.db.cleanup_expired_rules()

                # Получение списка процессов
                monitored = self.get_monitored_processes()

                # Проверка каждого процесса
                for proc in monitored:
                    proc_key = f"{proc['name']}_{proc['pid']}"

                    # Проверяем, заблокировано ли приложение
                    if self.db.is_app_blocked(proc['name'], proc['exe']):
                        if proc_key not in self._running_processes:
                            self._running_processes[proc_key] = {
                                'process': proc,
                                'blocked_at': datetime.now(),
                                'session_id': self.db.start_app_session(
                                    proc['name'], proc['exe']
                                )
                            }

                            # Гарантированное завершение
                            if self.kill_process_by_pid(proc['pid']):
                                if callback:
                                    await callback(proc, 'blocked')
                    else:
                        if proc_key not in self._running_processes:
                            self._running_processes[proc_key] = {
                                'process': proc,
                                'started_at': datetime.now(),
                                'session_id': self.db.start_app_session(
                                    proc['name'], proc['exe']
                                )
                            }

                # Очистка завершённых процессов из отслеживания
                current_pids = {p['pid'] for p in monitored}
                keys_to_remove = []
                for key, data in self._running_processes.items():
                    if data['process']['pid'] not in current_pids:
                        if 'session_id' in data:
                            self.db.end_app_session(data['session_id'])
                        keys_to_remove.append(key)

                for key in keys_to_remove:
                    del self._running_processes[key]

                await asyncio.sleep(interval)

            except Exception as e:
                logger.error(f"Ошибка в цикле мониторинга: {e}")
                await asyncio.sleep(interval)

    def stop_monitoring(self):
        """Остановка мониторинга"""
        self._stop_flag = True

    def get_tracked_processes(self) -> List[Dict]:
        """Получение списка отслеживаемых процессов"""
        return [data['process'] for data in self._running_processes.values()]


class AppScanner:
    """Класс для сканирования и обнаружения установленных приложений"""

    def __init__(self, db: Database):
        self.db = db
        self._games_cache = None
        self._browsers_cache = None

    def scan_games(self) -> List[Dict]:
        """Сканирование на установленные игры"""
        if self._games_cache is not None:
            return self._games_cache

        games = []

        steam_paths = [
            os.path.expandvars(r"%ProgramFiles(x86)%\Steam\steamapps\common"),
            os.path.expandvars(r"%ProgramFiles%\Steam\steamapps\common"),
            os.path.expanduser("~/.steam/steam/steamapps/common"),
        ]

        for steam_path in steam_paths:
            if os.path.exists(steam_path):
                for game_dir in os.listdir(steam_path):
                    full_path = os.path.join(steam_path, game_dir)
                    if os.path.isdir(full_path):
                        for root, _, files in os.walk(full_path):
                            for file in files:
                                if file.endswith('.exe'):
                                    games.append({
                                        'name': game_dir,
                                        'exe': os.path.join(root, file),
                                        'source': 'Steam'
                                    })
                                    break
                            break

        epic_path = os.path.expandvars(r"%ProgramFiles(x86)%\Epic Games")
        if os.path.exists(epic_path):
            for game_dir in os.listdir(epic_path):
                full_path = os.path.join(epic_path, game_dir)
                if os.path.isdir(full_path):
                    for root, _, files in os.walk(full_path):
                        for file in files:
                            if file.endswith('.exe'):
                                games.append({
                                    'name': game_dir,
                                    'exe': os.path.join(root, file),
                                    'source': 'Epic Games'
                                })
                                break
                        break

        standard_paths = ['C:\\Games', 'D:\\Games', 'E:\\Games']
        for path in standard_paths:
            if os.path.exists(path):
                for item in os.listdir(path):
                    full_path = os.path.join(path, item)
                    if os.path.isdir(full_path):
                        for root, _, files in os.walk(full_path):
                            for file in files:
                                if file.endswith('.exe'):
                                    app_type, _ = self._identify_app(
                                        item, os.path.join(root, file)
                                    )
                                    if app_type == 'game':
                                        games.append({
                                            'name': item,
                                            'exe': os.path.join(root, file),
                                            'source': 'Local'
                                        })
                                    break
                            break

        self._games_cache = games
        return games

    def scan_browsers(self) -> List[Dict]:
        """Сканирование на установленные браузеры"""
        if self._browsers_cache is not None:
            return self._browsers_cache

        browsers = []

        browser_paths = {
            'Google Chrome': [
                os.path.expandvars(
                    r"%ProgramFiles%\Google\Chrome\Application\chrome.exe"
                ),
                os.path.expandvars(
                    r"%ProgramFiles(x86)%\Google\Chrome\Application\chrome.exe"
                ),
                os.path.expandvars(
                    r"%LocalAppData%\Google\Chrome\Application\chrome.exe"
                ),
            ],
            'Mozilla Firefox': [
                os.path.expandvars(r"%ProgramFiles%\Mozilla Firefox\firefox.exe"),
                os.path.expandvars(r"%ProgramFiles(x86)%\Mozilla Firefox\firefox.exe"),
            ],
            'Microsoft Edge': [
                os.path.expandvars(
                    r"%ProgramFiles(x86)%\Microsoft\Edge\Application\msedge.exe"
                ),
                os.path.expandvars(
                    r"%ProgramFiles%\Microsoft\Edge\Application\msedge.exe"
                ),
            ],
            'Opera': [
                os.path.expandvars(r"%LocalAppData%\Programs\Opera\launcher.exe"),
                os.path.expandvars(r"%ProgramFiles%\Opera\launcher.exe"),
            ],
            'Brave': [
                os.path.expandvars(
                    r"%LocalAppData%\BraveSoftware\Brave-Browser\Application\brave.exe"
                ),
            ],
            'Yandex Browser': [
                os.path.expandvars(
                    r"%LocalAppData%\Yandex\YandexBrowser\Application\browser.exe"
                ),
            ],
            'Vivaldi': [
                os.path.expandvars(r"%LocalAppData%\Vivaldi\Application\vivaldi.exe"),
            ],
        }

        for browser_name, paths in browser_paths.items():
            for path in paths:
                if os.path.exists(path):
                    browsers.append({
                        'name': browser_name,
                        'exe': path,
                        'source': 'Installed'
                    })
                    break

        self._browsers_cache = browsers
        return browsers

    def _identify_app(self, name: str, exe_path: str) -> tuple:
        """Определение типа приложения"""
        name_lower = name.lower()

        for game in KNOWN_GAMES:
            if game.lower() in name_lower:
                return ('game', 'game')

        for browser in KNOWN_BROWSERS:
            if browser.lower() in name_lower:
                return ('browser', 'browser')

        return ('app', 'other')

    def clear_cache(self):
        """Очистка кэша сканирования"""
        self._games_cache = None
        self._browsers_cache = None