import datetime import pwd import time from defence360agent.contracts.messages import MessageType from defence360agent.contracts.config import ( int_from_envvar, ) from defence360agent.contracts.plugins import ( MessageSink, MessageSource, ) from defence360agent.utils import ( recurring_check, split_for_chunk, Scope, ) from defence360agent.subsys.panels.cpanel import cPanel from defence360agent.subsys.panels.hosting_panel import HostingPanel from defence360agent.contracts.myimunify_id import get_myimunify_users 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): hp = HostingPanel() if hp.NAME != cPanel.NAME: return domains = [] myimunify_users = await get_myimunify_users() for user in myimunify_users: username = user["username"] user_pwd = pwd.getpwnam(username) for domain_data in 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 = MessageType.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)