Island: Update register_agent errors

This commit is contained in:
Kekoa Kaaikala 2022-09-20 15:50:39 +00:00
parent 2dc4871a7d
commit 8ebcd2ea33
3 changed files with 48 additions and 2 deletions

View File

@ -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

View File

@ -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
"""

View File

@ -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)