Merge pull request #2380 from guardicore/2362-get-agents

2362 get agents

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/<uuid:agent_id>` endpoint. #2274
 

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:
         """

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}")
 

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:

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

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