1#@!#!123s

: / opt / imunify360 / venv / lib64 / python3.11 / site-packages / im360 / plugins / protector /
Filename : lfd.py
back
import logging
import os

from defence360agent import utils
from defence360agent.contracts import config as common_config
from defence360agent.contracts.messages import MessageType, Reject
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) == 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):
        if self._script_installed:
            await self._revert_script()

    @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")
        # setup lfd script
        old_script = os.path.realpath(
            csf.Config("BLOCK_REPORT").set(self.BLOCK_REPORT_SCRIPT)
        )
        if old_script and (old_script != 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()

    @expect(MessageType.SensorAlert, plugin_id=config.OssecSensor.PLUGIN_ID)
    async def ignore_ossec_alert(self, _):
        if (
            self._script_installed
            and strategy.Strategy.current == self.STRATEGY
        ):
            raise Reject("CSF is running")

    @expect(MessageType.SensorIncident, plugin_id="lfd")
    async def copy_lfd_incident_to_alert(self, message):
        """protector only"""
        alert = MessageType.SensorAlert.from_incident(message)
        await self._sink.process_message(alert)

    @expect(MessageType.SensorAlert, plugin_id="lfd")
    async def unblock_in_csf(self, message):
        logger.info(
            "Unblocking %s in CSF before adding to graylist",
            message["attackers_ip"],
        )
        await csf.unblock(message["attackers_ip"])