Mini Shell

#!/opt/imunify360/venv/bin/python3
import asyncio
import os
import itertools
from logging import getLogger
from pathlib import Path

from peewee import OperationalError
from playhouse.sqlite_ext import SqliteExtDatabase

from im360.contracts.config import IPSET_LISTS_PATH
from imav.malwarelib.subsys.malware import (
    subscribe_to_malware_action,
    HackerTrapHitsSaver,
)
from imav.internals.lazy_load import AVSource
from imav import server
from defence360agent.subsys import systemd_notifier
from defence360agent.api import health
from defence360agent.contracts.config import Model
from defence360agent.contracts.license import LicenseCLN
from defence360agent.contracts.plugins import BasePlugin
from defence360agent.internals.cln import subscribe_to_license_changes
from defence360agent.internals.iaid import IndependentAgentIDAPI
from defence360agent.internals.lazy_load import CoreSource
from defence360agent.model import instance
from defence360agent.router import Router
from defence360agent.application import app
from imav.run import AV_PLUGINS_PACKAGES
from defence360agent.utils import importer, Scope, fail_agent_service

from im360 import rpc_handlers
from im360.application.settings import configure
from im360.internals.core import ip_versions
from im360.internals.lazy_load import IM360Source
from im360.subsys import features

IM360_PLUGINS_PACKAGES = ("im360.plugins",)
logger = getLogger(__name__)


def _set_correct_fgw_plugin(plugin):
    #  TEMPORARY LOGIC FOR MIGRATE TO FGW PLUGIN
    #  HANDLE CASE WHEN HTTP VERSION OF SendToServerClient MUST BE USER
    if plugin.__name__ == "SendToServerFGW":
        flag_file_path = os.getenv(
            "IM360_FGW_FOLDER_PATH", "/var/imunify360/.dont.use.fgw.flag"
        )
        if Path(flag_file_path).exists():
            from defence360agent.plugins.client import SendToServer

            return SendToServer
    return plugin


def get_plugins() -> set:
    """Return plugins in unspecified order."""
    importer.load_packages(
        CoreSource.MESSAGES + AVSource.MESSAGES + IM360Source.MESSAGES
    )
    importer.load_packages(AV_PLUGINS_PACKAGES + IM360_PLUGINS_PACKAGES)
    # use lexicographical order (but don't rely on it in code)
    return sorted(
        [
            _set_correct_fgw_plugin(plugin)
            for plugin in BasePlugin.get_active_plugins()
            if plugin.SCOPE not in (Scope.AV, Scope.IM360_RESIDENT)
        ],
        key=lambda item: f"{item.__module__}.{item.__name__}",
    )


async def update_health_sensor():
    if LicenseCLN.is_valid():
        health.sensor.registered()
    else:
        health.sensor.unregistered()


async def attach_db(
    db: SqliteExtDatabase, db_path: str, schema_name: str
) -> None:
    max_attempts = 5
    for attempt in itertools.count(1):
        try:
            db.execute_sql("ATTACH ? AS ?", (db_path, schema_name))
            return
        except Exception as e:
            if isinstance(e, OperationalError) and attempt < max_attempts:
                await asyncio.sleep(5)
                continue

            logger.error(
                "Error attaching to database",
                extra={
                    "db_path": db_path,
                    "schema_name": schema_name,
                    "attempt": attempt,
                },
            )
            raise


async def setup_databases():
    instance.db.execute_sql(
        "ATTACH ? AS ?", (Model.PROACTIVE_PATH, "proactive")
    )
    # DB is created and migration applied in resident part
    # so we need to retry until it's ready
    # 5 attempts with 5 seconds sleep between them max 20 seconds, concurrently for each database
    async with asyncio.TaskGroup() as task_group:
        task_group.create_task(
            attach_db(instance.db, Model.RESIDENT_PATH, "resident")
        )
        task_group.create_task(
            attach_db(instance.db, IPSET_LISTS_PATH, "ipsetlists")
        )

    router = Router(
        instance.db,
        migrations_dirs=app.MIGRATIONS_DIRS,
        logger=logger,
    )
    # Do not run non-resident part if there are new migration in dirs
    # This behavior might be reproduced on update
    # non-resident and resident agents both would be restarted
    # as soon as update process is finished
    if router.diff:
        logger.warning("Tried to start while migrations are not applied")
        fail_agent_service()


async def init_actions():
    # Any uncaught exceptions here prevents the agent from
    # starting. Non-critical (agent can continue) functionality should
    # either be moved to a plugin or caught&log *all* its exceptions.

    # Also nothing should block here indefinitely (any operation that
    # may block should have a timeout).

    ip_versions.init()
    subscribe_to_license_changes(features.update_repos)
    subscribe_to_license_changes(features.update_im_email)
    subscribe_to_license_changes(update_health_sensor)

    await setup_databases()

    subscribe_to_malware_action("delete", HackerTrapHitsSaver.add_hit)
    subscribe_to_malware_action("cleanup", HackerTrapHitsSaver.add_hit)
    subscribe_to_license_changes(IndependentAgentIDAPI.reactivate)
    IndependentAgentIDAPI.add_initial_task()
    rpc_handlers.init()


def run():
    systemd_notifier.notify(systemd_notifier.AgentState.READY)
    configure()
    plugins = get_plugins()
    server.start(plugins, init_actions)


if __name__ == "__main__":
    run()
    logger.info("agent stopped")