From 088e020fee14b9184fae1780f941e80b920d4b16 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 20 Sep 2022 18:31:57 +0530 Subject: [PATCH] Agent: Catch IslandAPIError and raise RetrievalError in CachingAgentBinaryRepository._download_binary_from_island() --- monkey/infection_monkey/exploit/__init__.py | 2 +- .../exploit/caching_agent_binary_repository.py | 12 +++++------- 2 files changed, 6 insertions(+), 8 deletions(-) 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)