From 0516e1e0157c9907405f79dc8fe87fda3476c3c7 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Wed, 28 Sep 2022 16:20:46 +0000 Subject: [PATCH 1/7] Island: Add get_machines to IMachineRepository --- .../cc/repository/i_machine_repository.py | 9 +++++ .../cc/repository/mongo_machine_repository.py | 8 +++++ .../test_mongo_machine_repository.py | 35 +++++++++++++++---- 3 files changed, 46 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 85e8e71f3..7cea0bb02 100644 --- a/monkey/monkey_island/cc/repository/i_machine_repository.py +++ b/monkey/monkey_island/cc/repository/i_machine_repository.py @@ -53,6 +53,15 @@ class IMachineRepository(ABC): :raises RetrievalError: If an error occurs while attempting to retrieve the `Machine` """ + @abstractmethod + def get_machines(self) -> Sequence[Machine]: + """ + Get all machines in the repository + + :return: A sequence of all stored `Machine`s + :raises RetrievalError: If an error occurs while attempting to retrieve the `Machine`s + """ + @abstractmethod def get_machines_by_ip(self, ip: IPv4Address) -> Sequence[Machine]: """ diff --git a/monkey/monkey_island/cc/repository/mongo_machine_repository.py b/monkey/monkey_island/cc/repository/mongo_machine_repository.py index fab038694..75cbf2d46 100644 --- a/monkey/monkey_island/cc/repository/mongo_machine_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_machine_repository.py @@ -69,6 +69,14 @@ class MongoMachineRepository(IMachineRepository): return Machine(**machine_dict) + def get_machines(self) -> Sequence[Machine]: + try: + cursor = self._machines_collection.find({}, {MONGO_OBJECT_ID_KEY: False}) + except Exception as err: + raise RetrievalError(f"Error retrieving machines: {err}") + + return list(map(lambda m: Machine(**m), cursor)) + def get_machines_by_ip(self, ip: IPv4Address) -> Sequence[Machine]: ip_regex = "^" + str(ip).replace(".", "\\.") + "\\/.*$" query = {"network_interfaces": {"$elemMatch": {"$regex": ip_regex}}} 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 90d0af1f2..bbd35283b 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,6 +93,11 @@ def machine_repository(mongo_client) -> IMachineRepository: return MongoMachineRepository(mongo_client) +@pytest.fixture +def empty_machine_repository() -> IMachineRepository: + return MongoMachineRepository(mongomock.MongoClient()) + + def test_get_new_id__unique_id(machine_repository): new_machine_id = machine_repository.get_new_id() @@ -107,8 +112,7 @@ def test_get_new_id__multiple_unique_ids(machine_repository): assert id_1 != id_2 -def test_get_new_id__new_id_for_empty_repo(machine_repository): - empty_machine_repository = MongoMachineRepository(mongomock.MongoClient()) +def test_get_new_id__new_id_for_empty_repo(empty_machine_repository): id_1 = empty_machine_repository.get_new_id() id_2 = empty_machine_repository.get_new_id() @@ -202,7 +206,7 @@ def test_get_machine_by_hardware_id__retrieval_error(error_raising_machine_repos error_raising_machine_repository.get_machine_by_hardware_id(1) -def test_get_machine_by_ip(machine_repository): +def test_get_machines_by_ip(machine_repository): expected_machine = MACHINES[0] expected_machine_ip = expected_machine.network_interfaces[0].ip @@ -212,7 +216,7 @@ def test_get_machine_by_ip(machine_repository): assert retrieved_machines[0] == expected_machine -def test_get_machine_by_ip__multiple_results(machine_repository): +def test_get_machines_by_ip__multiple_results(machine_repository): search_ip = MACHINES[3].network_interfaces[0].ip retrieved_machines = machine_repository.get_machines_by_ip(search_ip) @@ -222,16 +226,35 @@ def test_get_machine_by_ip__multiple_results(machine_repository): assert MACHINES[3] in retrieved_machines -def test_get_machine_by_ip__not_found(machine_repository): +def test_get_machines_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): +def test_get_machines_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_get_machines(machine_repository): + retrieved_machines = machine_repository.get_machines() + + assert len(retrieved_machines) == len(MACHINES) + for machine in MACHINES: + assert machine in retrieved_machines + + +def test_get_machines__empty_repository(empty_machine_repository): + retrieved_machines = empty_machine_repository.get_machines() + + assert len(retrieved_machines) == 0 + + +def test_get_machines__retrieval_error(error_raising_machine_repository): + with pytest.raises(RetrievalError): + error_raising_machine_repository.get_machines() + + def test_reset(machine_repository): # Ensure the repository is not empty preexisting_machine = machine_repository.get_machine_by_id(MACHINES[0].id) From eeca5fbea2486cd12a9356c5940e88cf2206e576 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Fri, 30 Sep 2022 17:18:35 +0000 Subject: [PATCH 2/7] Island: Add resource for /api/machines endpoint --- monkey/monkey_island/cc/resources/__init__.py | 1 + monkey/monkey_island/cc/resources/machines.py | 16 ++++++++++++++++ 2 files changed, 17 insertions(+) create mode 100644 monkey/monkey_island/cc/resources/machines.py diff --git a/monkey/monkey_island/cc/resources/__init__.py b/monkey/monkey_island/cc/resources/__init__.py index e6c2db648..5802d35e6 100644 --- a/monkey/monkey_island/cc/resources/__init__.py +++ b/monkey/monkey_island/cc/resources/__init__.py @@ -12,3 +12,4 @@ from .agent_events import AgentEvents from .agents import Agents from .agent_signals import AgentSignals, TerminateAllAgents from .agent_logs import AgentLogs +from .machines import Machines diff --git a/monkey/monkey_island/cc/resources/machines.py b/monkey/monkey_island/cc/resources/machines.py new file mode 100644 index 000000000..3e1c1edb0 --- /dev/null +++ b/monkey/monkey_island/cc/resources/machines.py @@ -0,0 +1,16 @@ +from http import HTTPStatus + +from monkey_island.cc.repository import IMachineRepository +from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.resources.request_authentication import jwt_required + + +class Machines(AbstractResource): + urls = ["/api/machines"] + + def __init__(self, machine_repository: IMachineRepository): + self._machine_repository = machine_repository + + @jwt_required + def get(self): + return self._machine_repository.get_machines(), HTTPStatus.OK From f05f247417436f8b939ce8d04075818c81fa0dee Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Fri, 30 Sep 2022 17:19:36 +0000 Subject: [PATCH 3/7] Island: Hook up the /api/machines endpoint --- monkey/monkey_island/cc/app.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/monkey/monkey_island/cc/app.py b/monkey/monkey_island/cc/app.py index 5e0aaadb2..18af5f2a6 100644 --- a/monkey/monkey_island/cc/app.py +++ b/monkey/monkey_island/cc/app.py @@ -20,6 +20,7 @@ from monkey_island.cc.resources import ( ClearSimulationData, IPAddresses, IslandLog, + Machines, PBAFileDownload, PBAFileUpload, PropagationCredentials, @@ -173,6 +174,7 @@ def init_restful_endpoints(api: FlaskDIWrapper): api.add_resource(AgentBinaries) api.add_resource(NetMap) api.add_resource(Edge) + api.add_resource(Machines) api.add_resource(Node) api.add_resource(NodeStates) From a3d2d7f6a1f216eff71f103acb9ce50d96bf9aaf Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Fri, 30 Sep 2022 17:25:12 +0000 Subject: [PATCH 4/7] UT: Add tests for Machines resource --- .../cc/resources/test_machines.py | 80 +++++++++++++++++++ 1 file changed, 80 insertions(+) create mode 100644 monkey/tests/unit_tests/monkey_island/cc/resources/test_machines.py diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_machines.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_machines.py new file mode 100644 index 000000000..7aab32a1c --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_machines.py @@ -0,0 +1,80 @@ +from http import HTTPStatus +from ipaddress import IPv4Interface +from unittest.mock import MagicMock + +import pytest +from tests.common import StubDIContainer +from tests.unit_tests.monkey_island.conftest import get_url_for_resource + +from common import OperatingSystem +from monkey_island.cc.models import Machine +from monkey_island.cc.repository import IMachineRepository, RetrievalError +from monkey_island.cc.resources import Machines + +MACHINES_URL = get_url_for_resource(Machines) +MACHINES = [ + Machine( + id=1, + hardware_id=101, + island=True, + network_interfaces=[IPv4Interface("10.10.10.1")], + operating_system=OperatingSystem.WINDOWS, + ), + Machine( + id=2, + hardware_id=102, + island=False, + network_interfaces=[IPv4Interface("10.10.10.2/24")], + operating_system=OperatingSystem.LINUX, + ), +] + + +@pytest.fixture +def mock_machine_repository() -> IMachineRepository: + machine_repository = MagicMock(spec=IMachineRepository) + return machine_repository + + +@pytest.fixture +def flask_client(build_flask_client, mock_machine_repository): + container = StubDIContainer() + + container.register_instance(IMachineRepository, mock_machine_repository) + + with build_flask_client(container) as flask_client: + yield flask_client + + +def test_machines_get__status_ok(flask_client, mock_machine_repository): + mock_machine_repository.get_machines = MagicMock(return_value=MACHINES) + + resp = flask_client.get(MACHINES_URL) + + assert resp.status_code == HTTPStatus.OK + + +def test_machines_get__gets_machines(flask_client, mock_machine_repository): + mock_machine_repository.get_machines = MagicMock(return_value=MACHINES) + + resp = flask_client.get(MACHINES_URL) + + response_machines = [Machine(**m) for m in resp.json] + for machine in response_machines: + assert machine in MACHINES + + +def test_machines_get__gets_no_machines(flask_client, mock_machine_repository): + mock_machine_repository.get_machines = MagicMock(return_value=[]) + + resp = flask_client.get(MACHINES_URL) + + assert resp.json == [] + + +def test_machines_get__retrieval_error(flask_client, mock_machine_repository): + mock_machine_repository.get_machines = MagicMock(side_effect=RetrievalError) + + resp = flask_client.get(MACHINES_URL) + + assert resp.status_code == HTTPStatus.INTERNAL_SERVER_ERROR From 3409234a4d248f7077ff11c32c2f37bc813017c2 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Fri, 30 Sep 2022 17:25:49 +0000 Subject: [PATCH 5/7] UT: Address mypy errors due to get_url_for_resource --- monkey/tests/unit_tests/monkey_island/conftest.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/conftest.py b/monkey/tests/unit_tests/monkey_island/conftest.py index 094bbb2aa..6b47ccc34 100644 --- a/monkey/tests/unit_tests/monkey_island/conftest.py +++ b/monkey/tests/unit_tests/monkey_island/conftest.py @@ -1,7 +1,7 @@ import os import re from collections.abc import Callable -from typing import Set +from typing import Set, Type import flask_restful import pytest @@ -47,7 +47,7 @@ def mock_flask_resource_manager(container): return flask_resource_manager -def get_url_for_resource(resource: AbstractResource, **kwargs): +def get_url_for_resource(resource: Type[AbstractResource], **kwargs): chosen_url = None for url in resource.urls: if _get_url_keywords(url) == set(kwargs.keys()): From a2a6934a49c81c47ffcee3a775197f0a658faf15 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Fri, 30 Sep 2022 17:35:48 +0000 Subject: [PATCH 6/7] Changelog: Add entry for /api/machines --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8aeb5c3c4..1e33450fa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,6 +26,7 @@ Changelog](https://keepachangelog.com/en/1.0.0/). - `/api/agents` endpoint. - `/api/agent-signals` endpoint. #2261 - `/api/agent-logs/` endpoint. #2274 +- `/api/machines` endpoint. #2362 ### Changed - Reset workflow. Now it's possible to delete data gathered by agents without From df1baeebe05dcab220e9b6098c2e0fa668c093cc Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 30 Sep 2022 14:22:02 -0400 Subject: [PATCH 7/7] Island: Use list comprehension instead of map() --- .../monkey_island/cc/repository/mongo_machine_repository.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/repository/mongo_machine_repository.py b/monkey/monkey_island/cc/repository/mongo_machine_repository.py index 75cbf2d46..8e1934cc7 100644 --- a/monkey/monkey_island/cc/repository/mongo_machine_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_machine_repository.py @@ -75,7 +75,7 @@ class MongoMachineRepository(IMachineRepository): except Exception as err: raise RetrievalError(f"Error retrieving machines: {err}") - return list(map(lambda m: Machine(**m), cursor)) + return [Machine(**m) for m in cursor] def get_machines_by_ip(self, ip: IPv4Address) -> Sequence[Machine]: ip_regex = "^" + str(ip).replace(".", "\\.") + "\\/.*$" @@ -86,7 +86,7 @@ class MongoMachineRepository(IMachineRepository): except Exception as err: raise RetrievalError(f'Error retrieving machines with ip "{ip}": {err}') - machines = list(map(lambda m: Machine(**m), cursor)) + machines = [Machine(**m) for m in cursor] if len(machines) == 0: raise UnknownRecordError(f'No machines found with IP "{ip}"')