p15670423
/
monkey
1
0
Fork 5
Code Issues Pull Requests 3 Projects Releases Wiki Activity

Agent: Use OperatingSystems in CachingAgentRepository

This commit is contained in:
Ilija Lazoroski 2022-06-24 10:19:47 +02:00
parent f9445a2c76
commit d59dd81f43
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
monkey/infection_monkey/exploit/caching_agent_repository.py
View File

@ -5,6 +5,7 @@ from typing import Mapping
import requests import requests
from common import OperatingSystems
from common.common_consts.timeouts import MEDIUM_REQUEST_TIMEOUT from common.common_consts.timeouts import MEDIUM_REQUEST_TIMEOUT
from . import IAgentRepository from . import IAgentRepository
@ -22,13 +23,13 @@ class CachingAgentRepository(IAgentRepository):
self._proxies = proxies self._proxies = proxies
self._lock = threading.Lock() 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 # 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 # _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 # 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. # that only one request per OS will be sent to the island.
with self._lock: 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) @lru_cache(maxsize=None)
def _download_binary_from_island(self, os: str) -> bytes: def _download_binary_from_island(self, os: str) -> bytes: