From d59dd81f432edf8ba4aa007fc01d3af25194c182 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 24 Jun 2022 10:19:47 +0200 Subject: [PATCH] Agent: Use OperatingSystems in CachingAgentRepository --- monkey/infection_monkey/exploit/caching_agent_repository.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/monkey/infection_monkey/exploit/caching_agent_repository.py b/monkey/infection_monkey/exploit/caching_agent_repository.py index 0f86bbd9d..c23aca287 100644 --- a/monkey/infection_monkey/exploit/caching_agent_repository.py +++ b/monkey/infection_monkey/exploit/caching_agent_repository.py @@ -5,6 +5,7 @@ from typing import Mapping import requests +from common import OperatingSystems from common.common_consts.timeouts import MEDIUM_REQUEST_TIMEOUT from . import IAgentRepository @@ -22,13 +23,13 @@ class CachingAgentRepository(IAgentRepository): self._proxies = proxies self._lock = threading.Lock() - def get_agent_binary(self, os: str, architecture: str = None) -> io.BytesIO: + def get_agent_binary(self, os: OperatingSystems, architecture: str = None) -> io.BytesIO: # If multiple calls to get_agent_binary() are made simultaneously before the result of # _download_binary_from_island() is cached, then multiple requests will be sent to the # island. Add a mutex in front of the call to _download_agent_binary_from_island() so # that only one request per OS will be sent to the island. with self._lock: - return io.BytesIO(self._download_binary_from_island(os)) + return io.BytesIO(self._download_binary_from_island(os.value)) @lru_cache(maxsize=None) def _download_binary_from_island(self, os: str) -> bytes: