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

Agent: Use IIslandAPIClient.get_agent_signals in ControlChannel

This commit is contained in:
Ilija Lazoroski 2022-09-23 13:52:26 +02:00 committed by Mike Salvatore
parent 88c011e883
commit 3da90223fc
2 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
monkey/infection_monkey/master/control_channel.py
View File

@ -36,7 +36,7 @@ class ControlChannel(IControlChannel):
if not self._control_channel_server: if not self._control_channel_server:
logger.error("Agent should stop because it can't connect to the C&C server.") logger.error("Agent should stop because it can't connect to the C&C server.")
return True return True
return self._island_api_client.should_agent_stop(self._agent_id) return self._island_api_client.get_agent_signals(self._agent_id) is not None
@handle_island_api_errors @handle_island_api_errors
def get_config(self) -> AgentConfiguration: def get_config(self) -> AgentConfiguration:

4
monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py
View File

@ -35,14 +35,14 @@ def control_channel(island_api_client) -> ControlChannel:
def test_control_channel__should_agent_stop(control_channel, island_api_client): def test_control_channel__should_agent_stop(control_channel, island_api_client):
control_channel.should_agent_stop() control_channel.should_agent_stop()
assert island_api_client.should_agent_stop.called_once() assert island_api_client.get_agent_signals.called_once()
@pytest.mark.parametrize("api_error", CONTROL_CHANNEL_API_ERRORS) @pytest.mark.parametrize("api_error", CONTROL_CHANNEL_API_ERRORS)
def test_control_channel__should_agent_stop_raises_error( def test_control_channel__should_agent_stop_raises_error(
control_channel, island_api_client, api_error control_channel, island_api_client, api_error
): ):
island_api_client.should_agent_stop.side_effect = api_error() island_api_client.get_agent_signals.side_effect = api_error()
with pytest.raises(IslandCommunicationError): with pytest.raises(IslandCommunicationError):
control_channel.should_agent_stop() control_channel.should_agent_stop()