Merge pull request #2326 from guardicore/2256-initialize-i-machine-repository

2256 initialize i machine repository

monkey/monkey_island/cc/island_event_handlers/__init__.py

@@ -1 +1,2 @@
 from .reset_agent_configuration import reset_agent_configuration
+from .reset_machine_repository import reset_machine_repository

monkey/monkey_island/cc/island_event_handlers/reset_machine_repository.py

@@ -0,0 +1,14 @@
+from monkey_island.cc.repository import IMachineRepository, initialize_machine_repository
+
+
+class reset_machine_repository:
+    """
+    Callable class that handles reset and reinitialization of IMachineRepository
+    """
+
+    def __init__(self, machine_repository: IMachineRepository):
+        self._machine_repository = machine_repository
+
+    def __call__(self):
+        self._machine_repository.reset()
+        initialize_machine_repository(self._machine_repository)

monkey/monkey_island/cc/models/machine.py

@@ -20,6 +20,9 @@ class Machine(MutableInfectionMonkeyBaseModel):
     hardware_id: Optional[HardwareID]
     """An identifier generated by the agent that uniquely identifies a machine"""
 
+    island: bool = Field(default=False, allow_mutation=False)
+    """Whether or not the machine is an island (C&C server)"""
+
     network_interfaces: Sequence[IPv4Interface] = tuple()
     """The machine's networking interfaces"""
 

monkey/monkey_island/cc/repository/__init__.py

@@ -27,3 +27,5 @@ from .mongo_machine_repository import MongoMachineRepository
 from .mongo_agent_repository import MongoAgentRepository
 from .mongo_node_repository import MongoNodeRepository
 from .mongo_agent_event_repository import MongoAgentEventRepository
+
+from .utils import initialize_machine_repository

monkey/monkey_island/cc/repository/utils.py

@@ -0,0 +1,35 @@
+import platform
+from socket import gethostname
+from uuid import getnode
+
+from common import OperatingSystem
+from common.network.network_utils import get_network_interfaces
+from monkey_island.cc.models import Machine
+
+from . import IMachineRepository, UnknownRecordError
+
+
+def initialize_machine_repository(machine_repository: IMachineRepository):
+    """
+    Populate an IMachineRepository with island machine data
+
+    If the island is not already present in the IMachineRepository, add it.
+
+    :param machine_repository: The repository to populate
+    :raises StorageError: If an error occurs while attempting to store data in the repository
+    """
+    hardware_id = getnode()
+
+    try:
+        machine_repository.get_machine_by_hardware_id(hardware_id)
+    except UnknownRecordError:
+        machine = Machine(
+            id=machine_repository.get_new_id(),
+            hardware_id=hardware_id,
+            island=True,
+            network_interfaces=get_network_interfaces(),
+            operating_system=OperatingSystem(platform.system().lower()),
+            operating_system_version=platform.version(),
+            hostname=gethostname(),
+        )
+        machine_repository.upsert_machine(machine)

monkey/monkey_island/cc/services/initialize.py

@@ -43,6 +43,7 @@ from monkey_island.cc.repository import (
     MongoMachineRepository,
     MongoNodeRepository,
     RetrievalError,
+    initialize_machine_repository,
 )
 from monkey_island.cc.server_utils.consts import MONKEY_ISLAND_ABS_PATH
 from monkey_island.cc.server_utils.encryption import ILockableEncryptor, RepositoryEncryptor
@@ -114,7 +115,7 @@ def _register_repositories(container: DIContainer, data_dir: Path):
     container.register_instance(IAgentEventRepository, container.resolve(MongoAgentEventRepository))
 
     container.register_instance(INodeRepository, container.resolve(MongoNodeRepository))
-    container.register_instance(IMachineRepository, container.resolve(MongoMachineRepository))
+    container.register_instance(IMachineRepository, _build_machine_repository(container))
     container.register_instance(IAgentRepository, container.resolve(MongoAgentRepository))
 
 
