diff --git a/CHANGELOG.md b/CHANGELOG.md index 8aeb5c3c4..bd233801a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,7 +23,7 @@ Changelog](https://keepachangelog.com/en/1.0.0/). - `/api/agent-events` endpoint. #2155, #2300 - The ability to customize the file extension used by ransomware when encrypting files. #1242 -- `/api/agents` endpoint. +- `/api/agents` endpoint. #2362 - `/api/agent-signals` endpoint. #2261 - `/api/agent-logs/` endpoint. #2274 diff --git a/monkey/monkey_island/cc/repository/i_agent_repository.py b/monkey/monkey_island/cc/repository/i_agent_repository.py index 619738ab5..b56745222 100644 --- a/monkey/monkey_island/cc/repository/i_agent_repository.py +++ b/monkey/monkey_island/cc/repository/i_agent_repository.py @@ -19,6 +19,15 @@ class IAgentRepository(ABC): :raises StorageError: If an error occurs while attempting to store the `Agent` """ + @abstractmethod + def get_agents(self) -> Sequence[Agent]: + """ + Get all `Agents` stored in the repository + + :return: All agents in the repository + :raises RetrievalError: If an error occurs while attempting to retrieve the `Agents` + """ + @abstractmethod def get_agent_by_id(self, agent_id: AgentID) -> Agent: """ diff --git a/monkey/monkey_island/cc/repository/mongo_agent_repository.py b/monkey/monkey_island/cc/repository/mongo_agent_repository.py index 021cf34bc..527426ee3 100644 --- a/monkey/monkey_island/cc/repository/mongo_agent_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_agent_repository.py @@ -38,6 +38,14 @@ class MongoAgentRepository(IAgentRepository): f"but no agents were inserted" ) + def get_agents(self) -> Sequence[Agent]: + try: + cursor = self._agents_collection.find({}, {MONGO_OBJECT_ID_KEY: False}) + except Exception as err: + raise RetrievalError(f"Error retrieving agents: {err}") + + return [Agent(**a) for a in cursor] + def get_agent_by_id(self, agent_id: AgentID) -> Agent: try: agent_dict = self._agents_collection.find_one( @@ -54,7 +62,7 @@ class MongoAgentRepository(IAgentRepository): def get_running_agents(self) -> Sequence[Agent]: try: cursor = self._agents_collection.find({"stop_time": None}, {MONGO_OBJECT_ID_KEY: False}) - return list(map(lambda a: Agent(**a), cursor)) + return [Agent(**a) for a in cursor] except Exception as err: raise RetrievalError(f"Error retrieving running agents: {err}") diff --git a/monkey/monkey_island/cc/resources/agents.py b/monkey/monkey_island/cc/resources/agents.py index 9d6f701ec..84cb4f971 100644 --- a/monkey/monkey_island/cc/resources/agents.py +++ b/monkey/monkey_island/cc/resources/agents.py @@ -6,6 +6,7 @@ from flask import make_response, request from common import AgentRegistrationData from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic +from monkey_island.cc.repository import IAgentRepository from monkey_island.cc.resources.AbstractResource import AbstractResource logger = logging.getLogger(__name__) @@ -14,8 +15,12 @@ logger = logging.getLogger(__name__) class Agents(AbstractResource): urls = ["/api/agents"] - def __init__(self, island_event_queue: IIslandEventQueue): + def __init__(self, island_event_queue: IIslandEventQueue, agent_repository: IAgentRepository): self._island_event_queue = island_event_queue + self._agent_repository = agent_repository + + def get(self): + return self._agent_repository.get_agents(), HTTPStatus.OK def post(self): try: diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_repository.py index 62ebaed1c..51999d1c3 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_repository.py @@ -161,6 +161,26 @@ def test_upsert_agent__storage_error_insert_failed(error_raising_mock_mongo_clie agent_repository.upsert_agent(agent) +def test_get_agents__empty_repo(empty_agent_repository): + all_agents = empty_agent_repository.get_agents() + + assert len(all_agents) == 0 + + +def test_get_agents(agent_repository): + all_agents = agent_repository.get_agents() + + assert len(all_agents) == len(AGENTS) + + for agent in AGENTS: + assert agent in all_agents + + +def test_get_agents__retrieval_error(error_raising_agent_repository): + with pytest.raises(RetrievalError): + error_raising_agent_repository.get_agents() + + def test_get_agent_by_id(agent_repository): for i, expected_agent in enumerate(AGENTS): assert agent_repository.get_agent_by_id(expected_agent.id) == expected_agent diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agents.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agents.py index 97a9163de..605c8ecec 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agents.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agents.py @@ -6,11 +6,12 @@ import pytest from tests.common import StubDIContainer from tests.unit_tests.monkey_island.conftest import get_url_for_resource +from common.types import SocketAddress from monkey_island.cc.event_queue import IIslandEventQueue +from monkey_island.cc.models import Agent +from monkey_island.cc.repository import IAgentRepository from monkey_island.cc.resources import Agents -AGENTS_URL = get_url_for_resource(Agents) - AGENT_REGISTRATION_DICT = { "id": UUID("6bfd8b64-43d8-4449-8c70-d898aca74ad8"), "machine_hardware_id": 1, @@ -20,11 +21,37 @@ AGENT_REGISTRATION_DICT = { "network_interfaces": ["10.1.1.2/24"], } +AGENTS = ( + Agent( + id=UUID("12345678-1234-1234-1234-123456789abc"), + machine_id=2, + start_time=0, + stop_time=10, + cc_server=SocketAddress(ip="10.0.0.1", port=5000), + ), + Agent( + id=UUID("abcdef78-abcd-abcd-abcd-abcdef123456"), + machine_id=3, + start_time=5, + stop_time=15, + cc_server=SocketAddress(ip="10.0.0.1", port=5000), + ), +) + @pytest.fixture -def flask_client(build_flask_client): +def agent_repository() -> IAgentRepository: + agent_repository = MagicMock(spec=IAgentRepository) + agent_repository.get_agents = MagicMock(return_value=AGENTS) + + return agent_repository + + +@pytest.fixture +def flask_client(build_flask_client, agent_repository): container = StubDIContainer() container.register_instance(IIslandEventQueue, MagicMock(spec=IIslandEventQueue)) + container.register_instance(IAgentRepository, agent_repository) with build_flask_client(container) as flask_client: yield flask_client @@ -32,7 +59,7 @@ def flask_client(build_flask_client): def test_agent_registration(flask_client): resp = flask_client.post( - AGENTS_URL, + get_url_for_resource(Agents), json=AGENT_REGISTRATION_DICT, follow_redirects=True, ) @@ -46,9 +73,29 @@ def test_agent_registration_invalid_data(flask_client): agent_registration_dict["id"] = 1 resp = flask_client.post( - AGENTS_URL, + get_url_for_resource(Agents), json=agent_registration_dict, follow_redirects=True, ) assert resp.status_code == HTTPStatus.BAD_REQUEST + + +def test_get_agents__status_code(flask_client): + resp = flask_client.get( + get_url_for_resource(Agents), + follow_redirects=True, + ) + assert resp.status_code == HTTPStatus.OK + + +def test_get_agents__data(flask_client): + resp = flask_client.get( + get_url_for_resource(Agents), + follow_redirects=True, + ) + + agents = [Agent(**a) for a in resp.json] + assert len(agents) == len(AGENTS) + for a in agents: + assert a in AGENTS