From 1745b76122021e584c32d38ca316a889532cd6b4 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 1 Sep 2022 10:38:01 -0400 Subject: [PATCH] Island: Implement MongoAgentRepository.get_running_agents() --- .../cc/repository/mongo_agent_repository.py | 6 ++++- .../repository/test_mongo_agent_repository.py | 22 ++++++++++++++++++- 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/repository/mongo_agent_repository.py b/monkey/monkey_island/cc/repository/mongo_agent_repository.py index 02a4cfe57..5166e7168 100644 --- a/monkey/monkey_island/cc/repository/mongo_agent_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_agent_repository.py @@ -27,7 +27,11 @@ class MongoAgentRepository(IAgentRepository): return MongoAgentRepository._mongo_record_to_agent(agent_dict) def get_running_agents(self) -> Sequence[Agent]: - pass + try: + cursor = self._agents_collection.find({"stop_time": None}) + return list(map(MongoAgentRepository._mongo_record_to_agent, cursor)) + except Exception as err: + raise RetrievalError(f"Error retrieving running agents: {err}") @staticmethod def _mongo_record_to_agent(mongo_record: MutableMapping[str, Any]) -> Agent: 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 c8f267c6b..b4e556c77 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 @@ -14,7 +14,7 @@ from monkey_island.cc.repository import ( ) VICTIM_ZERO_ID = uuid4() -AGENTS = ( +RUNNING_AGENTS = ( Agent(id=VICTIM_ZERO_ID, machine_id=1, start_time=datetime.fromtimestamp(1661856718)), Agent( id=uuid4(), @@ -22,6 +22,8 @@ AGENTS = ( start_time=datetime.fromtimestamp(1661856818), parent_id=VICTIM_ZERO_ID, ), +) +STOPPED_AGENTS = ( Agent( id=uuid4(), machine_id=3, @@ -30,6 +32,10 @@ AGENTS = ( stop_time=datetime.fromtimestamp(1661856773), ), ) +AGENTS = ( + *RUNNING_AGENTS, + *STOPPED_AGENTS, +) @pytest.fixture @@ -47,6 +53,7 @@ def error_raising_mock_mongo_client() -> mongomock.MongoClient: # The first call to find() must succeed mongo_client.monkey_island.agents.find_one = MagicMock(side_effect=Exception("some exception")) + mongo_client.monkey_island.agents.find = MagicMock(side_effect=Exception("some exception")) return mongo_client @@ -69,3 +76,16 @@ def test_get_agent_by_id__not_found(agent_repository): def test_get_agent_by_id__retrieval_error(error_raising_agent_repository): with pytest.raises(RetrievalError): error_raising_agent_repository.get_agent_by_id(AGENTS[0].id) + + +def test_get_running_agents(agent_repository): + running_agents = agent_repository.get_running_agents() + + assert len(running_agents) == len(RUNNING_AGENTS) + for a in running_agents: + assert a in RUNNING_AGENTS + + +def test_get_running_agents__retrieval_error(error_raising_agent_repository): + with pytest.raises(RetrievalError): + error_raising_agent_repository.get_running_agents()