@@ -124,7 +125,7 @@ def _decorate_file_repository(file_repository: IFileRepository) -> IFileReposito
     )
 
 
-def _build_agent_binary_repository():
+def _build_agent_binary_repository() -> IAgentBinaryRepository:
     file_repository = _decorate_file_repository(LocalStorageFileRepository(AGENT_BINARIES_PATH))
     agent_binary_repository = AgentBinaryRepository(file_repository)
 
@@ -133,6 +134,13 @@ def _build_agent_binary_repository():
     return agent_binary_repository
 
 
+def _build_machine_repository(container: DIContainer) -> IMachineRepository:
+    machine_repository = container.resolve(MongoMachineRepository)
+    initialize_machine_repository(machine_repository)
+
+    return machine_repository
+
+
 def _setup_agent_event_serializers(container: DIContainer):
     agent_event_serializer_registry = AgentEventSerializerRegistry()
     register_common_agent_event_serializers(agent_event_serializer_registry)

monkey/monkey_island/cc/setup/island_event_handlers.py

@@ -2,12 +2,14 @@ from functools import partial
 
 from common import DIContainer
 from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic
-from monkey_island.cc.island_event_handlers import reset_agent_configuration
+from monkey_island.cc.island_event_handlers import (
+    reset_agent_configuration,
+    reset_machine_repository,
+)
 from monkey_island.cc.repository import (
     IAgentEventRepository,
     IAgentRepository,
     ICredentialsRepository,
-    IMachineRepository,
     INodeRepository,
 )
 from monkey_island.cc.services.database import Database
@@ -39,11 +41,14 @@ def _subscribe_clear_simulation_data_events(
         IslandEventTopic.CLEAR_SIMULATION_DATA, credentials_repository.remove_stolen_credentials
     )
 
+    island_event_queue.subscribe(
+        IslandEventTopic.CLEAR_SIMULATION_DATA, container.resolve(reset_machine_repository)
+    )
+
     for i_repository in [
         INodeRepository,
         IAgentEventRepository,
         IAgentRepository,
-        IMachineRepository,
     ]:
         repository = container.resolve(i_repository)
         island_event_queue.subscribe(IslandEventTopic.CLEAR_SIMULATION_DATA, repository.reset)

monkey/tests/unit_tests/monkey_island/cc/models/test_machine.py

@@ -12,6 +12,7 @@ MACHINE_OBJECT_DICT = MappingProxyType(
     {
         "id": 1,
         "hardware_id": uuid.getnode(),
+        "island": True,
         "network_interfaces": [IPv4Interface("10.0.0.1/24"), IPv4Interface("192.168.5.32/16")],
         "operating_system": OperatingSystem.WINDOWS,
         "operating_system_version": "eXtra Problems",
@@ -23,6 +24,7 @@ MACHINE_SIMPLE_DICT = MappingProxyType(
     {
         "id": 1,
         "hardware_id": uuid.getnode(),
+        "island": True,
         "network_interfaces": ["10.0.0.1/24", "192.168.5.32/16"],
         "operating_system": "windows",
         "operating_system_version": "eXtra Problems",
@@ -52,6 +54,7 @@ def test_to_dict():
     [
         ("id", "not-an-int"),
         ("hardware_id", "not-an-int"),
+        ("island", "not-a-bool"),
         ("network_interfaces", "not-a-list"),
         ("operating_system", 2.1),
         ("operating_system", "bsd"),
@@ -130,6 +133,21 @@ def test_hardware_id_default():
     assert m.hardware_id is None
 
 
+def test_island_immutable():
+    m = Machine(**MACHINE_OBJECT_DICT)
+    with pytest.raises(TypeError):
+        m.island = True
+
+
+def test_island_default():
+    missing_island_dict = MACHINE_OBJECT_DICT.copy()
+    del missing_island_dict["island"]
+
+    m = Machine(**missing_island_dict)
+
+    assert m.island is False
+
+
 def test_network_interfaces_set_valid_value():
     m = Machine(**MACHINE_OBJECT_DICT)