forked from p15670423/monkey
Island: Pass default agent configuration to repository constructor
This commit is contained in:
parent
8e875143ac
commit
1f00a13649
|
@ -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,
|
||||
)
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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}")
|
||||
|
||||
|
|
|
@ -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(
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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()
|
||||
|
|
|
@ -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()
|
||||
|
|
Loading…
Reference in New Issue