22

✏️ Editing: lfd.py

import logging
import os

from defence360agent import utils
from defence360agent.contracts import config as common_config
from defence360agent.contracts.messages import MessageType
from defence360agent.contracts.plugins import (
    MessageSink,
    MessageSource,
    expect,
)
from im360.contracts import config
from im360.internals import strategy
from im360.subsys import csf

logger = logging.getLogger(__name__)

class LFD(MessageSink, MessageSource):
    PROCESSING_ORDER = MessageSink.ProcessingOrder.LFD
    STRATEGY = strategy.Strategy.CSF_COOP_STRATEGY
    AVAILABLE_ON_FREEMIUM = False

BLOCK_REPORT_SCRIPT = os.path.join(
        common_config.Packaging.DATADIR, "scripts", "lfd_block.py"
    )
    USER_SCRIPT_LINK = os.path.join(
        common_config.Packaging.DATADIR, "scripts", "block_report_user"
    )

@property
    def _script_installed(self):
        try:
            current_script = csf.Config("BLOCK_REPORT").get()
        except (FileNotFoundError, NotADirectoryError):
            return False
        else:
            return os.path.realpath(current_script) == os.path.realpath(
                self.BLOCK_REPORT_SCRIPT
            )

async def create_sink(self, loop):
        self._loop = loop

async def create_source(self, loop, sink):
        self._sink = sink

async def shutdown(self):
        # TODO: we should revert script on shutdown of imunify360 service
        # or uninstall
        pass

@expect(MessageType.StrategyChange)
    async def on_strategy_change(self, message):
        await self._switch_state(
            message.strategy,
            config.CSFIntegration.ENABLED,
            self._script_installed,
        )

@expect(MessageType.ConfigUpdate)
    async def on_config_change(self, _):
        await self._switch_state(
            strategy.Strategy.current,
            config.CSFIntegration.ENABLED,
            self._script_installed,
        )

@utils.log_error_and_ignore()
    async def _switch_state(
        self, current_strategy, is_enabled, is_script_installed
    ):
        should_script_be_installed = is_enabled and (
            current_strategy == self.STRATEGY
        )
        if is_script_installed:
            if not should_script_be_installed:
                await self._revert_script()
        elif should_script_be_installed:
            await self._setup_script()

async def _setup_script(self):
        logger.info("Setup lfd BLOCK_REPORT to Imunify360 script")
        old_script = csf.Config("BLOCK_REPORT").set(self.BLOCK_REPORT_SCRIPT)
        if old_script and (
            os.path.realpath(old_script)
            != os.path.realpath(self.BLOCK_REPORT_SCRIPT)
        ):
            logger.info(
                "Creating a symlink %s to %s",
                self.USER_SCRIPT_LINK,
                old_script,
            )
            os.symlink(old_script, self.USER_SCRIPT_LINK)
        await csf.lfd_restart()

async def _revert_script(self):
        logger.info("Reverting lfd BLOCK_REPORT")
        if os.path.islink(self.USER_SCRIPT_LINK):
            user_script = os.readlink(self.USER_SCRIPT_LINK)
            os.remove(self.USER_SCRIPT_LINK)
        else:
            user_script = ""

try:
            csf.Config("BLOCK_REPORT").set(user_script)
        except FileNotFoundError as e:
            logger.warning("Failed to revert BLOCK_REPORT value: %s", e)
        else:
            if await csf.is_running():
                await csf.lfd_restart()

MR SKB