From a3d94d7a493578cacb23bd21cc657af30a83a5aa Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 23 Sep 2022 13:51:21 +0200 Subject: [PATCH 1/7] Agent: Add get_agent_signals to IIslandAPIClient --- .../island_api_client/i_island_api_client.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/monkey/infection_monkey/island_api_client/i_island_api_client.py b/monkey/infection_monkey/island_api_client/i_island_api_client.py index 8ecd98b49..025261e06 100644 --- a/monkey/infection_monkey/island_api_client/i_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/i_island_api_client.py @@ -1,5 +1,6 @@ from abc import ABC, abstractmethod -from typing import Sequence +from datetime import datetime +from typing import Optional, Sequence from common import AgentRegistrationData, OperatingSystem from common.agent_configuration import AgentConfiguration @@ -143,3 +144,15 @@ class IIslandAPIClient(ABC): :raises IslandAPITimeoutError: If the command timed out :return: Credentials """ + + @abstractmethod + def get_agent_signals(self, agent_id: str) -> Optional[datetime]: + """ + Get agent signals from the island + + :raises IslandAPIConnectionError: If the client could not connect to the island + :raises IslandAPIRequestError: If there was a problem with the client request + :raises IslandAPIRequestFailedError: If the server experienced an error + :raises IslandAPITimeoutError: If the command timed out + :return: Terminate datetime + """ From 88c011e883a0336f3beebb6686a2dca961937138 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 23 Sep 2022 13:51:50 +0200 Subject: [PATCH 2/7] Agent: Implement IIslandAPIClient.get_agent_signals in HTTPIslandAPIClient --- .../http_island_api_client.py | 15 ++++- .../test_http_island_api_client.py | 60 +++++++++++++++++++ 2 files changed, 74 insertions(+), 1 deletion(-) diff --git a/monkey/infection_monkey/island_api_client/http_island_api_client.py b/monkey/infection_monkey/island_api_client/http_island_api_client.py index 6cd1d86a1..5c65ebf6f 100644 --- a/monkey/infection_monkey/island_api_client/http_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/http_island_api_client.py @@ -1,8 +1,9 @@ import functools import json import logging +from datetime import datetime from pprint import pformat -from typing import List, Sequence +from typing import List, Optional, Sequence import requests @@ -199,6 +200,18 @@ class HTTPIslandAPIClient(IIslandAPIClient): return serialized_events + @handle_island_errors + @convert_json_error_to_island_api_error + def get_agent_signals(self, agent_id: str) -> Optional[datetime]: + url = f"{self._api_url}/agent-signals/{agent_id}" + response = requests.get( # noqa: DUO123 + url, + verify=False, + timeout=SHORT_REQUEST_TIMEOUT, + ) + response.raise_for_status() + return response.json()["terminate"] + class HTTPIslandAPIClientFactory(AbstractIslandAPIClientFactory): def __init__( diff --git a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py index 03117b006..94ad23cac 100644 --- a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py +++ b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py @@ -33,6 +33,8 @@ AGENT_REGISTRATION = AgentRegistrationData( network_interfaces=[], ) +TIMESTAMP = 123456789 + ISLAND_URI = f"https://{SERVER}/api?action=is-up" ISLAND_SEND_LOG_URI = f"https://{SERVER}/api/log" ISLAND_GET_PBA_FILE_URI = f"https://{SERVER}/api/pba/download/{PBA_FILE}" @@ -42,6 +44,7 @@ ISLAND_REGISTER_AGENT_URI = f"https://{SERVER}/api/agents" ISLAND_AGENT_STOP_URI = f"https://{SERVER}/api/monkey-control/needs-to-stop/{AGENT_ID}" ISLAND_GET_CONFIG_URI = f"https://{SERVER}/api/agent-configuration" ISLAND_GET_PROPAGATION_CREDENTIALS_URI = f"https://{SERVER}/api/propagation-credentials" +ISLAND_GET_AGENT_SIGNALS = f"https://{SERVER}/api/agent-signals/{AGENT_ID}" class Event1(AbstractAgentEvent): @@ -461,3 +464,60 @@ def test_island_api_client_get_credentials_for_propagation__bad_json(island_api_ with pytest.raises(IslandAPIRequestFailedError): m.get(ISLAND_GET_PROPAGATION_CREDENTIALS_URI, content=b"bad") island_api_client.get_credentials_for_propagation() + + +@pytest.mark.parametrize( + "actual_error, expected_error", + [ + (requests.exceptions.ConnectionError, IslandAPIConnectionError), + (TimeoutError, IslandAPITimeoutError), + ], +) +def test_island_api_client__get_agent_signals(island_api_client, actual_error, expected_error): + with requests_mock.Mocker() as m: + m.get(ISLAND_URI) + island_api_client.connect(SERVER) + + with pytest.raises(expected_error): + m.get(ISLAND_GET_AGENT_SIGNALS, exc=actual_error) + island_api_client.get_agent_signals(agent_id=AGENT_ID) + + +@pytest.mark.parametrize( + "status_code, expected_error", + [ + (401, IslandAPIRequestError), + (501, IslandAPIRequestFailedError), + ], +) +def test_island_api_client_get_agent_signals__status_code( + island_api_client, status_code, expected_error +): + with requests_mock.Mocker() as m: + m.get(ISLAND_URI) + island_api_client.connect(SERVER) + + with pytest.raises(expected_error): + m.get(ISLAND_GET_AGENT_SIGNALS, status_code=status_code) + island_api_client.get_agent_signals(agent_id=AGENT_ID) + + +def test_island_api_client_get_agent_signals(island_api_client): + with requests_mock.Mocker() as m: + m.get(ISLAND_URI) + island_api_client.connect(SERVER) + + m.get(ISLAND_GET_AGENT_SIGNALS, json={"terminate": TIMESTAMP}) + actual_terminate_timestamp = island_api_client.get_agent_signals(agent_id=AGENT_ID) + + assert actual_terminate_timestamp == TIMESTAMP + + +def test_island_api_client_get_agent_signals__bad_json(island_api_client): + with requests_mock.Mocker() as m: + m.get(ISLAND_URI) + island_api_client.connect(SERVER) + + with pytest.raises(IslandAPIError): + m.get(ISLAND_GET_AGENT_SIGNALS, json={"bogus": "vogus"}) + island_api_client.get_agent_signals(agent_id=AGENT_ID) From 3da90223fcabee9e288db5d3f83aa42a5a7cd557 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 23 Sep 2022 13:52:26 +0200 Subject: [PATCH 3/7] Agent: Use IIslandAPIClient.get_agent_signals in ControlChannel --- monkey/infection_monkey/master/control_channel.py | 2 +- .../infection_monkey/master/test_control_channel.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py index 48b827f63..947d6c0da 100644 --- a/monkey/infection_monkey/master/control_channel.py +++ b/monkey/infection_monkey/master/control_channel.py @@ -36,7 +36,7 @@ class ControlChannel(IControlChannel): if not self._control_channel_server: logger.error("Agent should stop because it can't connect to the C&C server.") 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 def get_config(self) -> AgentConfiguration: diff --git a/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py b/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py index 658635615..1da0d0713 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py +++ b/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py @@ -35,14 +35,14 @@ def control_channel(island_api_client) -> ControlChannel: def test_control_channel__should_agent_stop(control_channel, island_api_client): 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) def test_control_channel__should_agent_stop_raises_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): control_channel.should_agent_stop() From 67956358bd802ea65ddb901223a92446948f60b5 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 23 Sep 2022 14:11:57 +0200 Subject: [PATCH 4/7] Agent: Remove shoudl_agent_stop from IIslandAPIClient --- .../island_api_client/i_island_api_client.py | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/monkey/infection_monkey/island_api_client/i_island_api_client.py b/monkey/infection_monkey/island_api_client/i_island_api_client.py index 025261e06..34348aeef 100644 --- a/monkey/infection_monkey/island_api_client/i_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/i_island_api_client.py @@ -108,19 +108,6 @@ class IIslandAPIClient(ABC): :raises IslandAPITimeoutError: If the command timed out """ - @abstractmethod - def should_agent_stop(self, agent_id: str) -> bool: - """ - Check with the island to see if the agent should stop - - :param agent_id: The agent identifier for the agent to check - :raises IslandAPIConnectionError: If the client could not connect to the island - :raises IslandAPIRequestError: If there was a problem with the client request - :raises IslandAPIRequestFailedError: If the server experienced an error - :raises IslandAPITimeoutError: If the command timed out - :return: True if the agent should stop, otherwise False - """ - @abstractmethod def get_config(self) -> AgentConfiguration: """ From edf0593d4a2b29f1d3122849dff545fa3b07ff0b Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 23 Sep 2022 14:12:22 +0200 Subject: [PATCH 5/7] Agent: Remove should_agent_stop from HTTPIslandAPIClient --- .../http_island_api_client.py | 13 ------ .../test_http_island_api_client.py | 46 ------------------- 2 files changed, 59 deletions(-) diff --git a/monkey/infection_monkey/island_api_client/http_island_api_client.py b/monkey/infection_monkey/island_api_client/http_island_api_client.py index 5c65ebf6f..65e34df3c 100644 --- a/monkey/infection_monkey/island_api_client/http_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/http_island_api_client.py @@ -147,19 +147,6 @@ class HTTPIslandAPIClient(IIslandAPIClient): ) response.raise_for_status() - @handle_island_errors - @convert_json_error_to_island_api_error - def should_agent_stop(self, agent_id: str) -> bool: - url = f"{self._api_url}/monkey-control/needs-to-stop/{agent_id}" - response = requests.get( # noqa: DUO123 - url, - verify=False, - timeout=SHORT_REQUEST_TIMEOUT, - ) - response.raise_for_status() - - return response.json()["stop_agent"] - @handle_island_errors @convert_json_error_to_island_api_error def get_config(self) -> AgentConfiguration: diff --git a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py index 94ad23cac..376425696 100644 --- a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py +++ b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py @@ -328,52 +328,6 @@ def test_island_api_client_register_agent__status_code( island_api_client.register_agent(AGENT_REGISTRATION) -@pytest.mark.parametrize( - "actual_error, expected_error", - [ - (requests.exceptions.ConnectionError, IslandAPIConnectionError), - (TimeoutError, IslandAPITimeoutError), - ], -) -def test_island_api_client__should_agent_stop(island_api_client, actual_error, expected_error): - with requests_mock.Mocker() as m: - m.get(ISLAND_URI) - island_api_client.connect(SERVER) - - with pytest.raises(expected_error): - m.get(ISLAND_AGENT_STOP_URI, exc=actual_error) - island_api_client.should_agent_stop(AGENT_ID) - - -@pytest.mark.parametrize( - "status_code, expected_error", - [ - (401, IslandAPIRequestError), - (501, IslandAPIRequestFailedError), - ], -) -def test_island_api_client_should_agent_stop__status_code( - island_api_client, status_code, expected_error -): - with requests_mock.Mocker() as m: - m.get(ISLAND_URI) - island_api_client.connect(SERVER) - - with pytest.raises(expected_error): - m.get(ISLAND_AGENT_STOP_URI, status_code=status_code) - island_api_client.should_agent_stop(AGENT_ID) - - -def test_island_api_client_should_agent_stop__bad_json(island_api_client): - with requests_mock.Mocker() as m: - m.get(ISLAND_URI) - island_api_client.connect(SERVER) - - with pytest.raises(IslandAPIRequestFailedError): - m.get(ISLAND_AGENT_STOP_URI, content=b"bad") - island_api_client.should_agent_stop(AGENT_ID) - - @pytest.mark.parametrize( "actual_error, expected_error", [ From d1fc4fa7f42f7e86e5f8fd6b5077b6b834b56477 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 23 Sep 2022 15:49:41 +0200 Subject: [PATCH 6/7] UT: Parametrize HTTPIslandAPIClient get_agent_signals test --- .../island_api_client/test_http_island_api_client.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py index 376425696..9505e6649 100644 --- a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py +++ b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py @@ -456,15 +456,16 @@ def test_island_api_client_get_agent_signals__status_code( island_api_client.get_agent_signals(agent_id=AGENT_ID) -def test_island_api_client_get_agent_signals(island_api_client): +@pytest.mark.parametrize("expected_timestamp", [TIMESTAMP, None]) +def test_island_api_client_get_agent_signals(island_api_client, expected_timestamp): with requests_mock.Mocker() as m: m.get(ISLAND_URI) island_api_client.connect(SERVER) - m.get(ISLAND_GET_AGENT_SIGNALS, json={"terminate": TIMESTAMP}) + m.get(ISLAND_GET_AGENT_SIGNALS, json={"terminate": expected_timestamp}) actual_terminate_timestamp = island_api_client.get_agent_signals(agent_id=AGENT_ID) - assert actual_terminate_timestamp == TIMESTAMP + assert actual_terminate_timestamp == expected_timestamp def test_island_api_client_get_agent_signals__bad_json(island_api_client): From a314efb8d963d5c617fb844ead64d138ba1f33e4 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 23 Sep 2022 15:55:12 +0200 Subject: [PATCH 7/7] Agent: Reword get_agent_signals docstring --- .../island_api_client/i_island_api_client.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/monkey/infection_monkey/island_api_client/i_island_api_client.py b/monkey/infection_monkey/island_api_client/i_island_api_client.py index 34348aeef..c2a4dc899 100644 --- a/monkey/infection_monkey/island_api_client/i_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/i_island_api_client.py @@ -135,11 +135,12 @@ class IIslandAPIClient(ABC): @abstractmethod def get_agent_signals(self, agent_id: str) -> Optional[datetime]: """ - Get agent signals from the island + Gets an agent's signals from the island + :param agent_id: ID of the agent whose signals should be retrieved :raises IslandAPIConnectionError: If the client could not connect to the island :raises IslandAPIRequestError: If there was a problem with the client request :raises IslandAPIRequestFailedError: If the server experienced an error :raises IslandAPITimeoutError: If the command timed out - :return: Terminate datetime + :return: The relevant agent's terminate signal's timestamp """