Island: Refactor LocalMonkerRunService to use AgentBinaryRepository

* Register the data_dir in the DI container
* Construct LocalRun resource with the LocalMonkeyRunService

monkey/monkey_island/cc/resources/local_run.py

@@ -12,6 +12,9 @@ from monkey_island.cc.services.run_local_monkey import LocalMonkeyRunService
 class LocalRun(AbstractResource):
     urls = ["/api/local-monkey"]
 
+    def __init__(self, local_monkey_run_service: LocalMonkeyRunService):
+        self._local_monkey_run_service = local_monkey_run_service
+
     # API Spec: Both of these methods should be separated to their own resources
 
     # API Spec: This should be a REST endpoint, /api/monkeys or something
@@ -30,7 +33,7 @@ class LocalRun(AbstractResource):
     def post(self):
         body = json.loads(request.data)
         if body.get("action") == "run":
-            local_run = LocalMonkeyRunService.run_local_monkey()
+            local_run = self._local_monkey_run_service.run_local_monkey()
             # API Spec: Feels weird to return "error_text" even when "is_running" is True
             return jsonify(is_running=local_run[0], error_text=local_run[1])
 

monkey/monkey_island/cc/services/initialize.py

@@ -26,6 +26,11 @@ AGENT_BINARIES_PATH = Path(MONKEY_ISLAND_ABS_PATH) / "cc" / "binaries"
 
 def initialize_services(data_dir: Path) -> DIContainer:
     container = DIContainer()
+
+    # TODO: everything that is build with DI and expects Path in the constructor
+    # will use the same data_dir. Come up with a better way to inject
+    # the data_dir in the things that needed
+    container.register_instance(Path, data_dir)
     container.register_instance(AWSInstance, AWSInstance())
 
     container.register_instance(
@@ -33,10 +38,10 @@ def initialize_services(data_dir: Path) -> DIContainer:
     )
     container.register_instance(AWSService, container.resolve(AWSService))
     container.register_instance(IAgentBinaryRepository, _build_agent_binary_repository())
+    container.register_instance(LocalMonkeyRunService, container.resolve(LocalMonkeyRunService))
 
     # This is temporary until we get DI all worked out.
     PostBreachFilesService.initialize(container.resolve(IFileRepository))
-    LocalMonkeyRunService.initialize(data_dir)
     AuthenticationService.initialize(data_dir, JsonFileUserDatastore(data_dir))
     ReportService.initialize(container.resolve(AWSService))
 

monkey/monkey_island/cc/services/run_local_monkey.py

@@ -1,49 +1,58 @@
 import logging
-import os
 import platform
 import stat
 import subprocess
 from pathlib import Path
-from shutil import copyfile
+from shutil import copyfileobj
 
-from monkey_island.cc.server_utils.consts import ISLAND_PORT, MONKEY_ISLAND_ABS_PATH
+from monkey_island.cc.repository import AgentRetrievalError, IAgentBinaryRepository
+from monkey_island.cc.server_utils.consts import ISLAND_PORT
 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
+AGENT_NAMES = {"linux": "monkey-linux-64", "windows": "monkey-windows-64.exe"}
 
 
 class LocalMonkeyRunService:
-    DATA_DIR = None
+    def __init__(self, data_dir: Path, agent_binary_repository: IAgentBinaryRepository):
+        self._data_dir = data_dir
+        self._agent_binary_repository = agent_binary_repository
 
-    # TODO: A number of these services should be instance objects instead of
-    # static/singleton hybrids. At the moment, this requires invasive refactoring that's
-    # not a priority.
-    @classmethod
-    def initialize(cls, data_dir: Path):
-        cls.DATA_DIR = data_dir
-
-    @staticmethod
-    def run_local_monkey():
+    def run_local_monkey(self):
         # get the monkey executable suitable to run on the server
+        operating_system = platform.system().lower()
         try:
-            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)
+            agents = {
+                "linux": self._agent_binary_repository.get_linux_binary,
+                "windows": self._agent_binary_repository.get_windows_binary,
+            }
 
-        dest_path = LocalMonkeyRunService.DATA_DIR / src_path.name
+            agent_binary = agents[platform.system().lower()]()
+        except AgentRetrievalError as err:
+            logger.error(
+                f"No Agent can be retrieved for the specified operating system"
+                f'"{operating_system}"'
+            )
+            return False, str(err)
+        except KeyError as err:
+            logger.error(
+                f"No Agents are available for unsupported operating system" f'"{operating_system}"'
+            )
+            return False, str(err)
+        except Exception as err:
+            logger.error(f"Error running agent from island: {err}")
+            return False, str(err)
+
+        dest_path = self._data_dir / AGENT_NAMES[operating_system]
 
         # copy the executable to temp path (don't run the monkey from its current location as it may
         # delete itself)
         try:
-            copyfile(src_path, dest_path)
-            os.chmod(dest_path, stat.S_IRWXU | stat.S_IRWXG)
+            with open(dest_path, "wb") as dest_agent:
+                copyfileobj(agent_binary, dest_agent)
+
+            dest_path.chmod(stat.S_IRWXU | stat.S_IRWXG)
         except Exception as exc:
             logger.error("Copy file failed", exc_info=True)
             return False, "Copy file failed: %s" % exc
@@ -54,28 +63,9 @@ class LocalMonkeyRunService:
             port = ISLAND_PORT
 
             args = [str(dest_path), "m0nk3y", "-s", f"{ip}:{port}"]
-            subprocess.Popen(args, cwd=LocalMonkeyRunService.DATA_DIR)
+            subprocess.Popen(args, cwd=self._data_dir)
         except Exception as exc:
             logger.error("popen failed", exc_info=True)
             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