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)