MAINHACK

Viewing / Editing File: dos_protector_manager.py

File is not writable. Editing disabled.
"""DoS protector manager plugin.

It enables/disables dos-protection service on an imunify360 config change.

"""
import logging
from asyncio import Lock
from os import R_OK, X_OK, access
from pathlib import Path
from typing import Callable

from defence360agent import utils
from defence360agent.contracts import plugins
from defence360agent.contracts.config import SystemConfig
from defence360agent.contracts.messages import MessageType
from defence360agent.subsys.svcctl import _apply_cmd, _SystemctlBased
from im360.contracts.config import EnhancedDOS

__all__ = ["DosProtectorManager"]

logger = logging.getLogger(__name__)


class UnitCtl(_SystemctlBased):
    def __init__(self, service_name: str):
        self.SVC_CTL_BIN = self._find_bin(self.SVC_CTL_BIN)
        super().__init__(service_name)

    @staticmethod
    def _find_bin(bin_name: str) -> str:
        if utils.is_centos6_or_cloudlinux6():
            bin_name = 'true'
            logger.debug('Enhanced DoS is not available on CentOS/CL6.')
        # /usr/bin is for CentOS 7 and the other ones;
        # /bin is for Debian-based distros;
        for dir_path in ("/usr/bin", "/bin"):
            bin_path = Path(dir_path) / bin_name
            if access(bin_path, R_OK | X_OK):
                return str(bin_path.resolve())
        raise RuntimeError(f"Unable to find {bin_name}:")

    async def enable(self, *, now: bool):
        @_apply_cmd
        def run_enable_now():
            return [
                # fmt: off
                self.SVC_CTL_BIN,
                "enable", *(['--now'] if now else []),
                self._service_name,
                # fmt: on
            ]

        await run_enable_now()

        # WARN: Ubuntu 16.04 demonstrates very special behavior of the
        # `systemcl enable --now` command - if the unit is stopped it
        # wouldn't be started. We need to handle that case.
        # TODO: Remove this case on dropping support for Ubuntu 16.04.
        osinfo = {}
        try:
            utils.OsReleaseInfo.dict_from_file(osinfo)
        except (FileNotFoundError, PermissionError):
            return
        if osinfo.get('ID', '').lower() != 'ubuntu':
            return
        if osinfo.get('VERSION_ID', '') == '16.04':
            await self.restart()

    @_apply_cmd
    def disable(self, *, now: bool):
        return [
            # fmt: off
            self.SVC_CTL_BIN,
            "disable", *(['--now'] if now else []),
            self._service_name,
            # fmt: on
        ]

    @_apply_cmd
    def reload(self):
        return [self.SVC_CTL_BIN, "reload", self._service_name]


class DosProtectorManager(plugins.MessageSink):
    def __init__(self, *args, **kwargs):
        # fmt: off
        unit_func: Callable[[str], UnitCtl] = kwargs.pop(
            "mkunitctl_fn", lambda name: UnitCtl(name),
        )
        # fmt: on
        self.__lock = Lock()
        self.__config: EnhancedDOS = kwargs.pop("config", EnhancedDOS())
        self.__unitctl = unit_func("imunify360-dos-protection.service")
        # WARN: last_opts is set to None because we need to apply an action
        # (enable|disable|reload) on starting this plugin with no condition.
        # After that we process messages, compare the config and apply an
        # action only if needed.
        self.__last_opts = None
        super().__init__(*args, **kwargs)

    async def create_sink(self, loop):
        await self._ensure_consistent_dosprotector_state()  # on start up
        self.__last_opts = self.__config.as_dict()

    @plugins.expect(MessageType.ConfigUpdate)
    async def on_config_update(self, message: MessageType.ConfigUpdate):
        if not isinstance(message["conf"], SystemConfig):
            return
        if (cur_opts := self.__config.as_dict()) != self.__last_opts:
            await self._ensure_consistent_dosprotector_state()
            self.__last_opts = cur_opts

    @utils.log_error_and_ignore()
    async def _ensure_consistent_dosprotector_state(self):
        async with self.__lock:  # handle concurrent config updates
            should_be_running = self.__config.ENABLED
            unitctl = self.__unitctl
            is_running = await unitctl.is_active()
            if is_running is not should_be_running:
                if should_be_running:
                    logger.info(
                        "DosProtector is enabled in the config but it is not running."  # noqa: E501
                        " Enabling it..."
                    )
                    # enable on boot & start it right now
                    await unitctl.enable(now=True)
                    logger.info("Enabled DosProtector")
                else:
                    logger.info(
                        "DosProtector is not enabled in the config but it is running."  # noqa: E501
                        " Disabling it..."
                    )
                    # disable on boot & stop it right now
                    await unitctl.disable(now=True)
                    logger.info("Disabled DosProtector")
            else:
                if is_running:
                    await unitctl.reload()
                    logger.info("Reloading DosProtector after config update...")       # noqa: E501
Cancel / Back
[ MAINHACK ]

Viewing / Editing File: dos_protector_manager.py

Server Info

System Features

User Info
