Mini Shell

import asyncio
from logging import getLogger

from defence360agent import utils
from defence360agent.contracts.messages import MessageType
from defence360agent.contracts.plugins import MessageSink, expect, Scope
from defence360agent.files import FILES_DIR
from defence360agent.subsys.panels.base import (
    ModsecVendorsError,
    PanelException,
)
from defence360agent.subsys.persistent_state import register_lock_file
from defence360agent.utils import recurring_check
from defence360agent.utils.check_lock import check_lock
from defence360agent.utils.common import DAY
from im360.contracts.config import Modsec
from im360.subsys.panels.base import (
    ModsecImunifyVendorNotInstalled,
    ModsecNotInstalledVendors,
    is_modsec_locked,
    use_modsec_lock,
)
from im360.subsys.panels.hosting_panel import HostingPanel
from im360.subsys.panels.update_hooks import (
    update_account_compromise_prevention_rule_state,
    update_vendors,
)

logger = getLogger(__name__)

LOCK_FILE = register_lock_file("modsec-rules-check", Scope.IM360)


class ModsecRulesetChecker(MessageSink):
    VERSION_FILE = FILES_DIR / "modsec/v2/VERSION"

    def __init__(self):
        self._loop = None
        self._task = None
        self._install_rules_task = None
        self.panel = None

    async def create_sink(self, loop):
        self.panel = HostingPanel()
        self._loop = loop
        self._task = self._loop.create_task(self.reinstall_vendor_if_needed())
        self._install_rules_task = self._loop.create_task(self.update_rules())

    async def shutdown(self):
        self._task.cancel()
        self._install_rules_task.cancel()
        await self._task
        await self._install_rules_task

    @use_modsec_lock
    async def reinstall_vendor_if_needed(self):
        panel = HostingPanel(check_for_changes=True)

        token = panel.installing_settings_var.set(True)
        try:
            try:
                # We can't do anything if Panel doesn't support modsec
                # or generate specific panel error
                installed_im360_vendor = await panel.get_i360_vendor_name()
            except (
                ModsecNotInstalledVendors,
                ModsecImunifyVendorNotInstalled,
            ):
                # this means that we can't find any modsec vendors or
                # there is no im360 vendor
                installed_im360_vendor = None
            if (
                installed_im360_vendor is None
                or Modsec.RULESET.lower() not in installed_im360_vendor
            ):
                logger.info(
                    "Installed i360 vendor %s does not match expected type of "
                    "ruleset: %s\n Trying to reinstall modsec ruleset",
                    str(installed_im360_vendor),
                    str(Modsec.RULESET),
                )
                await panel.apply_modsec_files_update()
        except asyncio.CancelledError:
            raise
        except Exception as e:
            logger.error(
                "Something went wrong during reinstalling modsec ruleset: %s",
                e,
            )
        finally:
            panel.installing_settings_var.reset(token)

    @recurring_check(
        check_lock,
        check_period_first=True,
        check_lock_period=DAY / 2,
        lock_file=LOCK_FILE,
    )
    async def update_rules(self):
        """
        Reinstall ModSec rules if installed version is not the same as in files
        """
        await self._update_rules()

    async def _update_rules(self):
        if is_modsec_locked():
            # Already being updated via imunify files
            return

        try:
            installed_vendor_version = (
                await self.panel.get_i360_vendor_version()
            )
        except (ModsecVendorsError, PanelException):
            return
        available_vendor_version = self.get_vendor_version_from_files()
        if installed_vendor_version and available_vendor_version:
            if installed_vendor_version != available_vendor_version:
                logger.info(
                    "Reinstalling ModSec rules. Installed: %s. Available: %s",
                    installed_vendor_version,
                    available_vendor_version,
                )
                await update_vendors(None, is_updated=True)

    @classmethod
    def get_vendor_version_from_files(cls) -> str:
        try:
            with open(cls.VERSION_FILE) as f:
                return f.read().strip()
        except FileNotFoundError:
            return ""

    @expect(MessageType.ConfigUpdate)
    @utils.log_error_and_ignore()
    async def on_config_update(self, _):
        await self.reinstall_vendor_if_needed()

    @expect(MessageType.ConfigUpdate)
    async def check_cms_account_compromise_prevention(self, _):
        """
        Update *WP_REDIRECT_CONF* to apply changes
        to *cms_account_compromise_prevention* config setting.
        """
        await update_account_compromise_prevention_rule_state()