Agent: Implement IIslandAPIClient.get_agent_signals in

HTTPIslandAPIClient
This commit is contained in:
Ilija Lazoroski 2022-09-23 13:51:50 +02:00 committed by Mike Salvatore
parent a3d94d7a49
commit 88c011e883
2 changed files with 74 additions and 1 deletions

View File

@ -1,8 +1,9 @@
import functools import functools
import json import json
import logging import logging
from datetime import datetime
from pprint import pformat from pprint import pformat
from typing import List, Sequence from typing import List, Optional, Sequence
import requests import requests
@ -199,6 +200,18 @@ class HTTPIslandAPIClient(IIslandAPIClient):
return serialized_events 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): class HTTPIslandAPIClientFactory(AbstractIslandAPIClientFactory):
def __init__( def __init__(

View File

@ -33,6 +33,8 @@ AGENT_REGISTRATION = AgentRegistrationData(
network_interfaces=[], network_interfaces=[],
) )
TIMESTAMP = 123456789
ISLAND_URI = f"https://{SERVER}/api?action=is-up" ISLAND_URI = f"https://{SERVER}/api?action=is-up"
ISLAND_SEND_LOG_URI = f"https://{SERVER}/api/log" ISLAND_SEND_LOG_URI = f"https://{SERVER}/api/log"
ISLAND_GET_PBA_FILE_URI = f"https://{SERVER}/api/pba/download/{PBA_FILE}" 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_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_CONFIG_URI = f"https://{SERVER}/api/agent-configuration"
ISLAND_GET_PROPAGATION_CREDENTIALS_URI = f"https://{SERVER}/api/propagation-credentials" 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): class Event1(AbstractAgentEvent):
@ -461,3 +464,60 @@ def test_island_api_client_get_credentials_for_propagation__bad_json(island_api_
with pytest.raises(IslandAPIRequestFailedError): with pytest.raises(IslandAPIRequestFailedError):
m.get(ISLAND_GET_PROPAGATION_CREDENTIALS_URI, content=b"bad") m.get(ISLAND_GET_PROPAGATION_CREDENTIALS_URI, content=b"bad")
island_api_client.get_credentials_for_propagation() 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)