parent
94e88d2c8b
commit
e7459412e4
|
@ -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
|
||||
|
|
|
@ -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}"
|
||||
)
|
|
@ -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
|
||||
|
|
|
@ -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)
|
||||
|
|
Loading…
Reference in New Issue