Island: Pass default agent configuration to repository constructor

This commit is contained in:
Mike Salvatore 2022-06-22 07:35:06 -04:00
parent 8e875143ac
commit 1f00a13649
7 changed files with 43 additions and 21 deletions

View File

@ -2,4 +2,7 @@ from .agent_configuration import (
AgentConfiguration,
AgentConfigurationSchema,
)
from .default_agent_configuration import DEFAULT_AGENT_CONFIGURATION
from .default_agent_configuration import (
DEFAULT_AGENT_CONFIGURATION_JSON,
build_default_agent_configuration,
)

View File

@ -1,4 +1,6 @@
DEFAULT_AGENT_CONFIGURATION = """{
from . import AgentConfiguration, AgentConfigurationSchema
DEFAULT_AGENT_CONFIGURATION_JSON = """{
"keep_tunnel_open_time": 30,
"post_breach_actions": [
{
@ -215,3 +217,8 @@ DEFAULT_AGENT_CONFIGURATION = """{
}
}
"""
def build_default_agent_configuration() -> AgentConfiguration:
schema = AgentConfigurationSchema()
return schema.loads(DEFAULT_AGENT_CONFIGURATION_JSON)

View File

@ -1,10 +1,6 @@
import io
from common.configuration import (
DEFAULT_AGENT_CONFIGURATION,
AgentConfiguration,
AgentConfigurationSchema,
)
from common.configuration import AgentConfiguration, AgentConfigurationSchema
from monkey_island.cc import repository
from monkey_island.cc.repository import (
IAgentConfigurationRepository,
@ -16,7 +12,10 @@ AGENT_CONFIGURATION_FILE_NAME = "agent_configuration.json"
class FileAgentConfigurationRepository(IAgentConfigurationRepository):
def __init__(self, file_repository: IFileRepository):
def __init__(
self, default_agent_configuration: AgentConfiguration, file_repository: IFileRepository
):
self._default_agent_configuration = default_agent_configuration
self._file_repository = file_repository
self._schema = AgentConfigurationSchema()
@ -27,7 +26,7 @@ class FileAgentConfigurationRepository(IAgentConfigurationRepository):
return self._schema.loads(configuration_json)
except repository.FileNotFoundError:
return self._schema.loads(DEFAULT_AGENT_CONFIGURATION)
return self._default_agent_configuration
except Exception as err:
raise RetrievalError(f"Error retrieving the agent configuration: {err}")

View File

@ -3,6 +3,7 @@ from pathlib import Path
from common import DIContainer
from common.aws import AWSInstance
from common.configuration import AgentConfiguration, build_default_agent_configuration
from common.utils.file_utils import get_binary_io_sha256_hash
from monkey_island.cc.repository import (
AgentBinaryRepository,
@ -30,6 +31,9 @@ def initialize_services(data_dir: Path) -> DIContainer:
container = DIContainer()
container.register_convention(Path, "data_dir", data_dir)
container.register_convention(
AgentConfiguration, "default_agent_configuration", build_default_agent_configuration()
)
container.register_instance(AWSInstance, AWSInstance())
container.register_instance(

View File

@ -25,7 +25,7 @@ from tests.common.example_agent_configuration import (
from common import OperatingSystems
from common.configuration import (
DEFAULT_AGENT_CONFIGURATION,
DEFAULT_AGENT_CONFIGURATION_JSON,
AgentConfiguration,
AgentConfigurationSchema,
)
@ -180,6 +180,6 @@ def test_agent_configuration():
def test_default_agent_configuration():
schema = AgentConfigurationSchema()
config = schema.loads(DEFAULT_AGENT_CONFIGURATION)
config = schema.loads(DEFAULT_AGENT_CONFIGURATION_JSON)
assert isinstance(config, AgentConfiguration)

View File

@ -6,6 +6,8 @@ from typing import Callable, Dict
import pytest
from _pytest.monkeypatch import MonkeyPatch
from common.configuration import AgentConfiguration, build_default_agent_configuration
MONKEY_BASE_PATH = str(Path(__file__).parent.parent.parent)
sys.path.insert(0, MONKEY_BASE_PATH)
@ -54,3 +56,8 @@ def load_monkey_config(data_for_tests_dir) -> Callable[[str], Dict]:
return json.loads(open(config_path, "r").read())
return inner
@pytest.fixture
def default_agent_configuration() -> AgentConfiguration:
return build_default_agent_configuration()

View File

@ -2,12 +2,16 @@ import pytest
from tests.common.example_agent_configuration import AGENT_CONFIGURATION
from tests.monkey_island import OpenErrorFileRepository, SingleFileRepository
from common.configuration import DEFAULT_AGENT_CONFIGURATION, AgentConfigurationSchema
from common.configuration import AgentConfigurationSchema
from monkey_island.cc.repository import FileAgentConfigurationRepository, RetrievalError
def test_store_agent_config():
repository = FileAgentConfigurationRepository(SingleFileRepository())
@pytest.fixture
def repository(default_agent_configuration):
return FileAgentConfigurationRepository(default_agent_configuration, SingleFileRepository())
def test_store_agent_config(repository):
schema = AgentConfigurationSchema()
agent_configuration = schema.load(AGENT_CONFIGURATION)
@ -17,18 +21,16 @@ def test_store_agent_config():
assert retrieved_agent_configuration == agent_configuration
def test_get_default_agent_config():
repository = FileAgentConfigurationRepository(SingleFileRepository())
schema = AgentConfigurationSchema()
default_agent_configuration = schema.loads(DEFAULT_AGENT_CONFIGURATION)
def test_get_default_agent_config(repository, default_agent_configuration):
retrieved_agent_configuration = repository.get_configuration()
assert retrieved_agent_configuration == default_agent_configuration
def test_get_agent_config_retrieval_error():
repository = FileAgentConfigurationRepository(OpenErrorFileRepository())
def test_get_agent_config_retrieval_error(default_agent_configuration):
repository = FileAgentConfigurationRepository(
default_agent_configuration, OpenErrorFileRepository()
)
with pytest.raises(RetrievalError):
repository.get_configuration()