p34709852
/
monkey
forked from p15670423/monkey
1
0
Fork 1
Code Issues Pull Requests Projects Releases Wiki Activity

Island: Modify HTTPIslandAPIClient.get_agent_signals to return 

AgentSignals
This commit is contained in:
Ilija Lazoroski 2022-09-23 17:11:44 +02:00 committed by Mike Salvatore
parent a2be330d16
commit 6299529f4a
2 changed files with 11 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
monkey/infection_monkey/island_api_client/http_island_api_client.py
View File

@ -1,13 +1,12 @@
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, Optional, Sequence from typing import List, Sequence
import requests import requests
from common import AgentRegistrationData, OperatingSystem from common import AgentRegistrationData, AgentSignals, OperatingSystem
from common.agent_configuration import AgentConfiguration from common.agent_configuration import AgentConfiguration
from common.agent_event_serializers import AgentEventSerializerRegistry, JSONSerializable from common.agent_event_serializers import AgentEventSerializerRegistry, JSONSerializable
from common.agent_events import AbstractAgentEvent from common.agent_events import AbstractAgentEvent
@ -189,7 +188,7 @@ class HTTPIslandAPIClient(IIslandAPIClient):
@handle_island_errors @handle_island_errors
@convert_json_error_to_island_api_error @convert_json_error_to_island_api_error
def get_agent_signals(self, agent_id: str) -> Optional[datetime]: def get_agent_signals(self, agent_id: str) -> AgentSignals:
url = f"{self._api_url}/agent-signals/{agent_id}" url = f"{self._api_url}/agent-signals/{agent_id}"
response = requests.get( # noqa: DUO123 response = requests.get( # noqa: DUO123
url, url,
@ -197,7 +196,7 @@ class HTTPIslandAPIClient(IIslandAPIClient):
timeout=SHORT_REQUEST_TIMEOUT, timeout=SHORT_REQUEST_TIMEOUT,
) )
response.raise_for_status() response.raise_for_status()
return response.json()["terminate"] return AgentSignals(**response.json())
class HTTPIslandAPIClientFactory(AbstractIslandAPIClientFactory): class HTTPIslandAPIClientFactory(AbstractIslandAPIClientFactory):

13
monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py
View File

@ -4,7 +4,7 @@ import pytest
import requests import requests
import requests_mock import requests_mock
from common import OperatingSystem from common import AgentSignals, OperatingSystem
from common.agent_event_serializers import ( from common.agent_event_serializers import (
AgentEventSerializerRegistry, AgentEventSerializerRegistry,
PydanticAgentEventSerializer, PydanticAgentEventSerializer,
@ -456,16 +456,17 @@ def test_island_api_client_get_agent_signals__status_code(
island_api_client.get_agent_signals(agent_id=AGENT_ID) island_api_client.get_agent_signals(agent_id=AGENT_ID)
@pytest.mark.parametrize("expected_timestamp", [TIMESTAMP, None]) @pytest.mark.parametrize("timestamp", [TIMESTAMP, None])
def test_island_api_client_get_agent_signals(island_api_client, expected_timestamp): def test_island_api_client_get_agent_signals(island_api_client, timestamp):
expected_agent_signals = AgentSignals(terminate=timestamp)
with requests_mock.Mocker() as m: with requests_mock.Mocker() as m:
m.get(ISLAND_URI) m.get(ISLAND_URI)
island_api_client.connect(SERVER) island_api_client.connect(SERVER)
m.get(ISLAND_GET_AGENT_SIGNALS, json={"terminate": expected_timestamp}) m.get(ISLAND_GET_AGENT_SIGNALS, json={"terminate": timestamp})
actual_terminate_timestamp = island_api_client.get_agent_signals(agent_id=AGENT_ID) actual_agent_signals = island_api_client.get_agent_signals(agent_id=AGENT_ID)
assert actual_terminate_timestamp == expected_timestamp assert actual_agent_signals == expected_agent_signals
def test_island_api_client_get_agent_signals__bad_json(island_api_client): def test_island_api_client_get_agent_signals__bad_json(island_api_client):