Island: Refactor agent-binaries resource

monkey/monkey_island/cc/resources/agent_binaries.py

@@ -1,48 +1,40 @@
 import logging
-from pathlib import Path
 
-from flask import make_response, send_from_directory
+from flask import make_response, send_file
 
+from monkey_island.cc.repository import AgentRetrievalError, IAgentBinaryRepository
 from monkey_island.cc.resources.AbstractResource import AbstractResource
-from monkey_island.cc.server_utils.consts import MONKEY_ISLAND_ABS_PATH
 
 logger = logging.getLogger(__name__)
 
-AGENTS = {
-    "linux": "monkey-linux-64",
-    "windows": "monkey-windows-64.exe",
-}
-
-
-class UnsupportedOSError(Exception):
-    pass
-
 
 class AgentBinaries(AbstractResource):
     urls = ["/api/agent-binaries/<string:os>"]
 
+    def __init__(self, agent_binary_repository: IAgentBinaryRepository):
+        self._agent_binary_repository = agent_binary_repository
+
     # Used by monkey. can't secure.
     def get(self, os):
+        """
+        Gets the agent binary based on the OS
+
+        :param os: Operating systems. Supported OS are: 'linux' and 'windows'
+        :return: file-like object with a filename same as the OS
+        """
         try:
-            path = get_agent_executable_path(os)
-            return send_from_directory(path.parent, path.name)
-        except UnsupportedOSError as ex:
-            logger.error(ex)
-            return make_response({"error": str(ex)}, 404)
+            agent_binaries = {
+                "linux": self._agent_binary_repository.get_linux_binary,
+                "windows": self._agent_binary_repository.get_windows_binary,
+            }
 
+            file = agent_binaries[os]()
 
-def get_agent_executable_path(os: str) -> Path:
-    try:
-        agent_path = get_executable_full_path(AGENTS[os])
-        logger.debug(f'Local path for {os} executable is "{agent_path}"')
-        if not agent_path.is_file():
-            logger.error(f"File {agent_path} not found")
-
-        return agent_path
-    except KeyError:
-        logger.warning(f"No monkey executables could be found for the host os: {os}")
-        raise UnsupportedOSError(f'No Agents are available for unsupported operating system "{os}"')
-
-
-def get_executable_full_path(executable_filename: str) -> Path:
-    return Path(MONKEY_ISLAND_ABS_PATH) / "cc" / "binaries" / executable_filename
+            return send_file(file, mimetype="application/octet-stream")
+        except AgentRetrievalError as err:
+            logger.error(err)
+            return make_response({"error": str(err)}, 500)
+        except KeyError as err:
+            error_msg = f'No Agents are available for unsupported operating system "{os}": {err}'
+            logger.error(error_msg)
+            return make_response({"error": error_msg}, 404)

monkey/monkey_island/cc/services/run_local_monkey.py

@@ -6,12 +6,17 @@ import subprocess
 from pathlib import Path
 from shutil import copyfile
 
-from monkey_island.cc.resources.agent_binaries import get_agent_executable_path
-from monkey_island.cc.server_utils.consts import ISLAND_PORT
+from monkey_island.cc.server_utils.consts import ISLAND_PORT, MONKEY_ISLAND_ABS_PATH
 from monkey_island.cc.services.utils.network_utils import local_ip_addresses
 
 logger = logging.getLogger(__name__)
 
+AGENTS = {"linux": "monkey-linux-64", "windows": "monkey-windows-64.exe"}
+
+
+class UnsupportedOSError(Exception):
+    pass
+
 
 class LocalMonkeyRunService:
     DATA_DIR = None
@@ -27,7 +32,7 @@ class LocalMonkeyRunService:
     def run_local_monkey():
         # get the monkey executable suitable to run on the server
         try:
-            src_path = get_agent_executable_path(platform.system().lower())
+            src_path = LocalMonkeyRunService._get_agent_executable_path(platform.system().lower())
         except Exception as ex:
             logger.error(f"Error running agent from island: {ex}")
             return False, str(ex)
@@ -55,3 +60,22 @@ class LocalMonkeyRunService:
             return False, "popen failed: %s" % exc
 
         return True, ""
+
+    @staticmethod
+    def _get_agent_executable_path(os: str) -> Path:
+        try:
+            agent_path = LocalMonkeyRunService._get_executable_full_path(AGENTS[os])
+            logger.debug(f'Local path for {os} executable is "{agent_path}"')
+            if not agent_path.is_file():
+                logger.error(f"File {agent_path} not found")
+
+            return agent_path
+        except KeyError:
+            logger.warning(f"No monkey executable could be found for the host os: {os}")
+            raise UnsupportedOSError(
+                f'No Agents are available for unsupported operating system "{os}"'
+            )
+
+    @staticmethod
+    def _get_executable_full_path(executable_filename: str) -> Path:
+        return Path(MONKEY_ISLAND_ABS_PATH) / "cc" / "binaries" / executable_filename