diff --git a/monkey/infection_monkey/exploit/__init__.py b/monkey/infection_monkey/exploit/__init__.py index 195e880ad..22d1fc865 100644 --- a/monkey/infection_monkey/exploit/__init__.py +++ b/monkey/infection_monkey/exploit/__init__.py @@ -1,3 +1,3 @@ -from .i_agent_binary_repository import IAgentBinaryRepository +from .i_agent_binary_repository import IAgentBinaryRepository, RetrievalError from .caching_agent_binary_repository import CachingAgentBinaryRepository from .exploiter_wrapper import ExploiterWrapper diff --git a/monkey/infection_monkey/exploit/caching_agent_binary_repository.py b/monkey/infection_monkey/exploit/caching_agent_binary_repository.py index 46e752470..8af4b7a4d 100644 --- a/monkey/infection_monkey/exploit/caching_agent_binary_repository.py +++ b/monkey/infection_monkey/exploit/caching_agent_binary_repository.py @@ -2,12 +2,11 @@ import io import logging import threading from functools import lru_cache -from typing import Optional from common import OperatingSystem -from infection_monkey.island_api_client import IIslandAPIClient +from infection_monkey.island_api_client import IIslandAPIClient, IslandAPIError -from . import IAgentBinaryRepository +from . import IAgentBinaryRepository, RetrievalError logger = logging.getLogger(__name__) @@ -34,10 +33,9 @@ class CachingAgentBinaryRepository(IAgentBinaryRepository): return io.BytesIO(self._download_binary_from_island(operating_system)) @lru_cache(maxsize=None) - def _download_binary_from_island(self, operating_system: OperatingSystem) -> Optional[bytes]: + def _download_binary_from_island(self, operating_system: OperatingSystem) -> bytes: os_name = operating_system.value try: return self._island_api_client.get_agent_binary(os_name) - except Exception as exc: - logger.warning(f"Error connecting to control server: {exc}") - return None + except IslandAPIError as err: + raise RetrievalError(err)