forked from p15670423/monkey
Agent: Implement IIslandAPIClient.get_agent_signals in
HTTPIslandAPIClient
This commit is contained in:
parent
a3d94d7a49
commit
88c011e883
|
@ -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__(
|
||||
|
|
|
@ -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)
|
||||
|
|
Loading…
Reference in New Issue