p15670423
/
monkey
1
0
Fork 5
Code Issues Pull Requests 3 Projects Releases Wiki Activity

Island: Implement `get_progenitor` in MongoAgentRepository

This commit is contained in:
Ilija Lazoroski 2022-09-22 14:45:07 +02:00 committed by Mike Salvatore
parent d10c148533
commit c25e245a8e
2 changed files with 43 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -58,6 +58,14 @@ class MongoAgentRepository(IAgentRepository):
except Exception as err:
raise RetrievalError(f"Error retrieving running agents: {err}")
def get_progenitor(self, agent: Agent) -> Agent:
if agent.parent_id is None:
return agent
parent = self.get_agent_by_id(agent.parent_id)
return self.get_progenitor(parent)
def reset(self):
try:
self._agents_collection.drop()

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

@ -17,14 +17,29 @@ from monkey_island.cc.repository import (
)
VICTIM_ZERO_ID = uuid4()
VICTIM_TWO_ID = uuid4()
VICTIM_THREE_ID = uuid4()
PROGENITOR_AGENT = Agent(
id=VICTIM_ZERO_ID, machine_id=1, start_time=datetime.fromtimestamp(1661856718)
)
DESCENDANT_AGENT = Agent(
id=VICTIM_THREE_ID,
machine_id=4,
start_time=datetime.fromtimestamp(1661856868),
parent_id=VICTIM_TWO_ID,
)
RUNNING_AGENTS = (
Agent(id=VICTIM_ZERO_ID, machine_id=1, start_time=datetime.fromtimestamp(1661856718)),
PROGENITOR_AGENT,
Agent(
id=uuid4(),
id=VICTIM_TWO_ID,
machine_id=2,
start_time=datetime.fromtimestamp(1661856818),
parent_id=VICTIM_ZERO_ID,
),
DESCENDANT_AGENT,
)
STOPPED_AGENTS = (
Agent(
@ -172,6 +187,24 @@ def test_get_running_agents__retrieval_error(error_raising_agent_repository):
error_raising_agent_repository.get_running_agents()
@pytest.mark.parametrize("agent", [DESCENDANT_AGENT, PROGENITOR_AGENT])
def test_get_progenitor(agent_repository, agent):
actual_progenitor = agent_repository.get_progenitor(agent)
assert actual_progenitor == PROGENITOR_AGENT
def test_get_progenitor__id_not_found(agent_repository):
dummy_agent = Agent(id=uuid4(), machine_id=10, start_time=datetime.now(), parent_id=uuid4())
with pytest.raises(UnknownRecordError):
agent_repository.get_progenitor(dummy_agent)
def test_get_progenitor__retrieval_error(error_raising_agent_repository):
with pytest.raises(RetrievalError):
error_raising_agent_repository.get_progenitor(AGENTS[1])
def test_reset(agent_repository):
# Ensure the repository is not empty
for agent in AGENTS: