관리-도구

편집 파일: event_monitor.py

import asyncio
import contextlib
import json
from abc import ABC
from logging import getLogger
from pathlib import Path
from typing import Dict, List, Optional

from defence360agent.contracts.config import Core
from defence360agent.contracts.messages import MessageType
from defence360agent.contracts.plugins import MessageSource
from defence360agent.feature_management.plugins.native import (
    NativeFeatureManagementSettingsChange,
)
from defence360agent.plugins.event_monitor_message_processor import (
    EventProcessorBase,
    UserConfigProcessor,
)
from defence360agent.utils import recurring_check

logger = getLogger(__name__)

class EventMonitor(MessageSource, ABC):
    EVENT_DIR = Core.INBOX_HOOKS_DIR
    PATTERN = "*.*.*.*.json"

def __init__(self):
        self._loop = None
        self._sink = None
        self._processors: List[EventProcessorBase] = []
        self._processing_task = None

async def create_source(self, loop, sink):
        self._loop = loop
        self._sink = sink
        self._processors.append(NativeFeatureManagementSettingsChange(loop))
        self._processors.append(UserConfigProcessor(loop))
        self._processing_task = self._loop.create_task(
            self._check_inbox_folder_generate_events()
        )

async def shutdown(self):
        self._processing_task.cancel()
        with contextlib.suppress(asyncio.CancelledError):
            await self._processing_task

@staticmethod
    def _rmfile(file: Path):  # pragma: no cover
        try:
            file.unlink()
        except FileNotFoundError:
            pass  # do nothing if we cannot remove it, just skip it
        except Exception as e:
            logger.warning("Couldn't remove file %s %s", file, e)

@staticmethod
    def _from_json(file: Path) -> Dict:
        return json.loads(file.read_text())

def _event_to_message(self, file) -> Optional[MessageType.cPanelEvent]:
        try:
            username, hook, ts1, ts2, *_ = file.name.split(".")
            ts = float(ts1 + "." + ts2)
        except ValueError:
            logger.warning("hook-event-file detected with wrong name %s", file)
            return None
        try:
            return MessageType.cPanelEvent.from_hook_event(
                username=username,
                hook=hook,
                ts=ts,
                fields=self._from_json(file),
            )
        except FileNotFoundError:  # pragma: no cover
            # already deleted
            logger.warning("hook file disappeared %s", file)
        except json.JSONDecodeError:
            # wrong format or broken json
            logger.warning("hook file have broken json %s", file)
        return None

@recurring_check(30)
    async def _check_inbox_folder_generate_events(self):
        for file in Path(self.EVENT_DIR).glob("*.*.*.json"):
            try:
                message = self._event_to_message(file)
                if message is not None:
                    for processor in self._processors:
                        if await processor.is_enabled():
                            processor.add_message(message)
            except Exception as exc:  # pragma: no cover
                logger.error("Failed to process %s hook event", exc)
            finally:
                self._rmfile(file)
        for processor in self._processors:
            await processor.process_messages()