Agent: Add CachingAgentRepository

monkey/infection_monkey/exploit/__init__.py

@@ -1,2 +1,3 @@
 from .i_agent_repository import IAgentRepository
+from .caching_agent_repository import CachingAgentRepository
 from .exploiter_wrapper import ExploiterWrapper

monkey/infection_monkey/exploit/caching_agent_repository.py

@@ -0,0 +1,37 @@
+import io
+from functools import lru_cache
+from typing import Mapping
+
+import requests
+
+from common.common_consts.timeouts import MEDIUM_REQUEST_TIMEOUT
+
+from . import IAgentRepository
+
+
+class CachingAgentRepository(IAgentRepository):
+    """
+    CachingAgentRepository implements the IAgentRepository interface and downloads the requested
+    agent binary from the island on request. The agent binary is cached so that only one request is
+    actually sent to the island for each requested binary.
+    """
+
+    def __init__(self, island_url: str, proxies: Mapping[str, str]):
+        self._island_url = island_url
+        self._proxies = proxies
+
+    def get_agent_binary(self, os: str, _: str = None) -> io.BytesIO:
+        return io.BytesIO(self._download_binary_from_island(os))
+
+    @lru_cache(maxsize=None)
+    def _download_binary_from_island(self, os: str) -> bytes:
+        response = requests.get(  # noqa: DUO123
+            f"{self._island_url}/api/monkey/download/{os}",
+            verify=False,
+            proxies=self._proxies,
+            timeout=MEDIUM_REQUEST_TIMEOUT,
+        )
+
+        response.raise_for_status()
+
+        return response.content