OwlCyberSecurity - MANAGER

Edit File: service_manager.py

"""
This program is free software: you can redistribute it and/or modify it under
the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License,
or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. 
See the GNU General Public License for more details.

You should have received a copy of the GNU General Public License
 along with this program.  If not, see <https://www.gnu.org/licenses/>.

This software is also available under ImunifyAV commercial license,
see <https://www.imunify360.com/legal/eula>

Services manager plugin.

It enables/disables various service based on an imunify config change.

"""
import asyncio
import logging

from defence360agent import utils
from defence360agent.contracts import messages, plugins
from defence360agent.subsys import svcctl
from imav.contracts.config import Patchman
from imav.patchman.constants import PATCHMAN_SERVICE_NAME
from imav.patchman.state import is_installed as is_patchman_installed

__all__ = ["ServiceManager"]

logger = logging.getLogger(__name__)

class ServiceManager(plugins.MessageSink):
    """Service manager plugin: stop/start services based on config changes."""

SCOPE = utils.Scope.AV

def __init__(self, *, unitctl=None):
        self._lock = asyncio.Lock()
        self._services = [self._ensure_consistent_patchman_state]
        self._units = {
            "patchman": unitctl or svcctl.adaptor(PATCHMAN_SERVICE_NAME),
        }

async def create_sink(self, loop):
        # on startup ConfigUpdate message is sent to all plugins
        pass

async def _ensure_consistent_services_state(self):
        for service in self._services:
            await service()

@plugins.expect(messages.MessageType.ConfigUpdate)
    async def on_config_update(
        self, message_ignored: messages.MessageType.ConfigUpdate
    ):
        async with self._lock:  # handle concurrent config updates
            await self._ensure_consistent_services_state()

@plugins.expect(
        messages.MessageType.EnsureServiceState, service=PATCHMAN_SERVICE_NAME
    )
    async def ensure_patchman_state(
        self, message: messages.MessageType.EnsureServiceState
    ):
        async with self._lock:
            await self._ensure_consistent_patchman_state()

@utils.log_error_and_ignore()
    async def __ensure_service_status(
        self, unitctl, service_name, should_be_running, reload=False
    ):
        is_running = await unitctl.is_active()
        if is_running is not should_be_running:
            if should_be_running:
                logger.info(
                    "%s is enabled in the config but it is not"
                    " running. Enabling it...",
                    service_name,
                )
                # enable on boot & start it right now
                await unitctl.enable(now=True)
                logger.info("Enabled %s", service_name)
            else:
                logger.info(
                    "%s is not enabled in the config but it is"
                    " running. Disabling it...",
                    service_name,
                )
                # disable on boot & stop it right now
                await unitctl.disable(now=True)
                logger.info("Disabled %s", service_name)
        else:
            if is_running and reload:
                await unitctl.reload()
                logger.info(
                    "Reloading %s after config update...", service_name
                )  # noqa: E501

async def _ensure_consistent_patchman_state(self) -> None:
        if is_patchman_installed():
            await self.__ensure_service_status(
                self._units["patchman"],
                PATCHMAN_SERVICE_NAME,
                Patchman.ENABLED,
            )