import asyncio import time from contextlib import suppress from datetime import timedelta from logging import getLogger from typing import Union from defence360agent.contracts.messages import MessageType from defence360agent.contracts.plugins import MessageSource from defence360agent.subsys.backup_systems import ( get_current_backend, get_last_backup_timestamp, ) from defence360agent.subsys.persistent_state import load_state, save_state from defence360agent.utils import Scope, recurring_check logger = getLogger(__name__) SEND_INTERVAL = int(timedelta(hours=24).total_seconds()) RECURRING_CHECK_INTERVAL = 5 class BackupInfoSender(MessageSource): """Send user backup statistics to CH periodically""" SCOPE = Scope.IM360 async def create_source(self, loop, sink): self._loop = loop self._sink = sink self._send_event = asyncio.Event() self._last_send_timestamp = self.load_last_send_timestamp() self._check_task = self._loop.create_task( self._recurring_check_data_to_send() ) self._send_stat_task = self._loop.create_task( self._recurring_send_stat() ) async def shutdown(self): for task in [self._check_task, self._send_stat_task]: task.cancel() with suppress(asyncio.CancelledError): await task self.save_last_send_timestamp() @staticmethod def is_valid_timestamp(timestamp: Union[int, float]) -> bool: return isinstance(timestamp, (int, float)) and timestamp > 0 def save_last_send_timestamp(self, ts: Union[int, float] = None): timestamp = self._last_send_timestamp if ts is None else ts if not self.is_valid_timestamp(timestamp): logger.warning("Invalid timestamp: %s", timestamp) return save_state("BackupInfoSender", {"last_send_timestamp": timestamp}) def load_last_send_timestamp(self): timestamp = load_state("BackupInfoSender").get("last_send_timestamp") if not self.is_valid_timestamp(timestamp): logger.warning("Invalid timestamp loaded, resetting to 0") timestamp = 0 return timestamp @recurring_check(RECURRING_CHECK_INTERVAL) async def _recurring_check_data_to_send(self): if time.time() - self._last_send_timestamp >= SEND_INTERVAL: self._send_event.set() @recurring_check(0) async def _recurring_send_stat(self): await self._send_event.wait() try: await self._send_server_config() except Exception as e: logger.exception("Failed to collect backup info: %s", e) finally: # Ensure backup info is not sent too frequently, even after an error self._last_send_timestamp = time.time() self._send_event.clear() async def _send_server_config(self): confg_msg = MessageType.BackupInfo( backup_provider_type=get_current_backend(), last_backup_timestamp=await get_last_backup_timestamp(), ) await self._sink.process_message(confg_msg)