From 43a53112fc2d940cc91d4484fe9688daa9ce35eb Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 24 Aug 2022 08:55:35 -0400 Subject: [PATCH 01/22] Island: Use monkey_island.cc.models.Machine in IMachineRepository --- .../cc/repository/i_machine_repository.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/monkey/monkey_island/cc/repository/i_machine_repository.py b/monkey/monkey_island/cc/repository/i_machine_repository.py index 604d58068..79538e22f 100644 --- a/monkey/monkey_island/cc/repository/i_machine_repository.py +++ b/monkey/monkey_island/cc/repository/i_machine_repository.py @@ -1,20 +1,19 @@ from abc import ABC from typing import Optional, Sequence +from monkey_island.cc.models import Machine + class IMachineRepository(ABC): - # TODO define Machine object(ORM model) - def save_machine(self, machine: Machine): # noqa: F821 + def save_machine(self, machine: Machine): pass - # TODO define Machine object(ORM model) # TODO define or re-use machine state. - # TODO investigate where should the state be stored in edge or both edge and machine def get_machines( self, id: Optional[str] = None, ips: Optional[Sequence[str]] = None, state: Optional[MachineState] = None, # noqa: F821 - is_island: Optional[bool] = None, # noqa: F841 - ) -> Sequence[Machine]: # noqa: F821 + is_island: Optional[bool] = None, + ) -> Sequence[Machine]: pass From ebcfe5a9fc729ea17f48fd2d70de56c2bd5de91d Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 29 Aug 2022 14:35:45 -0400 Subject: [PATCH 02/22] Island: Add docstrings to Machine --- monkey/monkey_island/cc/models/machine.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/monkey/monkey_island/cc/models/machine.py b/monkey/monkey_island/cc/models/machine.py index 77000c39a..17771c016 100644 --- a/monkey/monkey_island/cc/models/machine.py +++ b/monkey/monkey_island/cc/models/machine.py @@ -12,12 +12,25 @@ MachineID: TypeAlias = PositiveInt class Machine(MutableInfectionMonkeyBaseModel): + """Represents machines, VMs, or other network nodes discovered by Infection Monkey""" + id: MachineID = Field(..., allow_mutation=False) + """Uniquely identifies the machine within the island""" + hardware_id: Optional[HardwareID] + """An identifier generated by the agent that uniquely identifies a machine""" + network_interfaces: Sequence[IPv4Interface] + """The machine's networking interfaces""" + operating_system: OperatingSystem + """The operating system the machine is running""" + operating_system_version: str + """The specific version of the operating system the machine is running""" + hostname: str + """The hostname of the machine""" _make_immutable_sequence = validator("network_interfaces", pre=True, allow_reuse=True)( make_immutable_sequence From 694cdca883db18977f9ccb4ce0a192754317d6ca Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 29 Aug 2022 15:00:45 -0400 Subject: [PATCH 03/22] Island: Make Machine.operating_system Optional --- monkey/monkey_island/cc/models/machine.py | 2 +- .../unit_tests/monkey_island/cc/models/test_machine.py | 9 +++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/models/machine.py b/monkey/monkey_island/cc/models/machine.py index 17771c016..686e677b6 100644 --- a/monkey/monkey_island/cc/models/machine.py +++ b/monkey/monkey_island/cc/models/machine.py @@ -23,7 +23,7 @@ class Machine(MutableInfectionMonkeyBaseModel): network_interfaces: Sequence[IPv4Interface] """The machine's networking interfaces""" - operating_system: OperatingSystem + operating_system: Optional[OperatingSystem] """The operating system the machine is running""" operating_system_version: str diff --git a/monkey/tests/unit_tests/monkey_island/cc/models/test_machine.py b/monkey/tests/unit_tests/monkey_island/cc/models/test_machine.py index 50cdaa61d..8414cd9d6 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/models/test_machine.py +++ b/monkey/tests/unit_tests/monkey_island/cc/models/test_machine.py @@ -84,6 +84,15 @@ def test_construct_invalid_field__value_error(key, value): Machine(**invalid_type_dict) +@pytest.mark.parametrize("field", ["hardware_id", "operating_system"]) +def test_optional_fields(field): + none_field_dict = MACHINE_SIMPLE_DICT.copy() + none_field_dict[field] = None + + # Raises exception_on_failure + Machine(**none_field_dict) + + def test_construct__extra_fields_forbidden(): extra_field_dict = MACHINE_SIMPLE_DICT.copy() extra_field_dict["extra_field"] = 99 # red balloons From 1de552ea94706ce75ff9464c1af7f1886bc6d73e Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 29 Aug 2022 15:14:03 -0400 Subject: [PATCH 04/22] Island: Add UnknownRecordError to repository.errors --- monkey/monkey_island/cc/repository/__init__.py | 2 +- monkey/monkey_island/cc/repository/errors.py | 6 ++++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/repository/__init__.py b/monkey/monkey_island/cc/repository/__init__.py index 075317b43..a57518abf 100644 --- a/monkey/monkey_island/cc/repository/__init__.py +++ b/monkey/monkey_island/cc/repository/__init__.py @@ -1,4 +1,4 @@ -from .errors import RemovalError, RetrievalError, StorageError +from .errors import RemovalError, RetrievalError, StorageError, UnknownRecordError from .i_file_repository import FileNotFoundError, IFileRepository diff --git a/monkey/monkey_island/cc/repository/errors.py b/monkey/monkey_island/cc/repository/errors.py index aeb9fa23c..986d3007a 100644 --- a/monkey/monkey_island/cc/repository/errors.py +++ b/monkey/monkey_island/cc/repository/errors.py @@ -20,3 +20,9 @@ class StorageError(RuntimeError): """ pass + + +class UnknownRecordError(RuntimeError): + """ + Raised when the repository does not contain any data matching the request. + """ From 6b083ca61b8e3596ff065cefadfbfa0df51f9f4f Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 29 Aug 2022 15:15:49 -0400 Subject: [PATCH 05/22] Island: Remove superfluous "pass" from repository.errors --- monkey/monkey_island/cc/repository/errors.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/monkey/monkey_island/cc/repository/errors.py b/monkey/monkey_island/cc/repository/errors.py index 986d3007a..a5c26fe15 100644 --- a/monkey/monkey_island/cc/repository/errors.py +++ b/monkey/monkey_island/cc/repository/errors.py @@ -3,24 +3,18 @@ class RemovalError(RuntimeError): Raised when a repository encounters an error while attempting to remove data. """ - pass - class RetrievalError(RuntimeError): """ Raised when a repository encounters an error while attempting to retrieve data. """ - pass - class StorageError(RuntimeError): """ Raised when a repository encounters an error while attempting to store data. """ - pass - class UnknownRecordError(RuntimeError): """ From eb3fe21b11d93f61a240946c9fc881ec1722dee1 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 29 Aug 2022 15:34:36 -0400 Subject: [PATCH 06/22] Island: Redefine IMachineRepository --- .../cc/repository/i_machine_repository.py | 80 +++++++++++++++---- vulture_allowlist.py | 7 +- 2 files changed, 71 insertions(+), 16 deletions(-) diff --git a/monkey/monkey_island/cc/repository/i_machine_repository.py b/monkey/monkey_island/cc/repository/i_machine_repository.py index 79538e22f..97bd3cf87 100644 --- a/monkey/monkey_island/cc/repository/i_machine_repository.py +++ b/monkey/monkey_island/cc/repository/i_machine_repository.py @@ -1,19 +1,71 @@ -from abc import ABC -from typing import Optional, Sequence +from abc import ABC, abstractmethod +from ipaddress import IPv4Address +from typing import Sequence -from monkey_island.cc.models import Machine +from common.types import HardwareID +from monkey_island.cc.models import Machine, MachineID class IMachineRepository(ABC): - def save_machine(self, machine: Machine): - pass + """A repository used to store and retrieve Machines""" - # TODO define or re-use machine state. - def get_machines( - self, - id: Optional[str] = None, - ips: Optional[Sequence[str]] = None, - state: Optional[MachineState] = None, # noqa: F821 - is_island: Optional[bool] = None, - ) -> Sequence[Machine]: - pass + @abstractmethod + def create_machine(self) -> Machine: + """ + Create a new `Machine` in the repository + + :return: A new `Machine` with a unique ID + :raises StorageError: If a new `Machine` could not be created + """ + + @abstractmethod + def update_machine(self, machine: Machine): + """ + Update an existing `Machine` in the repository + + :param machine: An updated Machine object to store in the repository + :raises UnknownRecordError: If the provided `Machine` does not exist in the repository + :raises StorageError: If an error occurred while attempting to store the `Machine` + """ + + @abstractmethod + def get_machine_by_id(self, id_: MachineID) -> Machine: + """ + Get a `Machine` by ID + + :param id: The ID of the `Machine` to be retrieved + :return: A `Machine` with a matching `id` + :raises UnknownRecordError: If a `Machine` with the specified `id` does not exist in the + repository + :raises RetrievalError: If an error occurred while attempting to retrieve the `Machine` + """ + + @abstractmethod + def get_machine_by_hardware_id(self, hardware_id: HardwareID) -> Machine: + """ + Get a `Machine` by ID + + :param hardware_id: The hardware ID of the `Machine` to be retrieved + :return: A `Machine` with a matching `hardware_id` + :raises UnknownRecordError: If a `Machine` with the specified `hardware_id` does not exist + in the repository + :raises RetrievalError: If an error occurred while attempting to retrieve the `Machine` + """ + + @abstractmethod + def get_machines_by_ip(self, ip: IPv4Address) -> Sequence[Machine]: + """ + Search for machines by IP address + + :param ip: The IP address to search for + :return: A sequence of Machines that have a network interface with a matching IP + :raises UnknownRecordError: If a `Machine` with the specified `ip` does not exist in the + repository + :raises RetrievalError: If an error occurred while attempting to retrieve the `Machine` + """ + + @abstractmethod + def reset(self): + """ + Removes all data from the repository + """ diff --git a/vulture_allowlist.py b/vulture_allowlist.py index b203a2bea..9da5fee84 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -244,8 +244,11 @@ IConfigRepository.get_config_field ILogRepository.get_logs ILogRepository.save_log ILogRepository.delete_log -IMachineRepository.save_machine -IMachineRepository.get_machines +IMachineRepository.create_machine +IMachineRepository.update_machine +IMachineRepository.get_machine_by_id +IMachineRepository.get_machine_by_hardware_id +IMachineRepository.get_machines_by_ip INetworkMapRepository.get_map INetworkMapRepository.save_netmap IReportRepository From da752e041b19a67839aac69a483335bb05f99161 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 29 Aug 2022 15:43:57 -0400 Subject: [PATCH 07/22] Project: Exclude vulture_allowlist.py from mypy checks --- .pre-commit-config.yaml | 1 + pyproject.toml | 1 + 2 files changed, 2 insertions(+) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 7a553388a..a4e93e4be 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -40,6 +40,7 @@ repos: hooks: - id: mypy additional_dependencies: [types-paramiko, types-python-dateutil, types-requests] + exclude: "vulture_allowlist.py" args: [--ignore-missing-imports] - repo: https://github.com/koalaman/shellcheck-precommit rev: v0.7.2 diff --git a/pyproject.toml b/pyproject.toml index 2573d56fd..813908fa6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,4 +30,5 @@ exclude = ["monkey/monkey_island/cc/ui/", "monkey/tests/", "monkey/monkey_island paths = ["."] [tool.mypy] +exclude = 'vulture_allowlist\.py' show_error_codes = true From 3eda8d640dcfcfedf4b5b480a590dcb5b04622c6 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 29 Aug 2022 15:57:16 -0400 Subject: [PATCH 08/22] Island: Export IMachineRepository from repository package --- monkey/monkey_island/cc/repository/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/monkey/monkey_island/cc/repository/__init__.py b/monkey/monkey_island/cc/repository/__init__.py index a57518abf..d9987e677 100644 --- a/monkey/monkey_island/cc/repository/__init__.py +++ b/monkey/monkey_island/cc/repository/__init__.py @@ -7,6 +7,7 @@ from .i_agent_configuration_repository import IAgentConfigurationRepository from .i_simulation_repository import ISimulationRepository from .i_credentials_repository import ICredentialsRepository from .i_user_repository import IUserRepository +from .i_machine_repository import IMachineRepository from .local_storage_file_repository import LocalStorageFileRepository From bf6125dd55c3839dd8f8742890b13b3e456b9ea9 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 29 Aug 2022 19:31:35 -0400 Subject: [PATCH 09/22] Island: Add MONGO_OBJECT_ID_KEY --- monkey/monkey_island/cc/repository/consts.py | 1 + 1 file changed, 1 insertion(+) create mode 100644 monkey/monkey_island/cc/repository/consts.py diff --git a/monkey/monkey_island/cc/repository/consts.py b/monkey/monkey_island/cc/repository/consts.py new file mode 100644 index 000000000..161768005 --- /dev/null +++ b/monkey/monkey_island/cc/repository/consts.py @@ -0,0 +1 @@ +MONGO_OBJECT_ID_KEY = "_id" From 10d8d8e7561e1ac02696ce8a3660484c45040fe7 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 29 Aug 2022 19:32:40 -0400 Subject: [PATCH 10/22] Island: Use MONGO_OBJECT_ID_KEY in MongoCredentialsRepository --- .../cc/repository/mongo_credentials_repository.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/repository/mongo_credentials_repository.py b/monkey/monkey_island/cc/repository/mongo_credentials_repository.py index ee8beb13c..25789f937 100644 --- a/monkey/monkey_island/cc/repository/mongo_credentials_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_credentials_repository.py @@ -7,6 +7,8 @@ from monkey_island.cc.repository import RemovalError, RetrievalError, StorageErr from monkey_island.cc.repository.i_credentials_repository import ICredentialsRepository from monkey_island.cc.server_utils.encryption import ILockableEncryptor +from .consts import MONGO_OBJECT_ID_KEY + class MongoCredentialsRepository(ICredentialsRepository): """ @@ -51,7 +53,7 @@ class MongoCredentialsRepository(ICredentialsRepository): collection_result = [] list_collection_result = list(collection.find({})) for encrypted_credentials in list_collection_result: - del encrypted_credentials["_id"] + del encrypted_credentials[MONGO_OBJECT_ID_KEY] plaintext_credentials = self._decrypt_credentials_mapping(encrypted_credentials) collection_result.append(Credentials.from_mapping(plaintext_credentials)) From da8ed9e6db85c346e0d1bcc7d39b228160485b87 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 29 Aug 2022 19:35:04 -0400 Subject: [PATCH 11/22] Island: Fix ICredentialsRepository import in MongoCredentialsRepository --- .../cc/repository/mongo_credentials_repository.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/repository/mongo_credentials_repository.py b/monkey/monkey_island/cc/repository/mongo_credentials_repository.py index 25789f937..351cac981 100644 --- a/monkey/monkey_island/cc/repository/mongo_credentials_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_credentials_repository.py @@ -3,8 +3,12 @@ from typing import Any, Dict, Mapping, Sequence from pymongo import MongoClient from common.credentials import Credentials -from monkey_island.cc.repository import RemovalError, RetrievalError, StorageError -from monkey_island.cc.repository.i_credentials_repository import ICredentialsRepository +from monkey_island.cc.repository import ( + ICredentialsRepository, + RemovalError, + RetrievalError, + StorageError, +) from monkey_island.cc.server_utils.encryption import ILockableEncryptor from .consts import MONGO_OBJECT_ID_KEY From 40601b955c591fcca359e91be7bb18fab559e77d Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 29 Aug 2022 20:18:01 -0400 Subject: [PATCH 12/22] Island: Add default value for Machine.network_interfaces --- monkey/monkey_island/cc/models/machine.py | 2 +- .../unit_tests/monkey_island/cc/models/test_machine.py | 9 +++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/models/machine.py b/monkey/monkey_island/cc/models/machine.py index 686e677b6..b09e8bf99 100644 --- a/monkey/monkey_island/cc/models/machine.py +++ b/monkey/monkey_island/cc/models/machine.py @@ -20,7 +20,7 @@ class Machine(MutableInfectionMonkeyBaseModel): hardware_id: Optional[HardwareID] """An identifier generated by the agent that uniquely identifies a machine""" - network_interfaces: Sequence[IPv4Interface] + network_interfaces: Sequence[IPv4Interface] = Field(default_factory=tuple) """The machine's networking interfaces""" operating_system: Optional[OperatingSystem] diff --git a/monkey/tests/unit_tests/monkey_island/cc/models/test_machine.py b/monkey/tests/unit_tests/monkey_island/cc/models/test_machine.py index 8414cd9d6..a2e712a25 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/models/test_machine.py +++ b/monkey/tests/unit_tests/monkey_island/cc/models/test_machine.py @@ -141,6 +141,15 @@ def test_network_interfaces_sequence_is_immutable(): assert not isinstance(m.network_interfaces, MutableSequence) +def test_network_interfaces_default(): + missing_network_interfaces_dict = MACHINE_OBJECT_DICT.copy() + del missing_network_interfaces_dict["network_interfaces"] + + m = Machine(**missing_network_interfaces_dict) + + assert len(m.network_interfaces) == 0 + + def test_operating_system_set_valid_value(): m = Machine(**MACHINE_OBJECT_DICT) From 3e2244cd62e6fd912a2a40483b87879390b763ad Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 29 Aug 2022 20:23:56 -0400 Subject: [PATCH 13/22] UT: Add test_operating_system_default_value() --- .../unit_tests/monkey_island/cc/models/test_machine.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/monkey/tests/unit_tests/monkey_island/cc/models/test_machine.py b/monkey/tests/unit_tests/monkey_island/cc/models/test_machine.py index a2e712a25..a81f1823f 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/models/test_machine.py +++ b/monkey/tests/unit_tests/monkey_island/cc/models/test_machine.py @@ -164,6 +164,15 @@ def test_operating_system_set_invalid_value(): m.operating_system = "MacOS" +def test_operating_system_default_value(): + missing_operating_system_dict = MACHINE_OBJECT_DICT.copy() + del missing_operating_system_dict["operating_system"] + + m = Machine(**missing_operating_system_dict) + + assert m.operating_system is None + + def test_set_operating_system_version(): m = Machine(**MACHINE_OBJECT_DICT) From 5d51b40475545efc703d8f1da4b8b6c9e58af905 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 29 Aug 2022 20:25:33 -0400 Subject: [PATCH 14/22] Island: Add default value for Machine.operating_system_version --- monkey/monkey_island/cc/models/machine.py | 2 +- .../unit_tests/monkey_island/cc/models/test_machine.py | 9 +++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/models/machine.py b/monkey/monkey_island/cc/models/machine.py index b09e8bf99..31c5d41ad 100644 --- a/monkey/monkey_island/cc/models/machine.py +++ b/monkey/monkey_island/cc/models/machine.py @@ -26,7 +26,7 @@ class Machine(MutableInfectionMonkeyBaseModel): operating_system: Optional[OperatingSystem] """The operating system the machine is running""" - operating_system_version: str + operating_system_version: str = Field(default="") """The specific version of the operating system the machine is running""" hostname: str diff --git a/monkey/tests/unit_tests/monkey_island/cc/models/test_machine.py b/monkey/tests/unit_tests/monkey_island/cc/models/test_machine.py index a81f1823f..e60069ed6 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/models/test_machine.py +++ b/monkey/tests/unit_tests/monkey_island/cc/models/test_machine.py @@ -180,6 +180,15 @@ def test_set_operating_system_version(): m.operating_system_version = "1234" +def test_operating_system_version_default_value(): + missing_operating_system_version_dict = MACHINE_OBJECT_DICT.copy() + del missing_operating_system_version_dict["operating_system_version"] + + m = Machine(**missing_operating_system_version_dict) + + assert m.operating_system_version == "" + + def test_set_hostname(): m = Machine(**MACHINE_OBJECT_DICT) From ac1cda40a815b63d98ecc0e69de78352c189c2d8 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 29 Aug 2022 20:27:57 -0400 Subject: [PATCH 15/22] Island: Add default value for Machine.hostname --- monkey/monkey_island/cc/models/machine.py | 2 +- .../unit_tests/monkey_island/cc/models/test_machine.py | 9 +++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/models/machine.py b/monkey/monkey_island/cc/models/machine.py index 31c5d41ad..4c5925390 100644 --- a/monkey/monkey_island/cc/models/machine.py +++ b/monkey/monkey_island/cc/models/machine.py @@ -29,7 +29,7 @@ class Machine(MutableInfectionMonkeyBaseModel): operating_system_version: str = Field(default="") """The specific version of the operating system the machine is running""" - hostname: str + hostname: str = Field(default="") """The hostname of the machine""" _make_immutable_sequence = validator("network_interfaces", pre=True, allow_reuse=True)( diff --git a/monkey/tests/unit_tests/monkey_island/cc/models/test_machine.py b/monkey/tests/unit_tests/monkey_island/cc/models/test_machine.py index e60069ed6..db8fb5825 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/models/test_machine.py +++ b/monkey/tests/unit_tests/monkey_island/cc/models/test_machine.py @@ -194,3 +194,12 @@ def test_set_hostname(): # Raises exception_on_failure m.operating_system_version = "wopr" + + +def test_hostname_default_value(): + missing_hostname_dict = MACHINE_OBJECT_DICT.copy() + del missing_hostname_dict["hostname"] + + m = Machine(**missing_hostname_dict) + + assert m.hostname == "" From 5713d1c99b33b1fac5b2764cd813a4deb70ce5aa Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 29 Aug 2022 20:30:07 -0400 Subject: [PATCH 16/22] UT: Add test_hardware_id_default() --- .../unit_tests/monkey_island/cc/models/test_machine.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/monkey/tests/unit_tests/monkey_island/cc/models/test_machine.py b/monkey/tests/unit_tests/monkey_island/cc/models/test_machine.py index db8fb5825..49eeae9a7 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/models/test_machine.py +++ b/monkey/tests/unit_tests/monkey_island/cc/models/test_machine.py @@ -121,6 +121,15 @@ def test_hardware_id_validate_on_set(): m.hardware_id = -50 +def test_hardware_id_default(): + missing_hardware_id_dict = MACHINE_OBJECT_DICT.copy() + del missing_hardware_id_dict["hardware_id"] + + m = Machine(**missing_hardware_id_dict) + + assert m.hardware_id is None + + def test_network_interfaces_set_valid_value(): m = Machine(**MACHINE_OBJECT_DICT) From 0adf9d84677de77cb2c85915a3d3f1cb91f2c91b Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 30 Aug 2022 05:03:47 -0400 Subject: [PATCH 17/22] Island: Add MongoMachineRepository --- .../monkey_island/cc/repository/__init__.py | 1 + .../cc/repository/mongo_machine_repository.py | 99 ++++++++ .../test_mongo_machine_repository.py | 238 ++++++++++++++++++ 3 files changed, 338 insertions(+) create mode 100644 monkey/monkey_island/cc/repository/mongo_machine_repository.py create mode 100644 monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_machine_repository.py diff --git a/monkey/monkey_island/cc/repository/__init__.py b/monkey/monkey_island/cc/repository/__init__.py index d9987e677..caef77b9b 100644 --- a/monkey/monkey_island/cc/repository/__init__.py +++ b/monkey/monkey_island/cc/repository/__init__.py @@ -20,3 +20,4 @@ from .file_agent_configuration_repository import FileAgentConfigurationRepositor from .file_simulation_repository import FileSimulationRepository from .json_file_user_repository import JSONFileUserRepository from .mongo_credentials_repository import MongoCredentialsRepository +from .mongo_machine_repository import MongoMachineRepository diff --git a/monkey/monkey_island/cc/repository/mongo_machine_repository.py b/monkey/monkey_island/cc/repository/mongo_machine_repository.py new file mode 100644 index 000000000..ba638e53a --- /dev/null +++ b/monkey/monkey_island/cc/repository/mongo_machine_repository.py @@ -0,0 +1,99 @@ +from ipaddress import IPv4Address +from threading import Lock +from typing import Any, MutableMapping, Sequence + +from pymongo import MongoClient + +from common.types import HardwareID +from monkey_island.cc.models import Machine, MachineID + +from . import IMachineRepository, RetrievalError, StorageError, UnknownRecordError +from .consts import MONGO_OBJECT_ID_KEY + + +class MongoMachineRepository(IMachineRepository): + """A repository used to store and retrieve Machines in MongoDB""" + + def __init__(self, mongo_client: MongoClient): + self._machines_collection = mongo_client.monkey_island.machines + self._id_lock = Lock() + self._next_id = self._get_biggest_id() + + def _get_biggest_id(self) -> MachineID: + try: + return self._machines_collection.find().sort("id", -1).limit(1)[0]["id"] + except IndexError: + return 0 + + def create_machine(self) -> Machine: + try: + next_id = self._get_next_id() + new_machine = Machine(id=next_id) + self._machines_collection.insert_one(new_machine.dict(simplify=True)) + + return new_machine + except Exception as err: + raise StorageError(f"Error creating a new machine: {err}") + + def _get_next_id(self) -> MachineID: + with self._id_lock: + self._next_id += 1 + return self._next_id + + def update_machine(self, machine: Machine): + try: + result = self._machines_collection.replace_one( + {"id": machine.id}, machine.dict(simplify=True) + ) + except Exception as err: + raise StorageError(f'Error updating machine with ID "{machine.id}": {err}') + + if result.matched_count == 0: + raise UnknownRecordError(f"Unknown machine: id == {machine.id}") + + if result.modified_count != 1: + raise StorageError( + f'Error updating machine with ID "{machine.id}": Expected to update 1 machines, ' + f"but updated {result.modified_count}" + ) + + def get_machine_by_id(self, id_: MachineID) -> Machine: + return self._find_one("id", "machine ID", id_) + + def get_machine_by_hardware_id(self, hardware_id: HardwareID) -> Machine: + return self._find_one("hardware_id", "hardware ID", hardware_id) + + def _find_one(self, key: str, key_display_name: str, search_value: Any) -> Machine: + try: + machine_dict = self._machines_collection.find_one({key: search_value}) + except Exception as err: + raise RetrievalError(f'Error retrieving machine with "{key} == {search_value}": {err}') + + if machine_dict is None: + raise UnknownRecordError(f'Unknown {key_display_name} "{search_value}"') + + return MongoMachineRepository._mongo_record_to_machine(machine_dict) + + def get_machines_by_ip(self, ip: IPv4Address) -> Sequence[Machine]: + ip_regex = "^" + str(ip).replace(".", "\\.") + "\\/.*$" + query = {"network_interfaces": {"$elemMatch": {"$regex": ip_regex}}} + + try: + cursor = self._machines_collection.find(query) + except Exception as err: + raise RetrievalError(f'Error retrieving machines with ip "{ip}": {err}') + + machines = list(map(MongoMachineRepository._mongo_record_to_machine, cursor)) + + if len(machines) == 0: + raise UnknownRecordError(f'No machines found with IP "{ip}"') + + return machines + + @staticmethod + def _mongo_record_to_machine(mongo_record: MutableMapping[str, Any]) -> Machine: + del mongo_record[MONGO_OBJECT_ID_KEY] + return Machine(**mongo_record) + + def reset(self): + self._machines_collection.drop() diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_machine_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_machine_repository.py new file mode 100644 index 000000000..9ff00c64b --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_machine_repository.py @@ -0,0 +1,238 @@ +from ipaddress import IPv4Interface +from itertools import chain, repeat +from unittest.mock import MagicMock + +import mongomock +import pytest + +from common import OperatingSystem +from monkey_island.cc.models import Machine +from monkey_island.cc.repository import ( + IMachineRepository, + MongoMachineRepository, + RetrievalError, + StorageError, + UnknownRecordError, +) + +MACHINES = ( + Machine( + id=1, + hardware_id=12345, + network_interfaces=[IPv4Interface("192.168.1.10/24")], + operating_system=OperatingSystem.LINUX, + operating_system_version="Ubuntu 22.04", + hostname="wopr", + ), + Machine( + id=2, + hardware_id=67890, + network_interfaces=[IPv4Interface("192.168.1.11/24"), IPv4Interface("192.168.1.12/24")], + operating_system=OperatingSystem.WINDOWS, + operating_system_version="eXtra Problems", + hostname="hal", + ), + Machine( + id=3, + hardware_id=112345, + network_interfaces=[IPv4Interface("192.168.1.13/24"), IPv4Interface("192.168.1.14/24")], + operating_system=OperatingSystem.WINDOWS, + operating_system_version="Vista", + hostname="smith", + ), + Machine( + id=4, + hardware_id=167890, + network_interfaces=[IPv4Interface("192.168.1.14/24")], + operating_system=OperatingSystem.LINUX, + operating_system_version="CentOS Linux 8", + hostname="skynet", + ), +) + + +@pytest.fixture +def mongo_client() -> mongomock.MongoClient: + client = mongomock.MongoClient() + client.monkey_island.machines.insert_many((m.dict(simplify=True) for m in MACHINES)) + return client + + +@pytest.fixture +def error_raising_mock_mongo_client() -> mongomock.MongoClient: + mongo_client = MagicMock(spec=mongomock.MongoClient) + mongo_client.monkey_island = MagicMock(spec=mongomock.Database) + mongo_client.monkey_island.machines = MagicMock(spec=mongomock.Collection) + + # The first call to find() must succeed + mongo_client.monkey_island.machines.find = MagicMock( + side_effect=chain([MagicMock()], repeat(Exception("some exception"))) + ) + mongo_client.monkey_island.machines.find_one = MagicMock( + side_effect=Exception("some exception") + ) + mongo_client.monkey_island.machines.insert_one = MagicMock( + side_effect=Exception("some exception") + ) + mongo_client.monkey_island.machines.replace_one = MagicMock( + side_effect=Exception("some exception") + ) + + return mongo_client + + +@pytest.fixture +def error_raising_machine_repository(error_raising_mock_mongo_client) -> IMachineRepository: + return MongoMachineRepository(error_raising_mock_mongo_client) + + +@pytest.fixture +def machine_repository(mongo_client) -> IMachineRepository: + return MongoMachineRepository(mongo_client) + + +def test_create_machine__unique_id(machine_repository): + new_machine = machine_repository.create_machine() + + for m in MACHINES: + assert m.id != new_machine.id + + +def test_create_machine__multiple_unique_ids(machine_repository): + new_machine_1 = machine_repository.create_machine() + new_machine_2 = machine_repository.create_machine() + + assert new_machine_1.id != new_machine_2.id + + +def test_create_machine__new_id_for_empty_repo(machine_repository): + empty_machine_repository = MongoMachineRepository(mongomock.MongoClient()) + new_machine_1 = empty_machine_repository.create_machine() + new_machine_2 = empty_machine_repository.create_machine() + + assert new_machine_1.id != new_machine_2.id + + +def test_create_machine__storage_error(error_raising_machine_repository): + with pytest.raises(StorageError): + error_raising_machine_repository.create_machine() + + +def test_update_machine(machine_repository): + machine = machine_repository.get_machine_by_id(1) + + machine.operating_system = OperatingSystem.WINDOWS + machine.hostname = "viki" + machine.network_interfaces = [IPv4Interface("10.0.0.1/16")] + + machine_repository.update_machine(machine) + + assert machine_repository.get_machine_by_id(1) == machine + + +def test_update_machine__not_found(machine_repository): + machine = Machine(id=99) + + with pytest.raises(UnknownRecordError): + machine_repository.update_machine(machine) + + +def test_update_machine__storage_error_exception(error_raising_machine_repository): + machine = MACHINES[0] + + with pytest.raises(StorageError): + error_raising_machine_repository.update_machine(machine) + + +def test_update_machine__storage_error_update_failed(error_raising_mock_mongo_client): + mock_result = MagicMock() + mock_result.matched_count = 1 + mock_result.modified_count = 0 + + error_raising_mock_mongo_client.monkey_island.machines.replace_one = MagicMock( + return_value=mock_result + ) + machine_repository = MongoMachineRepository(error_raising_mock_mongo_client) + + machine = MACHINES[0] + with pytest.raises(StorageError): + machine_repository.update_machine(machine) + + +def test_get_machine_by_id(machine_repository): + for i, expected_machine in enumerate(MACHINES, start=1): + assert machine_repository.get_machine_by_id(i) == expected_machine + + +def test_get_machine_by_id__not_found(machine_repository): + with pytest.raises(UnknownRecordError): + machine_repository.get_machine_by_id(9999) + + +def test_get_machine_by_id__retrieval_error(error_raising_machine_repository): + with pytest.raises(RetrievalError): + error_raising_machine_repository.get_machine_by_id(1) + + +def test_get_machine_by_hardware_id(machine_repository): + for hardware_id, expected_machine in ((machine.hardware_id, machine) for machine in MACHINES): + assert machine_repository.get_machine_by_hardware_id(hardware_id) == expected_machine + + +def test_get_machine_by_hardware_id__not_found(machine_repository): + with pytest.raises(UnknownRecordError): + machine_repository.get_machine_by_hardware_id(9999888887777) + + +def test_get_machine_by_hardware_id__retrieval_error(error_raising_machine_repository): + with pytest.raises(RetrievalError): + error_raising_machine_repository.get_machine_by_hardware_id(1) + + +def test_get_machine_by_ip(machine_repository): + expected_machine = MACHINES[0] + expected_machine_ip = expected_machine.network_interfaces[0].ip + + retrieved_machines = machine_repository.get_machines_by_ip(expected_machine_ip) + + assert len(retrieved_machines) == 1 + assert retrieved_machines[0] == expected_machine + + +def test_get_machine_by_ip__multiple_results(machine_repository): + search_ip = MACHINES[3].network_interfaces[0].ip + + retrieved_machines = machine_repository.get_machines_by_ip(search_ip) + + assert len(retrieved_machines) == 2 + assert MACHINES[2] in retrieved_machines + assert MACHINES[3] in retrieved_machines + + +def test_get_machine_by_ip__not_found(machine_repository): + with pytest.raises(UnknownRecordError): + machine_repository.get_machines_by_ip("1.1.1.1") + + +def test_get_machine_by_ip__retrieval_error(error_raising_machine_repository): + with pytest.raises(RetrievalError): + error_raising_machine_repository.get_machines_by_ip("1.1.1.1") + + +def test_reset(machine_repository): + # Ensure the repository is not empty + preexisting_machine = machine_repository.get_machine_by_id(MACHINES[0].id) + assert isinstance(preexisting_machine, Machine) + + machine_repository.reset() + + with pytest.raises(UnknownRecordError): + machine_repository.get_machine_by_id(MACHINES[0].id) + + +def test_usable_after_reset(machine_repository): + machine_repository.reset() + + new_machine = machine_repository.create_machine() + + assert new_machine == machine_repository.get_machine_by_id(new_machine.id) From 383cfdfefe83de67909e54606154261f74d5d304 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 30 Aug 2022 07:30:41 -0400 Subject: [PATCH 18/22] Island: Rename `id_` parameter to `machine_id` in IMachineRepository --- monkey/monkey_island/cc/repository/i_machine_repository.py | 4 ++-- .../monkey_island/cc/repository/mongo_machine_repository.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/monkey/monkey_island/cc/repository/i_machine_repository.py b/monkey/monkey_island/cc/repository/i_machine_repository.py index 97bd3cf87..6d6832669 100644 --- a/monkey/monkey_island/cc/repository/i_machine_repository.py +++ b/monkey/monkey_island/cc/repository/i_machine_repository.py @@ -29,11 +29,11 @@ class IMachineRepository(ABC): """ @abstractmethod - def get_machine_by_id(self, id_: MachineID) -> Machine: + def get_machine_by_id(self, machine_id: MachineID) -> Machine: """ Get a `Machine` by ID - :param id: The ID of the `Machine` to be retrieved + :param machine_id: The ID of the `Machine` to be retrieved :return: A `Machine` with a matching `id` :raises UnknownRecordError: If a `Machine` with the specified `id` does not exist in the repository diff --git a/monkey/monkey_island/cc/repository/mongo_machine_repository.py b/monkey/monkey_island/cc/repository/mongo_machine_repository.py index ba638e53a..e3d59169a 100644 --- a/monkey/monkey_island/cc/repository/mongo_machine_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_machine_repository.py @@ -57,8 +57,8 @@ class MongoMachineRepository(IMachineRepository): f"but updated {result.modified_count}" ) - def get_machine_by_id(self, id_: MachineID) -> Machine: - return self._find_one("id", "machine ID", id_) + def get_machine_by_id(self, machine_id: MachineID) -> Machine: + return self._find_one("id", "machine ID", machine_id) def get_machine_by_hardware_id(self, hardware_id: HardwareID) -> Machine: return self._find_one("hardware_id", "hardware ID", hardware_id) From bf5e54ebc9b164e7b389ab4d002514c8aaa4fd13 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 30 Aug 2022 07:34:20 -0400 Subject: [PATCH 19/22] Island: Raise RemovalError from IMachineRepository.reset() --- monkey/monkey_island/cc/repository/i_machine_repository.py | 3 +++ .../cc/repository/mongo_machine_repository.py | 7 +++++-- .../cc/repository/test_mongo_machine_repository.py | 7 +++++++ 3 files changed, 15 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/repository/i_machine_repository.py b/monkey/monkey_island/cc/repository/i_machine_repository.py index 6d6832669..568564009 100644 --- a/monkey/monkey_island/cc/repository/i_machine_repository.py +++ b/monkey/monkey_island/cc/repository/i_machine_repository.py @@ -68,4 +68,7 @@ class IMachineRepository(ABC): def reset(self): """ Removes all data from the repository + + :raises RemovalError: If an error occurred while attempting to remove all `Machines` from + the repository """ diff --git a/monkey/monkey_island/cc/repository/mongo_machine_repository.py b/monkey/monkey_island/cc/repository/mongo_machine_repository.py index e3d59169a..318153e51 100644 --- a/monkey/monkey_island/cc/repository/mongo_machine_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_machine_repository.py @@ -7,7 +7,7 @@ from pymongo import MongoClient from common.types import HardwareID from monkey_island.cc.models import Machine, MachineID -from . import IMachineRepository, RetrievalError, StorageError, UnknownRecordError +from . import IMachineRepository, RemovalError, RetrievalError, StorageError, UnknownRecordError from .consts import MONGO_OBJECT_ID_KEY @@ -96,4 +96,7 @@ class MongoMachineRepository(IMachineRepository): return Machine(**mongo_record) def reset(self): - self._machines_collection.drop() + try: + self._machines_collection.drop() + except Exception as err: + raise RemovalError(f"Error resetting the repository: {err}") diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_machine_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_machine_repository.py index 9ff00c64b..7851f8e20 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_machine_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_machine_repository.py @@ -10,6 +10,7 @@ from monkey_island.cc.models import Machine from monkey_island.cc.repository import ( IMachineRepository, MongoMachineRepository, + RemovalError, RetrievalError, StorageError, UnknownRecordError, @@ -77,6 +78,7 @@ def error_raising_mock_mongo_client() -> mongomock.MongoClient: mongo_client.monkey_island.machines.replace_one = MagicMock( side_effect=Exception("some exception") ) + mongo_client.monkey_island.machines.drop = MagicMock(side_effect=Exception("some exception")) return mongo_client @@ -236,3 +238,8 @@ def test_usable_after_reset(machine_repository): new_machine = machine_repository.create_machine() assert new_machine == machine_repository.get_machine_by_id(new_machine.id) + + +def test_reset__removal_error(error_raising_machine_repository): + with pytest.raises(RemovalError): + error_raising_machine_repository.reset() From b538842e84d87fec315b3386c0fa6996e1d989c4 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 30 Aug 2022 09:32:47 -0400 Subject: [PATCH 20/22] Island: Remove display_name from MongoMachineRepository._find_one() --- .../cc/repository/mongo_machine_repository.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/monkey/monkey_island/cc/repository/mongo_machine_repository.py b/monkey/monkey_island/cc/repository/mongo_machine_repository.py index 318153e51..cea459c25 100644 --- a/monkey/monkey_island/cc/repository/mongo_machine_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_machine_repository.py @@ -58,19 +58,19 @@ class MongoMachineRepository(IMachineRepository): ) def get_machine_by_id(self, machine_id: MachineID) -> Machine: - return self._find_one("id", "machine ID", machine_id) + return self._find_one("id", machine_id) def get_machine_by_hardware_id(self, hardware_id: HardwareID) -> Machine: - return self._find_one("hardware_id", "hardware ID", hardware_id) + return self._find_one("hardware_id", hardware_id) - def _find_one(self, key: str, key_display_name: str, search_value: Any) -> Machine: + def _find_one(self, key: str, search_value: Any) -> Machine: try: machine_dict = self._machines_collection.find_one({key: search_value}) except Exception as err: raise RetrievalError(f'Error retrieving machine with "{key} == {search_value}": {err}') if machine_dict is None: - raise UnknownRecordError(f'Unknown {key_display_name} "{search_value}"') + raise UnknownRecordError(f'Unknown machine with "{key} == {search_value}"') return MongoMachineRepository._mongo_record_to_machine(machine_dict) From 81128a4842ced999f229cc6ad45923e7892846aa Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 30 Aug 2022 09:34:17 -0400 Subject: [PATCH 21/22] Island: Don't use Field() for simple defaults in Machine --- monkey/monkey_island/cc/models/machine.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/monkey/monkey_island/cc/models/machine.py b/monkey/monkey_island/cc/models/machine.py index 4c5925390..5a8c5e12b 100644 --- a/monkey/monkey_island/cc/models/machine.py +++ b/monkey/monkey_island/cc/models/machine.py @@ -20,16 +20,16 @@ class Machine(MutableInfectionMonkeyBaseModel): hardware_id: Optional[HardwareID] """An identifier generated by the agent that uniquely identifies a machine""" - network_interfaces: Sequence[IPv4Interface] = Field(default_factory=tuple) + network_interfaces: Sequence[IPv4Interface] = tuple() """The machine's networking interfaces""" operating_system: Optional[OperatingSystem] """The operating system the machine is running""" - operating_system_version: str = Field(default="") + operating_system_version: str = "" """The specific version of the operating system the machine is running""" - hostname: str = Field(default="") + hostname: str = "" """The hostname of the machine""" _make_immutable_sequence = validator("network_interfaces", pre=True, allow_reuse=True)( From ba7dab26d7d1ab951602cf64ac3f09fd457ac546 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 31 Aug 2022 10:14:57 -0400 Subject: [PATCH 22/22] Island: Refactor how Machine objects are managed by IMachineRepository - Replace `create_machine()` with `get_new_id()` - Replace `update_machine()` with `upsert_machine()` Benefits: The repository doesn't store Machine objects that only have the ID populated (unless that is the caller's desire). Upsert instead of update allows the interface to be more permissive. --- .../cc/repository/i_machine_repository.py | 17 ++--- .../cc/repository/mongo_machine_repository.py | 31 ++++----- .../test_mongo_machine_repository.py | 67 +++++++++++-------- vulture_allowlist.py | 4 +- 4 files changed, 63 insertions(+), 56 deletions(-) diff --git a/monkey/monkey_island/cc/repository/i_machine_repository.py b/monkey/monkey_island/cc/repository/i_machine_repository.py index 568564009..d60346eb2 100644 --- a/monkey/monkey_island/cc/repository/i_machine_repository.py +++ b/monkey/monkey_island/cc/repository/i_machine_repository.py @@ -10,21 +10,22 @@ class IMachineRepository(ABC): """A repository used to store and retrieve Machines""" @abstractmethod - def create_machine(self) -> Machine: + def get_new_id(self) -> MachineID: """ - Create a new `Machine` in the repository + Generates a new, unique `MachineID` - :return: A new `Machine` with a unique ID - :raises StorageError: If a new `Machine` could not be created + :return: A new, unique `MachineID` """ @abstractmethod - def update_machine(self, machine: Machine): + def upsert_machine(self, machine: Machine): """ - Update an existing `Machine` in the repository + Upsert (insert or update) a `Machine` - :param machine: An updated Machine object to store in the repository - :raises UnknownRecordError: If the provided `Machine` does not exist in the repository + Insert the `Machine` if no `Machine` with a matching ID exists in the repository. If the + `Machine` already exists, update it. + + :param machine: The `Machine` to be inserted or updated :raises StorageError: If an error occurred while attempting to store the `Machine` """ diff --git a/monkey/monkey_island/cc/repository/mongo_machine_repository.py b/monkey/monkey_island/cc/repository/mongo_machine_repository.py index cea459c25..4d1a36470 100644 --- a/monkey/monkey_island/cc/repository/mongo_machine_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_machine_repository.py @@ -25,36 +25,29 @@ class MongoMachineRepository(IMachineRepository): except IndexError: return 0 - def create_machine(self) -> Machine: - try: - next_id = self._get_next_id() - new_machine = Machine(id=next_id) - self._machines_collection.insert_one(new_machine.dict(simplify=True)) - - return new_machine - except Exception as err: - raise StorageError(f"Error creating a new machine: {err}") - - def _get_next_id(self) -> MachineID: + def get_new_id(self) -> MachineID: with self._id_lock: self._next_id += 1 return self._next_id - def update_machine(self, machine: Machine): + def upsert_machine(self, machine: Machine): try: result = self._machines_collection.replace_one( - {"id": machine.id}, machine.dict(simplify=True) + {"id": machine.id}, machine.dict(simplify=True), upsert=True ) except Exception as err: raise StorageError(f'Error updating machine with ID "{machine.id}": {err}') - if result.matched_count == 0: - raise UnknownRecordError(f"Unknown machine: id == {machine.id}") - - if result.modified_count != 1: + if result.matched_count != 0 and result.modified_count != 1: raise StorageError( - f'Error updating machine with ID "{machine.id}": Expected to update 1 machines, ' - f"but updated {result.modified_count}" + f'Error updating machine with ID "{machine.id}": Expected to update 1 machine, ' + f"but {result.modified_count} were updated" + ) + + if result.matched_count == 0 and result.upserted_id is None: + raise StorageError( + f'Error inserting machine with ID "{machine.id}": Expected to insert 1 machine, ' + f"but no machines were inserted" ) def get_machine_by_id(self, machine_id: MachineID) -> Machine: diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_machine_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_machine_repository.py index 7851f8e20..90d0af1f2 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_machine_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_machine_repository.py @@ -93,60 +93,56 @@ def machine_repository(mongo_client) -> IMachineRepository: return MongoMachineRepository(mongo_client) -def test_create_machine__unique_id(machine_repository): - new_machine = machine_repository.create_machine() +def test_get_new_id__unique_id(machine_repository): + new_machine_id = machine_repository.get_new_id() for m in MACHINES: - assert m.id != new_machine.id + assert m.id != new_machine_id -def test_create_machine__multiple_unique_ids(machine_repository): - new_machine_1 = machine_repository.create_machine() - new_machine_2 = machine_repository.create_machine() +def test_get_new_id__multiple_unique_ids(machine_repository): + id_1 = machine_repository.get_new_id() + id_2 = machine_repository.get_new_id() - assert new_machine_1.id != new_machine_2.id + assert id_1 != id_2 -def test_create_machine__new_id_for_empty_repo(machine_repository): +def test_get_new_id__new_id_for_empty_repo(machine_repository): empty_machine_repository = MongoMachineRepository(mongomock.MongoClient()) - new_machine_1 = empty_machine_repository.create_machine() - new_machine_2 = empty_machine_repository.create_machine() + id_1 = empty_machine_repository.get_new_id() + id_2 = empty_machine_repository.get_new_id() - assert new_machine_1.id != new_machine_2.id + assert id_1 != id_2 -def test_create_machine__storage_error(error_raising_machine_repository): - with pytest.raises(StorageError): - error_raising_machine_repository.create_machine() - - -def test_update_machine(machine_repository): +def test_upsert_machine__update(machine_repository): machine = machine_repository.get_machine_by_id(1) machine.operating_system = OperatingSystem.WINDOWS machine.hostname = "viki" machine.network_interfaces = [IPv4Interface("10.0.0.1/16")] - machine_repository.update_machine(machine) + machine_repository.upsert_machine(machine) assert machine_repository.get_machine_by_id(1) == machine -def test_update_machine__not_found(machine_repository): - machine = Machine(id=99) +def test_upsert_machine__insert(machine_repository): + new_machine = Machine(id=99, hardware_id=8675309) - with pytest.raises(UnknownRecordError): - machine_repository.update_machine(machine) + machine_repository.upsert_machine(new_machine) + + assert machine_repository.get_machine_by_id(99) == new_machine -def test_update_machine__storage_error_exception(error_raising_machine_repository): +def test_upsert_machine__storage_error_exception(error_raising_machine_repository): machine = MACHINES[0] with pytest.raises(StorageError): - error_raising_machine_repository.update_machine(machine) + error_raising_machine_repository.upsert_machine(machine) -def test_update_machine__storage_error_update_failed(error_raising_mock_mongo_client): +def test_upsert_machine__storage_error_update_failed(error_raising_mock_mongo_client): mock_result = MagicMock() mock_result.matched_count = 1 mock_result.modified_count = 0 @@ -158,7 +154,22 @@ def test_update_machine__storage_error_update_failed(error_raising_mock_mongo_cl machine = MACHINES[0] with pytest.raises(StorageError): - machine_repository.update_machine(machine) + machine_repository.upsert_machine(machine) + + +def test_upsert_machine__storage_error_insert_failed(error_raising_mock_mongo_client): + mock_result = MagicMock() + mock_result.matched_count = 0 + mock_result.upserted_id = None + + error_raising_mock_mongo_client.monkey_island.machines.replace_one = MagicMock( + return_value=mock_result + ) + machine_repository = MongoMachineRepository(error_raising_mock_mongo_client) + + machine = MACHINES[0] + with pytest.raises(StorageError): + machine_repository.upsert_machine(machine) def test_get_machine_by_id(machine_repository): @@ -235,7 +246,9 @@ def test_reset(machine_repository): def test_usable_after_reset(machine_repository): machine_repository.reset() - new_machine = machine_repository.create_machine() + new_id = machine_repository.get_new_id() + new_machine = Machine(id=new_id) + machine_repository.upsert_machine(new_machine) assert new_machine == machine_repository.get_machine_by_id(new_machine.id) diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 9da5fee84..688d2f123 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -244,8 +244,8 @@ IConfigRepository.get_config_field ILogRepository.get_logs ILogRepository.save_log ILogRepository.delete_log -IMachineRepository.create_machine -IMachineRepository.update_machine +IMachineRepository.get_new_id +IMachineRepository.upsert_machine IMachineRepository.get_machine_by_id IMachineRepository.get_machine_by_hardware_id IMachineRepository.get_machines_by_ip