p41275903
/
monkey
forked from p15670423/monkey
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Island: Implement MongoAgentRepository.get_running_agents()

This commit is contained in:
Mike Salvatore 2022-09-01 10:38:01 -04:00
parent ba228e56b6
commit 1745b76122
2 changed files with 26 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
monkey/monkey_island/cc/repository/mongo_agent_repository.py
View File

@ -27,7 +27,11 @@ class MongoAgentRepository(IAgentRepository):
return MongoAgentRepository._mongo_record_to_agent(agent_dict) return MongoAgentRepository._mongo_record_to_agent(agent_dict)
def get_running_agents(self) -> Sequence[Agent]: 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 @staticmethod
def _mongo_record_to_agent(mongo_record: MutableMapping[str, Any]) -> Agent: def _mongo_record_to_agent(mongo_record: MutableMapping[str, Any]) -> Agent:

22
monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_repository.py
View File

@ -14,7 +14,7 @@ from monkey_island.cc.repository import (
) )
VICTIM_ZERO_ID = uuid4() VICTIM_ZERO_ID = uuid4()
AGENTS = ( RUNNING_AGENTS = (
Agent(id=VICTIM_ZERO_ID, machine_id=1, start_time=datetime.fromtimestamp(1661856718)), Agent(id=VICTIM_ZERO_ID, machine_id=1, start_time=datetime.fromtimestamp(1661856718)),
Agent( Agent(
id=uuid4(), id=uuid4(),
@ -22,6 +22,8 @@ AGENTS = (
start_time=datetime.fromtimestamp(1661856818), start_time=datetime.fromtimestamp(1661856818),
parent_id=VICTIM_ZERO_ID, parent_id=VICTIM_ZERO_ID,
), ),
)
STOPPED_AGENTS = (
Agent( Agent(
id=uuid4(), id=uuid4(),
machine_id=3, machine_id=3,
@ -30,6 +32,10 @@ AGENTS = (
stop_time=datetime.fromtimestamp(1661856773), stop_time=datetime.fromtimestamp(1661856773),
), ),
) )
AGENTS = (
*RUNNING_AGENTS,
*STOPPED_AGENTS,
)
@pytest.fixture @pytest.fixture
@ -47,6 +53,7 @@ def error_raising_mock_mongo_client() -> mongomock.MongoClient:
# The first call to find() must succeed # 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_one = MagicMock(side_effect=Exception("some exception"))
mongo_client.monkey_island.agents.find = MagicMock(side_effect=Exception("some exception"))
return mongo_client 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): def test_get_agent_by_id__retrieval_error(error_raising_agent_repository):
with pytest.raises(RetrievalError): with pytest.raises(RetrievalError):
error_raising_agent_repository.get_agent_by_id(AGENTS[0].id) 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()