"""Plugin that detects whether IDS has set on non-resident agent start.""" import asyncio import contextlib import logging from defence360agent import utils from defence360agent.utils import safe_cancel_task from defence360agent.contracts import plugins, messages from im360.internals import strategy from im360.contracts import config logger = logging.getLogger(__name__) class StrategyGetter(plugins.MessageSource): """Send StrategyChange message on non-resident agent start.""" def __init__(self, **kwargs): super().__init__(**kwargs) self.__task = None async def create_source(self, loop, sink): self._loop = loop self._sink = sink self.__task = loop.create_task(self._detect_strategy_change()) async def shutdown(self): if self.__task is None or self.__task.done(): return await safe_cancel_task(self.__task) @utils.abort_agent_on(Exception) async def _detect_strategy_change(self): """Check whether IDS has changed periodically.""" with contextlib.suppress(asyncio.CancelledError): while True: await self._refresh_strategy() await asyncio.sleep( config.Subsys.THIRD_PARTY_IDS_CHECK_TIMEOUT ) async def _refresh_strategy(self): try: new_strategy = strategy.Strategy.get() except asyncio.CancelledError: pass except Exception as e: logger.error("Failed to get strategy, %s", e) else: if new_strategy != strategy.Strategy.current: await self._on_strategy_changed(new_strategy) async def _on_strategy_changed(self, new_strategy): logger.info( "Got new strategy: %s -> %s", strategy.Strategy.current, new_strategy, ) strategy.Strategy.current = new_strategy await self._sink.process_message( messages.MessageType.StrategyChange(strategy=new_strategy) )