Island: Refactor monkey download to take OS and return agent file

monkey/monkey_island/cc/app.py

@@ -134,8 +134,7 @@ def init_api_resources(api):
     api.add_resource(ConfigurationImport, "/api/configuration/import")
     api.add_resource(
         MonkeyDownload,
-        "/api/monkey/download",
-        "/api/monkey/download/<string:path>",
+        "/api/monkey/download/<string:host_os>",
     )
     api.add_resource(NetMap, "/api/netmap")
     api.add_resource(Edge, "/api/netmap/edge")

monkey/monkey_island/cc/resources/monkey_download.py

@@ -1,63 +1,34 @@
 import hashlib
-import json
 import logging
-import os
+from pathlib import Path
 
 import flask_restful
-from flask import request, send_from_directory
+from flask import make_response, send_from_directory
 
 from monkey_island.cc.server_utils.consts import MONKEY_ISLAND_ABS_PATH
 
 logger = logging.getLogger(__name__)
 
-MONKEY_DOWNLOADS = [
-    {
-        "type": "linux",
-        "filename": "monkey-linux-64",
-    },
-    {
-        "type": "windows",
-        "filename": "monkey-windows-64.exe",
-    },
-]
+AGENTS = {
+    "linux": "monkey-linux-64",
+    "windows": "monkey-windows-64.exe",
+}
 
 
-def get_monkey_executable(host_os):
-    for download in MONKEY_DOWNLOADS:
-        if host_os == download.get("type"):
-            logger.info(f"Monkey exec found for os: {host_os}")
-            return download
-    logger.warning(f"No monkey executables could be found for the host os: {host_os}")
-    return None
+class UnsupportedOSError(Exception):
+    pass
 
 
 class MonkeyDownload(flask_restful.Resource):
 
     # Used by monkey. can't secure.
-    def get(self, path):
-        return send_from_directory(os.path.join(MONKEY_ISLAND_ABS_PATH, "cc", "binaries"), path)
-
-    # Used by monkey. can't secure.
-    def post(self):
-        host_json = json.loads(request.data)
-        host_os = host_json.get("os")
-        if host_os:
-            result = get_monkey_executable(host_os.get("type"))
-
-            if result:
-                # change resulting from new base path
-                executable_filename = result["filename"]
-                real_path = MonkeyDownload.get_executable_full_path(executable_filename)
-                if os.path.isfile(real_path):
-                    result["size"] = os.path.getsize(real_path)
-                    return result
-
-        return {}
-
-    @staticmethod
-    def get_executable_full_path(executable_filename):
-        real_path = os.path.join(MONKEY_ISLAND_ABS_PATH, "cc", "binaries", executable_filename)
-        return real_path
+    def get(self, host_os):
+        try:
+            path = get_agent_executable_path(host_os)
+            return send_from_directory(path.parent, path.name)
+        except UnsupportedOSError as ex:
+            logger.error(ex)
+            return make_response({"error": str(ex)}, 404)
 
     @staticmethod
     def log_executable_hashes():
@@ -65,16 +36,30 @@ class MonkeyDownload(flask_restful.Resource):
         Logs all the hashes of the monkey executables for debugging ease (can check what Monkey
         version you have etc.).
         """
-        filenames = set([x["filename"] for x in MONKEY_DOWNLOADS])
+        filenames = set(AGENTS.values())
         for filename in filenames:
-            filepath = MonkeyDownload.get_executable_full_path(filename)
-            if os.path.isfile(filepath):
+            filepath = get_executable_full_path(filename)
+            if filepath.is_file():
                 with open(filepath, "rb") as monkey_exec_file:
                     file_contents = monkey_exec_file.read()
-                    logger.debug(
-                        "{} hashes:\nSHA-256 {}".format(
-                            filename, hashlib.sha256(file_contents).hexdigest()
-                        )
-                    )
+                    file_sha256_hash = filename, hashlib.sha256(file_contents).hexdigest()
+                    logger.debug(f"{filename} hash:\nSHA-256 {file_sha256_hash}")
             else:
-                logger.debug("No monkey executable for {}.".format(filepath))
+                logger.debug(f"No monkey executable for {filepath}")
+
+
+def get_agent_executable_path(host_os: str) -> Path:
+    try:
+        agent_path = get_executable_full_path(AGENTS[host_os])
+        logger.debug(f"Monkey exec found for os: {host_os}, {agent_path}")
+
+        return agent_path
+    except KeyError:
+        logger.warning(f"No monkey executables could be found for the host os: {host_os}")
+        raise UnsupportedOSError(
+            f'No Agents are available for unsupported operating system "{host_os}"'
+        )
+
+
+def get_executable_full_path(executable_filename: str) -> Path:
+    return Path(MONKEY_ISLAND_ABS_PATH) / "cc" / "binaries" / executable_filename

monkey/monkey_island/cc/services/run_local_monkey.py

@@ -5,8 +5,8 @@ import stat
 import subprocess
 from shutil import copyfile
 
-from monkey_island.cc.resources.monkey_download import get_monkey_executable
-from monkey_island.cc.server_utils.consts import ISLAND_PORT, MONKEY_ISLAND_ABS_PATH
+from monkey_island.cc.resources.monkey_download import get_agent_executable_path
+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__)
@@ -25,12 +25,13 @@ class LocalMonkeyRunService:
     @staticmethod
     def run_local_monkey():
         # get the monkey executable suitable to run on the server
-        result = get_monkey_executable(platform.system().lower())
-        if not result:
-            return False, "OS Type not found"
+        try:
+            src_path = get_agent_executable_path(platform.system().lower())
+        except Exception as ex:
+            logger.error(f"Error running agent from island: {ex}")
+            return False, str(ex)
 
-        src_path = os.path.join(MONKEY_ISLAND_ABS_PATH, "cc", "binaries", result["filename"])
-        dest_path = os.path.join(LocalMonkeyRunService.DATA_DIR, result["filename"])
+        dest_path = LocalMonkeyRunService.DATA_DIR / src_path.name
 
         # copy the executable to temp path (don't run the monkey from its current location as it may
         # delete itself)