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 17e2f9284..f277ad9c0 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 @@ -135,6 +135,7 @@ class HTTPIslandAPIClient(IIslandAPIClient): response.raise_for_status() + @handle_island_errors def register_agent(self, agent_registration_data: AgentRegistrationData): url = f"https://{agent_registration_data.cc_server}/api/agents" response = requests.post( # noqa: DUO123 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 409a1b265..2f5edc8f4 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 @@ -104,6 +104,8 @@ class IIslandAPIClient(ABC): :param agent_registration_data: Information about the agent to register with 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 """ 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 1d0ac17fc..d4acc1c79 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 @@ -10,6 +10,7 @@ from common.agent_event_serializers import ( PydanticAgentEventSerializer, ) from common.agent_events import AbstractAgentEvent +from common.agent_registration_data import AgentRegistrationData from infection_monkey.island_api_client import ( HTTPIslandAPIClient, IslandAPIConnectionError, @@ -22,14 +23,22 @@ from infection_monkey.island_api_client import ( SERVER = "1.1.1.1:9999" PBA_FILE = "dummy.pba" WINDOWS = "windows" +AGENT_ID = UUID("80988359-a1cd-42a2-9b47-5b94b37cd673") +AGENT_REGISTRATION = AgentRegistrationData( + id=AGENT_ID, + machine_hardware_id=1, + start_time=0, + parent_id=None, + cc_server=SERVER, + network_interfaces=[], +) 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}" ISLAND_GET_AGENT_BINARY_URI = f"https://{SERVER}/api/agent-binaries/{WINDOWS}" ISLAND_SEND_EVENTS_URI = f"https://{SERVER}/api/agent-events" - -AGENT_ID = UUID("80988359-a1cd-42a2-9b47-5b94b37cd673") +ISLAND_REGISTER_AGENT_URI = f"https://{SERVER}/api/agents" class Event1(AbstractAgentEvent): @@ -275,3 +284,37 @@ def test_island_api_client_send_events__status_code(island_api_client, status_co with pytest.raises(expected_error): m.post(ISLAND_SEND_EVENTS_URI, status_code=status_code) island_api_client.send_events(events=[Event1(source=AGENT_ID, a=1)]) + + +@pytest.mark.parametrize( + "actual_error, expected_error", + [ + (requests.exceptions.ConnectionError, IslandAPIConnectionError), + (TimeoutError, IslandAPITimeoutError), + ], +) +def test_island_api_client__register_agent(actual_error, expected_error): + with requests_mock.Mocker() as m: + m.get(ISLAND_URI) + island_api_client = HTTPIslandAPIClient(SERVER) + + with pytest.raises(expected_error): + m.post(ISLAND_REGISTER_AGENT_URI, exc=actual_error) + island_api_client.register_agent(AGENT_REGISTRATION) + + +@pytest.mark.parametrize( + "status_code, expected_error", + [ + (401, IslandAPIRequestError), + (501, IslandAPIRequestFailedError), + ], +) +def test_island_api_client_register_agent__status_code(status_code, expected_error): + with requests_mock.Mocker() as m: + m.get(ISLAND_URI) + island_api_client = HTTPIslandAPIClient(SERVER) + + with pytest.raises(expected_error): + m.post(ISLAND_REGISTER_AGENT_URI, status_code=status_code) + island_api_client.register_agent(AGENT_REGISTRATION)