import datetime import logging import pwd import time from typing import AsyncIterator from defence360agent.contracts.config import ( int_from_envvar, ) from defence360agent.contracts.messages import DomainList from defence360agent.contracts.myimunify_id import get_myimunify_users from defence360agent.contracts.plugins import ( MessageSink, MessageSource, ) from defence360agent.subsys.panels.hosting_panel import HostingPanel from defence360agent.utils import ( Scope, recurring_check, split_for_chunk, ) logger = logging.getLogger(__name__) class SendDomainList(MessageSink, MessageSource): SCOPE = Scope.AV_IM360 def __init__(self, period=None): self._task = None if period: self._period = period else: self._period = int_from_envvar( "IMUNIFY360_SEND_DOMAIN_PERIOD", int(datetime.timedelta(days=1).total_seconds()), ) async def create_sink(self, loop): """MessageSink method""" async def create_source(self, loop, sink): self._loop = loop self._sink = sink self._task = self._loop.create_task( recurring_check(self._period)(self._send_domain_list)() ) async def shutdown(self): if self._task is not None: self._task, t = None, self._task t.cancel() await t def _panel_domain_type_to_imunify(self, panel_type: str) -> str: return { "main": "primary", "parked": "alias", }.get(panel_type, panel_type) async def _create_domain_list_msg(self) -> AsyncIterator[DomainList]: hp = HostingPanel() logger.info("HostingsPanel: %s", hp.NAME) domains = [] myimunify_users = await get_myimunify_users() for user in myimunify_users: username = user["username"] user_pwd = pwd.getpwnam(username) for domain_data in await hp.get_user_domains_details(username): domains.append( { "username": username, "docroot": domain_data.docroot, "name": domain_data.domain, "securesite_user_id": user["myimunify_id"], "uid": user_pwd.pw_uid, "type": self._panel_domain_type_to_imunify( domain_data.type ), } ) timestamp = time.time() for chunk in split_for_chunk(domains, chunk_size=3000): msg = DomainList() msg["timestamp"] = timestamp msg["domains"] = chunk yield msg async def _send_domain_list(self): async for msg in self._create_domain_list_msg(): await self._sink.process_message(msg)