From c888c84e64288101cdcb6da971a6dacf95b762d8 Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Mon, 28 Feb 2022 15:28:48 -0500
Subject: [PATCH] Agent: Add CachingAgentRepository

---
 monkey/infection_monkey/exploit/__init__.py   |  1 +
 .../exploit/caching_agent_repository.py       | 37 +++++++++++++++++++
 2 files changed, 38 insertions(+)
 create mode 100644 monkey/infection_monkey/exploit/caching_agent_repository.py

diff --git a/monkey/infection_monkey/exploit/__init__.py b/monkey/infection_monkey/exploit/__init__.py
index 105d7947e..7e5733502 100644
--- a/monkey/infection_monkey/exploit/__init__.py
+++ b/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
diff --git a/monkey/infection_monkey/exploit/caching_agent_repository.py b/monkey/infection_monkey/exploit/caching_agent_repository.py
new file mode 100644
index 000000000..2e52990b9
--- /dev/null
+++ b/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