forked from p15670423/monkey
Island: Update register_agent errors
This commit is contained in:
parent
2dc4871a7d
commit
8ebcd2ea33
|
@ -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
|
||||
|
|
|
@ -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
|
||||
"""
|
||||
|
||||
|
|
|
@ -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)
|
||||
|
|
Loading…
Reference in New Issue