Agent: Catch IslandAPIError and raise RetrievalError in CachingAgentBinaryRepository._download_binary_from_island()

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

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)