From e7459412e43ccaf771df389ed1a4649c6e5379c3 Mon Sep 17 00:00:00 2001
From: Ilija Lazoroski <ilija.la@live.com>
Date: Mon, 6 Jun 2022 18:31:52 +0200
Subject: [PATCH] Island: Add concrete AgentBinaryRepository

Issue #1974
---
 .../monkey_island/cc/repository/__init__.py   |  2 +-
 .../cc/repository/agent_binary_repository.py  | 29 +++++++++++++++++++
 .../repository/i_agent_binary_repository.py   |  4 +++
 .../monkey_island/cc/services/initialize.py   | 16 +++++++++-
 4 files changed, 49 insertions(+), 2 deletions(-)
 create mode 100644 monkey/monkey_island/cc/repository/agent_binary_repository.py

diff --git a/monkey/monkey_island/cc/repository/__init__.py b/monkey/monkey_island/cc/repository/__init__.py
index e3cf5c2fc..88236f771 100644
--- a/monkey/monkey_island/cc/repository/__init__.py
+++ b/monkey/monkey_island/cc/repository/__init__.py
@@ -1,2 +1,2 @@
 from .file_storage import FileRetrievalError, IFileRepository, LocalStorageFileRepository
-from .i_agent_binary_repository import IAgentBinaryRepository
+from .i_agent_binary_repository import IAgentBinaryRepository, AgentRetrivalError
diff --git a/monkey/monkey_island/cc/repository/agent_binary_repository.py b/monkey/monkey_island/cc/repository/agent_binary_repository.py
new file mode 100644
index 000000000..ec414c95c
--- /dev/null
+++ b/monkey/monkey_island/cc/repository/agent_binary_repository.py
@@ -0,0 +1,29 @@
+from typing import BinaryIO
+
+from . import AgentRetrivalError, FileRetrivalError, IAgentBinaryRepository, IFileRepository
+
+LINUX_AGENT_FILE_NAME = "monkey-linux-64"
+WINDOWS_AGENT_FILE_NAME = "monkey-windows-64.exe"
+
+
+class AgentBinaryRepository(IAgentBinaryRepository):
+    def __init__(self, file_repository: IFileRepository):
+        self._file_repository = file_repository
+
+    def get_linux_binary(self) -> BinaryIO:
+        try:
+            agent_binary = self._file_repository.open_file(LINUX_AGENT_FILE_NAME)
+            return agent_binary
+        except FileRetrivalError as err:
+            raise AgentRetrivalError(
+                f"An error occurred while retrieving the Linux agent binary: {err}"
+            )
+
+    def get_windows_binary(self) -> BinaryIO:
+        try:
+            agent_binary = self._file_repository.open_file(WINDOWS_AGENT_FILE_NAME)
+            return agent_binary
+        except FileRetrivalError as err:
+            raise AgentRetrivalError(
+                f"An error occurred while retrieving the Windows agent binary: {err}"
+            )
diff --git a/monkey/monkey_island/cc/repository/i_agent_binary_repository.py b/monkey/monkey_island/cc/repository/i_agent_binary_repository.py
index a1d858e05..aac33176c 100644
--- a/monkey/monkey_island/cc/repository/i_agent_binary_repository.py
+++ b/monkey/monkey_island/cc/repository/i_agent_binary_repository.py
@@ -2,6 +2,10 @@ import abc
 from typing import BinaryIO
 
 
+class AgentRetrivalError(IOError):
+    pass
+
+
 class IAgentBinaryRepository(metaclass=abc.ABCMeta):
     """
     A repository that retrieves the agent binaries
diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py
index 431eb68d3..40ff4e69b 100644
--- a/monkey/monkey_island/cc/services/initialize.py
+++ b/monkey/monkey_island/cc/services/initialize.py
@@ -2,7 +2,13 @@ from pathlib import Path
 
 from common import DIContainer
 from common.aws import AWSInstance
-from monkey_island.cc.repository import IFileRepository, LocalStorageFileRepository
+from monkey_island.cc.repository import (
+    AgentBinaryRepository,
+    IAgentBinaryRepository,
+    IFileRepository,
+    LocalStorageFileRepository,
+)
+from monkey_island.cc.server_utils.consts import MONKEY_ISLAND_ABS_PATH
 from monkey_island.cc.services import AWSService
 from monkey_island.cc.services.post_breach_files import PostBreachFilesService
 from monkey_island.cc.services.run_local_monkey import LocalMonkeyRunService
@@ -10,6 +16,8 @@ from monkey_island.cc.services.run_local_monkey import LocalMonkeyRunService
 from . import AuthenticationService, JsonFileUserDatastore
 from .reporting.report import ReportService
 
+AGENT_BINARIES_PATH = Path(MONKEY_ISLAND_ABS_PATH) / "cc" / "binaries"
+
 
 def initialize_services(data_dir: Path) -> DIContainer:
     container = DIContainer()
@@ -19,6 +27,7 @@ def initialize_services(data_dir: Path) -> DIContainer:
         IFileRepository, LocalStorageFileRepository(data_dir / "custom_pbas")
     )
     container.register_instance(AWSService, container.resolve(AWSService))
+    container.register_instance(IAgentBinaryRepository, _build_agent_binary_repository())
 
     # This is temporary until we get DI all worked out.
     PostBreachFilesService.initialize(container.resolve(IFileRepository))
@@ -27,3 +36,8 @@ def initialize_services(data_dir: Path) -> DIContainer:
     ReportService.initialize(container.resolve(AWSService))
 
     return container
+
+
+def _build_agent_binary_repository():
+    file_repository = LocalStorageFileRepository(AGENT_BINARIES_PATH)
+    return AgentBinaryRepository(file_repository)