p59342861
/
monkey
forked from p15670423/monkey
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

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

This commit is contained in:
Shreya Malviya 2022-09-20 18:31:57 +05:30
parent 0e9397b283
commit 088e020fee
2 changed files with 6 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
monkey/infection_monkey/exploit/__init__.py
View File

@ -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 .caching_agent_binary_repository import CachingAgentBinaryRepository
from .exploiter_wrapper import ExploiterWrapper from .exploiter_wrapper import ExploiterWrapper

12
monkey/infection_monkey/exploit/caching_agent_binary_repository.py
View File

@ -2,12 +2,11 @@ import io
import logging import logging
import threading import threading
from functools import lru_cache from functools import lru_cache
from typing import Optional
from common import OperatingSystem 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__) logger = logging.getLogger(__name__)
@ -34,10 +33,9 @@ class CachingAgentBinaryRepository(IAgentBinaryRepository):
return io.BytesIO(self._download_binary_from_island(operating_system)) return io.BytesIO(self._download_binary_from_island(operating_system))
@lru_cache(maxsize=None) @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 os_name = operating_system.value
try: try:
return self._island_api_client.get_agent_binary(os_name) return self._island_api_client.get_agent_binary(os_name)
except Exception as exc: except IslandAPIError as err:
logger.warning(f"Error connecting to control server: {exc}") raise RetrievalError(err)
return None