From 54ef77698c391cb21423386d05fddb56671cf863 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Mon, 19 Sep 2022 16:08:42 +0000 Subject: [PATCH 01/20] Agent: Add register_agent to IslandAPIClient --- .../http_island_api_client.py | 27 ++++++++++- .../island_api_client/i_island_api_client.py | 14 +++++- .../master/control_channel.py | 24 ++++------ .../master/test_control_channel.py | 45 +++++++++++++++++++ 4 files changed, 92 insertions(+), 18 deletions(-) create mode 100644 monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py diff --git a/monkey/infection_monkey/island_api_client/http_island_api_client.py b/monkey/infection_monkey/island_api_client/http_island_api_client.py index a0bebc5a0..407b6562e 100644 --- a/monkey/infection_monkey/island_api_client/http_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/http_island_api_client.py @@ -4,10 +4,14 @@ from typing import List, Sequence import requests -from common import OperatingSystem +from common import AgentRegistrationData, OperatingSystem from common.agent_event_serializers import AgentEventSerializerRegistry, JSONSerializable from common.agent_events import AbstractAgentEvent -from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT, MEDIUM_REQUEST_TIMEOUT +from common.common_consts.timeouts import ( + LONG_REQUEST_TIMEOUT, + MEDIUM_REQUEST_TIMEOUT, + SHORT_REQUEST_TIMEOUT, +) from . import ( AbstractIslandAPIClientFactory, @@ -116,6 +120,25 @@ class HTTPIslandAPIClient(IIslandAPIClient): response.raise_for_status() + def register_agent(self, agent_registration_data: AgentRegistrationData): + try: + url = f"https://{agent_registration_data.cc_server}/api/agents" + response = requests.post( # noqa: DUO123 + url, + json=agent_registration_data.dict(simplify=True), + verify=False, + timeout=SHORT_REQUEST_TIMEOUT, + ) + response.raise_for_status() + except ( + requests.exceptions.ConnectionError, + requests.exceptions.TooManyRedirects, + requests.exceptions.HTTPError, + ) as e: + raise IslandAPIConnectionError(e) + except requests.exceptions.Timeout as e: + raise IslandAPITimeoutError(e) + def _serialize_events(self, events: Sequence[AbstractAgentEvent]) -> JSONSerializable: serialized_events: List[JSONSerializable] = [] diff --git a/monkey/infection_monkey/island_api_client/i_island_api_client.py b/monkey/infection_monkey/island_api_client/i_island_api_client.py index 5bebc79c1..cc32555dd 100644 --- a/monkey/infection_monkey/island_api_client/i_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/i_island_api_client.py @@ -4,6 +4,8 @@ from typing import Optional, Sequence from common import OperatingSystem from common.agent_events import AbstractAgentEvent +from common import AgentRegistrationData + class IIslandAPIClient(ABC): """ @@ -74,7 +76,6 @@ class IIslandAPIClient(ABC): :raises IslandAPITimeoutError: If a timeout occurs while attempting to connect to the island :raises IslandAPIError: If an unexpected error occurs while attempting to retrieve the agent binary - """ @abstractmethod @@ -92,3 +93,14 @@ class IIslandAPIClient(ABC): :raises IslandAPIError: If an unexpected error occurs while attempting to send events to the island """ + + @abstractmethod + def register_agent(self, agent_registration_data: AgentRegistrationData): + """ + Register an agent with the Island + + :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 IslandAPITimeoutError: If the command timed out + """ diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py index 76be63b5d..cd4496a9d 100644 --- a/monkey/infection_monkey/master/control_channel.py +++ b/monkey/infection_monkey/master/control_channel.py @@ -13,6 +13,11 @@ from common.common_consts.timeouts import SHORT_REQUEST_TIMEOUT from common.credentials import Credentials from common.network.network_utils import get_network_interfaces from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError +from infection_monkey.island_api_client import ( + IIslandAPIClient, + IslandAPIConnectionError, + IslandAPITimeoutError, +) from infection_monkey.utils import agent_process from infection_monkey.utils.ids import get_agent_id, get_machine_id @@ -22,9 +27,10 @@ logger = logging.getLogger(__name__) class ControlChannel(IControlChannel): - def __init__(self, server: str, agent_id: str): + def __init__(self, server: str, agent_id: str, api_client: IIslandAPIClient): self._agent_id = agent_id self._control_channel_server = server + self._island_api_client = api_client def register_agent(self, parent: Optional[UUID] = None): agent_registration_data = AgentRegistrationData( @@ -38,20 +44,8 @@ class ControlChannel(IControlChannel): ) try: - url = f"https://{self._control_channel_server}/api/agents" - response = requests.post( # noqa: DUO123 - url, - json=agent_registration_data.dict(simplify=True), - verify=False, - timeout=SHORT_REQUEST_TIMEOUT, - ) - response.raise_for_status() - except ( - requests.exceptions.ConnectionError, - requests.exceptions.Timeout, - requests.exceptions.TooManyRedirects, - requests.exceptions.HTTPError, - ) as e: + self._island_api_client.register_agent(agent_registration_data) + except (IslandAPIConnectionError, IslandAPITimeoutError) as e: raise IslandCommunicationError(e) def should_agent_stop(self) -> bool: diff --git a/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py b/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py new file mode 100644 index 000000000..75a3eb149 --- /dev/null +++ b/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py @@ -0,0 +1,45 @@ +from unittest.mock import MagicMock + +import pytest + +from infection_monkey.i_control_channel import IslandCommunicationError +from infection_monkey.island_api_client import ( + IIslandAPIClient, + IslandAPIConnectionError, + IslandAPITimeoutError, +) +from infection_monkey.master.control_channel import ControlChannel + + +@pytest.fixture +def island_api_client() -> IIslandAPIClient: + client = MagicMock() + return client + + +@pytest.fixture +def control_channel(island_api_client) -> ControlChannel: + return ControlChannel("server", "agent-id", island_api_client) + + +def test_control_channel__register_agent(control_channel, island_api_client): + control_channel.register_agent() + assert island_api_client.register_agent.called_once() + + +def test_control_channel__register_agent_raises_on_connection_error( + control_channel, island_api_client +): + island_api_client.register_agent.side_effect = IslandAPIConnectionError() + + with pytest.raises(IslandCommunicationError): + control_channel.register_agent() + + +def test_control_channel__register_agent_raises_on_timeout_error( + control_channel, island_api_client +): + island_api_client.register_agent.side_effect = IslandAPITimeoutError() + + with pytest.raises(IslandCommunicationError): + control_channel.register_agent() From 92e793c2cd4e3026cc3a1b2853f6e2dd69aeeaf2 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Mon, 19 Sep 2022 17:09:04 +0000 Subject: [PATCH 02/20] Agent: Add should_agent_stop to IslandAPIClient --- .../http_island_api_client.py | 28 +++++++++++ .../island_api_client/i_island_api_client.py | 14 ++++++ .../master/control_channel.py | 25 ++++------ .../master/test_control_channel.py | 48 ++++++++++++++++++- 4 files changed, 97 insertions(+), 18 deletions(-) diff --git a/monkey/infection_monkey/island_api_client/http_island_api_client.py b/monkey/infection_monkey/island_api_client/http_island_api_client.py index 407b6562e..97c5404d0 100644 --- a/monkey/infection_monkey/island_api_client/http_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/http_island_api_client.py @@ -1,4 +1,5 @@ import functools +import json import logging from typing import List, Sequence @@ -139,6 +140,33 @@ class HTTPIslandAPIClient(IIslandAPIClient): except requests.exceptions.Timeout as e: raise IslandAPITimeoutError(e) + def should_agent_stop(self, island_server: str, agent_id: str) -> bool: + try: + url = f"https://{island_server}/api/monkey-control" f"/needs-to-stop/{agent_id}" + response = requests.get( # noqa: DUO123 + url, + verify=False, + timeout=SHORT_REQUEST_TIMEOUT, + ) + response.raise_for_status() + + json_response = json.loads(response.content.decode()) + return json_response["stop_agent"] + except ( + requests.exceptions.ConnectionError, + requests.exceptions.TooManyRedirects, + ) as e: + raise IslandAPIConnectionError(e) + except requests.exceptions.Timeout as e: + raise IslandAPITimeoutError(e) + except requests.exceptions.HTTPError as e: + if e.errno >= 500: + raise IslandAPIRequestFailedError(e) + else: + raise IslandAPIRequestError(e) + except json.JSONDecodeError as e: + raise IslandAPIRequestFailedError(e) + def _serialize_events(self, events: Sequence[AbstractAgentEvent]) -> JSONSerializable: serialized_events: List[JSONSerializable] = [] diff --git a/monkey/infection_monkey/island_api_client/i_island_api_client.py b/monkey/infection_monkey/island_api_client/i_island_api_client.py index cc32555dd..cbfbbf306 100644 --- a/monkey/infection_monkey/island_api_client/i_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/i_island_api_client.py @@ -104,3 +104,17 @@ class IIslandAPIClient(ABC): :raises IslandAPIConnectionError: If the client could not connect to the island :raises IslandAPITimeoutError: If the command timed out """ + + @abstractmethod + def should_agent_stop(self, island_server: str, agent_id: str) -> bool: + """ + Check with the island to see if the agent should stop + + :param island_server: The server to query + :param agent_id: The agent identifier for the agent to check + :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 + :return: True if the agent should stop, otherwise False + """ diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py index cd4496a9d..f61091e8f 100644 --- a/monkey/infection_monkey/master/control_channel.py +++ b/monkey/infection_monkey/master/control_channel.py @@ -16,6 +16,8 @@ from infection_monkey.i_control_channel import IControlChannel, IslandCommunicat from infection_monkey.island_api_client import ( IIslandAPIClient, IslandAPIConnectionError, + IslandAPIRequestError, + IslandAPIRequestFailedError, IslandAPITimeoutError, ) from infection_monkey.utils import agent_process @@ -53,25 +55,14 @@ class ControlChannel(IControlChannel): logger.error("Agent should stop because it can't connect to the C&C server.") return True try: - url = ( - f"https://{self._control_channel_server}/api/monkey-control" - f"/needs-to-stop/{self._agent_id}" + return self._island_api_client.should_agent_stop( + self._control_channel_server, self._agent_id ) - response = requests.get( # noqa: DUO123 - url, - verify=False, - timeout=SHORT_REQUEST_TIMEOUT, - ) - response.raise_for_status() - - json_response = json.loads(response.content.decode()) - return json_response["stop_agent"] except ( - json.JSONDecodeError, - requests.exceptions.ConnectionError, - requests.exceptions.Timeout, - requests.exceptions.TooManyRedirects, - requests.exceptions.HTTPError, + IslandAPIConnectionError, + IslandAPIRequestError, + IslandAPIRequestFailedError, + IslandAPITimeoutError, ) as e: raise IslandCommunicationError(e) diff --git a/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py b/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py index 75a3eb149..534ef157b 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py +++ b/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py @@ -6,10 +6,15 @@ from infection_monkey.i_control_channel import IslandCommunicationError from infection_monkey.island_api_client import ( IIslandAPIClient, IslandAPIConnectionError, + IslandAPIRequestError, + IslandAPIRequestFailedError, IslandAPITimeoutError, ) from infection_monkey.master.control_channel import ControlChannel +SERVER = "server" +AGENT_ID = "agent" + @pytest.fixture def island_api_client() -> IIslandAPIClient: @@ -19,7 +24,7 @@ def island_api_client() -> IIslandAPIClient: @pytest.fixture def control_channel(island_api_client) -> ControlChannel: - return ControlChannel("server", "agent-id", island_api_client) + return ControlChannel(SERVER, AGENT_ID, island_api_client) def test_control_channel__register_agent(control_channel, island_api_client): @@ -43,3 +48,44 @@ def test_control_channel__register_agent_raises_on_timeout_error( with pytest.raises(IslandCommunicationError): control_channel.register_agent() + + +def test_control_channel__should_agent_stop(control_channel, island_api_client): + control_channel.should_agent_stop() + assert island_api_client.should_agent_stop.called_once() + + +def test_control_channel__should_agent_stop_raises_on_connection_error( + control_channel, island_api_client +): + island_api_client.should_agent_stop.side_effect = IslandAPIConnectionError() + + with pytest.raises(IslandCommunicationError): + control_channel.should_agent_stop() + + +def test_control_channel__should_agent_stop_raises_on_timeout_error( + control_channel, island_api_client +): + island_api_client.should_agent_stop.side_effect = IslandAPITimeoutError() + + with pytest.raises(IslandCommunicationError): + control_channel.should_agent_stop() + + +def test_control_channel__should_agent_stop_raises_on_request_error( + control_channel, island_api_client +): + island_api_client.should_agent_stop.side_effect = IslandAPIRequestError() + + with pytest.raises(IslandCommunicationError): + control_channel.should_agent_stop() + + +def test_control_channel__should_agent_stop_raises_on_request_failed_error( + control_channel, island_api_client +): + island_api_client.should_agent_stop.side_effect = IslandAPIRequestFailedError() + + with pytest.raises(IslandCommunicationError): + control_channel.should_agent_stop() From d6795492a4ba9059cd43cdd936c2720b1a2e5e17 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Mon, 19 Sep 2022 19:53:22 +0000 Subject: [PATCH 03/20] Agent: Add get_config to IslandAPIClient --- .../http_island_api_client.py | 31 +++++++++++++++++++ .../island_api_client/i_island_api_client.py | 14 +++++++++ .../master/control_channel.py | 24 +++----------- .../master/test_control_channel.py | 21 +++++++++++++ 4 files changed, 71 insertions(+), 19 deletions(-) diff --git a/monkey/infection_monkey/island_api_client/http_island_api_client.py b/monkey/infection_monkey/island_api_client/http_island_api_client.py index 97c5404d0..f92fcc946 100644 --- a/monkey/infection_monkey/island_api_client/http_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/http_island_api_client.py @@ -1,11 +1,13 @@ import functools import json import logging +from pprint import pformat from typing import List, Sequence import requests from common import AgentRegistrationData, OperatingSystem +from common.agent_configuration import AgentConfiguration from common.agent_event_serializers import AgentEventSerializerRegistry, JSONSerializable from common.agent_events import AbstractAgentEvent from common.common_consts.timeouts import ( @@ -167,6 +169,35 @@ class HTTPIslandAPIClient(IIslandAPIClient): except json.JSONDecodeError as e: raise IslandAPIRequestFailedError(e) + def get_config(self, island_server: str) -> AgentConfiguration: + try: + response = requests.get( # noqa: DUO123 + f"https://{island_server}/api/agent-configuration", + verify=False, + timeout=SHORT_REQUEST_TIMEOUT, + ) + response.raise_for_status() + + config_dict = json.loads(response.text) + + logger.debug(f"Received configuration:\n{pformat(config_dict)}") + + return AgentConfiguration(**config_dict) + except ( + requests.exceptions.ConnectionError, + requests.exceptions.TooManyRedirects, + ) as e: + raise IslandAPIConnectionError(e) + except requests.exceptions.Timeout as e: + raise IslandAPITimeoutError(e) + except requests.exceptions.HTTPError as e: + if e.errno >= 500: + raise IslandAPIRequestFailedError(e) + else: + raise IslandAPIRequestError(e) + except json.JSONDecodeError as e: + raise IslandAPIRequestFailedError(e) + def _serialize_events(self, events: Sequence[AbstractAgentEvent]) -> JSONSerializable: serialized_events: List[JSONSerializable] = [] diff --git a/monkey/infection_monkey/island_api_client/i_island_api_client.py b/monkey/infection_monkey/island_api_client/i_island_api_client.py index cbfbbf306..064511a73 100644 --- a/monkey/infection_monkey/island_api_client/i_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/i_island_api_client.py @@ -5,6 +5,7 @@ from common import OperatingSystem from common.agent_events import AbstractAgentEvent from common import AgentRegistrationData +from common.agent_configuration import AgentConfiguration class IIslandAPIClient(ABC): @@ -118,3 +119,16 @@ class IIslandAPIClient(ABC): :raises IslandAPITimeoutError: If the command timed out :return: True if the agent should stop, otherwise False """ + + @abstractmethod + def get_config(self, island_server: str) -> AgentConfiguration: + """ + Get agent configuration from the island + + :param island_server: The server to query + :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 + :return: Agent configuration + """ diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py index f61091e8f..d01a6a233 100644 --- a/monkey/infection_monkey/master/control_channel.py +++ b/monkey/infection_monkey/master/control_channel.py @@ -1,6 +1,4 @@ -import json import logging -from pprint import pformat from typing import Optional, Sequence from uuid import UUID @@ -68,24 +66,12 @@ class ControlChannel(IControlChannel): def get_config(self) -> AgentConfiguration: try: - response = requests.get( # noqa: DUO123 - f"https://{self._control_channel_server}/api/agent-configuration", - verify=False, - timeout=SHORT_REQUEST_TIMEOUT, - ) - response.raise_for_status() - - config_dict = json.loads(response.text) - - logger.debug(f"Received configuration:\n{pformat(config_dict)}") - - return AgentConfiguration(**config_dict) + return self._island_api_client.get_config(self._control_channel_server) except ( - json.JSONDecodeError, - requests.exceptions.ConnectionError, - requests.exceptions.Timeout, - requests.exceptions.TooManyRedirects, - requests.exceptions.HTTPError, + IslandAPIConnectionError, + IslandAPIRequestError, + IslandAPIRequestFailedError, + IslandAPITimeoutError, ) as e: raise IslandCommunicationError(e) diff --git a/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py b/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py index 534ef157b..dbb5ffb75 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py +++ b/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py @@ -89,3 +89,24 @@ def test_control_channel__should_agent_stop_raises_on_request_failed_error( with pytest.raises(IslandCommunicationError): control_channel.should_agent_stop() + + +def test_control_channel__get_config(control_channel, island_api_client): + control_channel.get_config() + assert island_api_client.get_config.called_once() + + +@pytest.mark.parametrize( + "api_error", + [ + IslandAPIConnectionError, + IslandAPIRequestError, + IslandAPIRequestFailedError, + IslandAPITimeoutError, + ], +) +def test_control_channel__get_config_raises_error(control_channel, island_api_client, api_error): + island_api_client.get_config.side_effect = api_error() + + with pytest.raises(IslandCommunicationError): + control_channel.get_config() From b260dcc5cb05c9d362feb381e9d8e2ce256599a1 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Mon, 19 Sep 2022 20:06:43 +0000 Subject: [PATCH 04/20] Agent: Add get_credentials_for_propagation to IslandAPIClient --- .../http_island_api_client.py | 26 ++++++++++++++++ .../island_api_client/i_island_api_client.py | 14 +++++++++ .../master/control_channel.py | 23 ++++---------- .../master/test_control_channel.py | 31 +++++++++++++------ 4 files changed, 68 insertions(+), 26 deletions(-) diff --git a/monkey/infection_monkey/island_api_client/http_island_api_client.py b/monkey/infection_monkey/island_api_client/http_island_api_client.py index f92fcc946..e5e338c4e 100644 --- a/monkey/infection_monkey/island_api_client/http_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/http_island_api_client.py @@ -15,6 +15,7 @@ from common.common_consts.timeouts import ( MEDIUM_REQUEST_TIMEOUT, SHORT_REQUEST_TIMEOUT, ) +from common.credentials import Credentials from . import ( AbstractIslandAPIClientFactory, @@ -198,6 +199,31 @@ class HTTPIslandAPIClient(IIslandAPIClient): except json.JSONDecodeError as e: raise IslandAPIRequestFailedError(e) + def get_credentials_for_propagation(self, island_server: str) -> Sequence[Credentials]: + try: + response = requests.get( # noqa: DUO123 + f"https://{island_server}/api/propagation-credentials", + verify=False, + timeout=SHORT_REQUEST_TIMEOUT, + ) + response.raise_for_status() + + return [Credentials(**credentials) for credentials in response.json()] + except ( + requests.exceptions.ConnectionError, + requests.exceptions.TooManyRedirects, + ) as e: + raise IslandAPIConnectionError(e) + except requests.exceptions.Timeout as e: + raise IslandAPITimeoutError(e) + except requests.exceptions.HTTPError as e: + if e.errno >= 500: + raise IslandAPIRequestFailedError(e) + else: + raise IslandAPIRequestError(e) + except json.JSONDecodeError as e: + raise IslandAPIRequestFailedError(e) + def _serialize_events(self, events: Sequence[AbstractAgentEvent]) -> JSONSerializable: serialized_events: List[JSONSerializable] = [] diff --git a/monkey/infection_monkey/island_api_client/i_island_api_client.py b/monkey/infection_monkey/island_api_client/i_island_api_client.py index 064511a73..409a1b265 100644 --- a/monkey/infection_monkey/island_api_client/i_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/i_island_api_client.py @@ -6,6 +6,7 @@ from common.agent_events import AbstractAgentEvent from common import AgentRegistrationData from common.agent_configuration import AgentConfiguration +from common.credentials import Credentials class IIslandAPIClient(ABC): @@ -132,3 +133,16 @@ class IIslandAPIClient(ABC): :raises IslandAPITimeoutError: If the command timed out :return: Agent configuration """ + + @abstractmethod + def get_credentials_for_propagation(self, island_server: str) -> Sequence[Credentials]: + """ + Get credentials from the island + + :param island_server: The server to query + :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 + :return: Credentials + """ diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py index d01a6a233..cce8791e1 100644 --- a/monkey/infection_monkey/master/control_channel.py +++ b/monkey/infection_monkey/master/control_channel.py @@ -2,12 +2,10 @@ import logging from typing import Optional, Sequence from uuid import UUID -import requests from urllib3 import disable_warnings from common import AgentRegistrationData from common.agent_configuration import AgentConfiguration -from common.common_consts.timeouts import SHORT_REQUEST_TIMEOUT from common.credentials import Credentials from common.network.network_utils import get_network_interfaces from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError @@ -76,23 +74,14 @@ class ControlChannel(IControlChannel): raise IslandCommunicationError(e) def get_credentials_for_propagation(self) -> Sequence[Credentials]: - propagation_credentials_url = ( - f"https://{self._control_channel_server}/api/propagation-credentials" - ) try: - response = requests.get( # noqa: DUO123 - propagation_credentials_url, - verify=False, - timeout=SHORT_REQUEST_TIMEOUT, + return self._island_api_client.get_credentials_for_propagation( + self._control_channel_server ) - response.raise_for_status() - - return [Credentials(**credentials) for credentials in response.json()] except ( - requests.exceptions.JSONDecodeError, - requests.exceptions.ConnectionError, - requests.exceptions.Timeout, - requests.exceptions.TooManyRedirects, - requests.exceptions.HTTPError, + IslandAPIConnectionError, + IslandAPIRequestError, + IslandAPIRequestFailedError, + IslandAPITimeoutError, ) as e: raise IslandCommunicationError(e) diff --git a/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py b/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py index dbb5ffb75..ff5c276f3 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py +++ b/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py @@ -14,6 +14,12 @@ from infection_monkey.master.control_channel import ControlChannel SERVER = "server" AGENT_ID = "agent" +CONTROL_CHANNEL_API_ERRORS = [ + IslandAPIConnectionError, + IslandAPIRequestError, + IslandAPIRequestFailedError, + IslandAPITimeoutError, +] @pytest.fixture @@ -96,17 +102,24 @@ def test_control_channel__get_config(control_channel, island_api_client): assert island_api_client.get_config.called_once() -@pytest.mark.parametrize( - "api_error", - [ - IslandAPIConnectionError, - IslandAPIRequestError, - IslandAPIRequestFailedError, - IslandAPITimeoutError, - ], -) +@pytest.mark.parametrize("api_error", CONTROL_CHANNEL_API_ERRORS) def test_control_channel__get_config_raises_error(control_channel, island_api_client, api_error): island_api_client.get_config.side_effect = api_error() with pytest.raises(IslandCommunicationError): control_channel.get_config() + + +def test_control_channel__get_credentials_for_propagation(control_channel, island_api_client): + control_channel.get_credentials_for_propagation() + assert island_api_client.get_credentials_for_propagation.called_once() + + +@pytest.mark.parametrize("api_error", CONTROL_CHANNEL_API_ERRORS) +def test_control_channel__get_credentials_for_propagation_raises_error( + control_channel, island_api_client, api_error +): + island_api_client.get_credentials_for_propagation.side_effect = api_error() + + with pytest.raises(IslandCommunicationError): + control_channel.get_credentials_for_propagation() From 42633c066f3f96fd5a7fd53cb2861286b3da7a25 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Mon, 19 Sep 2022 20:19:29 +0000 Subject: [PATCH 05/20] Agent: Reduce duplication due to error handling --- .../http_island_api_client.py | 139 +++++++----------- 1 file changed, 51 insertions(+), 88 deletions(-) diff --git a/monkey/infection_monkey/island_api_client/http_island_api_client.py b/monkey/infection_monkey/island_api_client/http_island_api_client.py index e5e338c4e..17e2f9284 100644 --- a/monkey/infection_monkey/island_api_client/http_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/http_island_api_client.py @@ -35,7 +35,7 @@ def handle_island_errors(fn): def decorated(*args, **kwargs): try: return fn(*args, **kwargs) - except requests.exceptions.ConnectionError as err: + except (requests.exceptions.ConnectionError, requests.exceptions.TooManyRedirects) as err: raise IslandAPIConnectionError(err) except requests.exceptions.HTTPError as err: if 400 <= err.response.status_code < 500: @@ -54,6 +54,17 @@ def handle_island_errors(fn): return decorated +def convert_json_error_to_island_api_error(fn): + @functools.wraps(fn) + def wrapper(*args, **kwargs): + try: + fn(*args, **kwargs) + except json.JSONDecodeError as e: + raise IslandAPIRequestFailedError(e) + + return wrapper + + class HTTPIslandAPIClient(IIslandAPIClient): """ A client for the Island's HTTP API @@ -125,104 +136,56 @@ class HTTPIslandAPIClient(IIslandAPIClient): response.raise_for_status() def register_agent(self, agent_registration_data: AgentRegistrationData): - try: - url = f"https://{agent_registration_data.cc_server}/api/agents" - response = requests.post( # noqa: DUO123 - url, - json=agent_registration_data.dict(simplify=True), - verify=False, - timeout=SHORT_REQUEST_TIMEOUT, - ) - response.raise_for_status() - except ( - requests.exceptions.ConnectionError, - requests.exceptions.TooManyRedirects, - requests.exceptions.HTTPError, - ) as e: - raise IslandAPIConnectionError(e) - except requests.exceptions.Timeout as e: - raise IslandAPITimeoutError(e) + url = f"https://{agent_registration_data.cc_server}/api/agents" + response = requests.post( # noqa: DUO123 + url, + json=agent_registration_data.dict(simplify=True), + verify=False, + timeout=SHORT_REQUEST_TIMEOUT, + ) + response.raise_for_status() + @handle_island_errors + @convert_json_error_to_island_api_error def should_agent_stop(self, island_server: str, agent_id: str) -> bool: - try: - url = f"https://{island_server}/api/monkey-control" f"/needs-to-stop/{agent_id}" - response = requests.get( # noqa: DUO123 - url, - verify=False, - timeout=SHORT_REQUEST_TIMEOUT, - ) - response.raise_for_status() + url = f"https://{island_server}/api/monkey-control" f"/needs-to-stop/{agent_id}" + response = requests.get( # noqa: DUO123 + url, + verify=False, + timeout=SHORT_REQUEST_TIMEOUT, + ) + response.raise_for_status() - json_response = json.loads(response.content.decode()) - return json_response["stop_agent"] - except ( - requests.exceptions.ConnectionError, - requests.exceptions.TooManyRedirects, - ) as e: - raise IslandAPIConnectionError(e) - except requests.exceptions.Timeout as e: - raise IslandAPITimeoutError(e) - except requests.exceptions.HTTPError as e: - if e.errno >= 500: - raise IslandAPIRequestFailedError(e) - else: - raise IslandAPIRequestError(e) - except json.JSONDecodeError as e: - raise IslandAPIRequestFailedError(e) + json_response = json.loads(response.content.decode()) + return json_response["stop_agent"] + @handle_island_errors + @convert_json_error_to_island_api_error def get_config(self, island_server: str) -> AgentConfiguration: - try: - response = requests.get( # noqa: DUO123 - f"https://{island_server}/api/agent-configuration", - verify=False, - timeout=SHORT_REQUEST_TIMEOUT, - ) - response.raise_for_status() + response = requests.get( # noqa: DUO123 + f"https://{island_server}/api/agent-configuration", + verify=False, + timeout=SHORT_REQUEST_TIMEOUT, + ) + response.raise_for_status() - config_dict = json.loads(response.text) + config_dict = json.loads(response.text) - logger.debug(f"Received configuration:\n{pformat(config_dict)}") + logger.debug(f"Received configuration:\n{pformat(config_dict)}") - return AgentConfiguration(**config_dict) - except ( - requests.exceptions.ConnectionError, - requests.exceptions.TooManyRedirects, - ) as e: - raise IslandAPIConnectionError(e) - except requests.exceptions.Timeout as e: - raise IslandAPITimeoutError(e) - except requests.exceptions.HTTPError as e: - if e.errno >= 500: - raise IslandAPIRequestFailedError(e) - else: - raise IslandAPIRequestError(e) - except json.JSONDecodeError as e: - raise IslandAPIRequestFailedError(e) + return AgentConfiguration(**config_dict) + @handle_island_errors + @convert_json_error_to_island_api_error def get_credentials_for_propagation(self, island_server: str) -> Sequence[Credentials]: - try: - response = requests.get( # noqa: DUO123 - f"https://{island_server}/api/propagation-credentials", - verify=False, - timeout=SHORT_REQUEST_TIMEOUT, - ) - response.raise_for_status() + response = requests.get( # noqa: DUO123 + f"https://{island_server}/api/propagation-credentials", + verify=False, + timeout=SHORT_REQUEST_TIMEOUT, + ) + response.raise_for_status() - return [Credentials(**credentials) for credentials in response.json()] - except ( - requests.exceptions.ConnectionError, - requests.exceptions.TooManyRedirects, - ) as e: - raise IslandAPIConnectionError(e) - except requests.exceptions.Timeout as e: - raise IslandAPITimeoutError(e) - except requests.exceptions.HTTPError as e: - if e.errno >= 500: - raise IslandAPIRequestFailedError(e) - else: - raise IslandAPIRequestError(e) - except json.JSONDecodeError as e: - raise IslandAPIRequestFailedError(e) + return [Credentials(**credentials) for credentials in response.json()] def _serialize_events(self, events: Sequence[AbstractAgentEvent]) -> JSONSerializable: serialized_events: List[JSONSerializable] = [] From 51ce974c4391cfdcf9cc046719169625503615d1 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Mon, 19 Sep 2022 20:51:05 +0000 Subject: [PATCH 06/20] Agent: Reduce duplication due to error handling --- .../master/control_channel.py | 63 +++++++++---------- 1 file changed, 28 insertions(+), 35 deletions(-) diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py index cce8791e1..6a972b6fa 100644 --- a/monkey/infection_monkey/master/control_channel.py +++ b/monkey/infection_monkey/master/control_channel.py @@ -1,4 +1,5 @@ import logging +from functools import wraps from typing import Optional, Sequence from uuid import UUID @@ -24,12 +25,30 @@ disable_warnings() # noqa: DUO131 logger = logging.getLogger(__name__) +def handle_island_api_errors(func): + @wraps(func) + def wrapper(*args, **kwargs): + try: + print(args) + func(*args, **kwargs) + except ( + IslandAPIConnectionError, + IslandAPIRequestError, + IslandAPIRequestFailedError, + IslandAPITimeoutError, + ) as e: + raise IslandCommunicationError(e) + + return wrapper + + class ControlChannel(IControlChannel): def __init__(self, server: str, agent_id: str, api_client: IIslandAPIClient): self._agent_id = agent_id self._control_channel_server = server self._island_api_client = api_client + @handle_island_api_errors def register_agent(self, parent: Optional[UUID] = None): agent_registration_data = AgentRegistrationData( id=get_agent_id(), @@ -41,47 +60,21 @@ class ControlChannel(IControlChannel): network_interfaces=get_network_interfaces(), ) - try: - self._island_api_client.register_agent(agent_registration_data) - except (IslandAPIConnectionError, IslandAPITimeoutError) as e: - raise IslandCommunicationError(e) + self._island_api_client.register_agent(agent_registration_data) + @handle_island_api_errors def should_agent_stop(self) -> bool: if not self._control_channel_server: logger.error("Agent should stop because it can't connect to the C&C server.") return True - try: - return self._island_api_client.should_agent_stop( - self._control_channel_server, self._agent_id - ) - except ( - IslandAPIConnectionError, - IslandAPIRequestError, - IslandAPIRequestFailedError, - IslandAPITimeoutError, - ) as e: - raise IslandCommunicationError(e) + return self._island_api_client.should_agent_stop( + self._control_channel_server, self._agent_id + ) + @handle_island_api_errors def get_config(self) -> AgentConfiguration: - try: - return self._island_api_client.get_config(self._control_channel_server) - except ( - IslandAPIConnectionError, - IslandAPIRequestError, - IslandAPIRequestFailedError, - IslandAPITimeoutError, - ) as e: - raise IslandCommunicationError(e) + return self._island_api_client.get_config(self._control_channel_server) + @handle_island_api_errors def get_credentials_for_propagation(self) -> Sequence[Credentials]: - try: - return self._island_api_client.get_credentials_for_propagation( - self._control_channel_server - ) - except ( - IslandAPIConnectionError, - IslandAPIRequestError, - IslandAPIRequestFailedError, - IslandAPITimeoutError, - ) as e: - raise IslandCommunicationError(e) + return self._island_api_client.get_credentials_for_propagation(self._control_channel_server) From 3e96a30a0999fc6c0f98a59cdb7821107d9b7ba6 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Mon, 19 Sep 2022 20:54:26 +0000 Subject: [PATCH 07/20] UT: Parametrize errors in tests --- .../master/test_control_channel.py | 50 +++---------------- 1 file changed, 8 insertions(+), 42 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py b/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py index ff5c276f3..2521442c6 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py +++ b/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py @@ -38,19 +38,11 @@ def test_control_channel__register_agent(control_channel, island_api_client): assert island_api_client.register_agent.called_once() -def test_control_channel__register_agent_raises_on_connection_error( - control_channel, island_api_client +@pytest.mark.parametrize("api_error", [IslandAPIConnectionError, IslandAPITimeoutError]) +def test_control_channel__register_agent_raises_error( + control_channel, island_api_client, api_error ): - island_api_client.register_agent.side_effect = IslandAPIConnectionError() - - with pytest.raises(IslandCommunicationError): - control_channel.register_agent() - - -def test_control_channel__register_agent_raises_on_timeout_error( - control_channel, island_api_client -): - island_api_client.register_agent.side_effect = IslandAPITimeoutError() + island_api_client.register_agent.side_effect = api_error() with pytest.raises(IslandCommunicationError): control_channel.register_agent() @@ -61,37 +53,11 @@ def test_control_channel__should_agent_stop(control_channel, island_api_client): assert island_api_client.should_agent_stop.called_once() -def test_control_channel__should_agent_stop_raises_on_connection_error( - control_channel, island_api_client +@pytest.mark.parametrize("api_error", CONTROL_CHANNEL_API_ERRORS) +def test_control_channel__should_agent_stop_raises_error( + control_channel, island_api_client, api_error ): - island_api_client.should_agent_stop.side_effect = IslandAPIConnectionError() - - with pytest.raises(IslandCommunicationError): - control_channel.should_agent_stop() - - -def test_control_channel__should_agent_stop_raises_on_timeout_error( - control_channel, island_api_client -): - island_api_client.should_agent_stop.side_effect = IslandAPITimeoutError() - - with pytest.raises(IslandCommunicationError): - control_channel.should_agent_stop() - - -def test_control_channel__should_agent_stop_raises_on_request_error( - control_channel, island_api_client -): - island_api_client.should_agent_stop.side_effect = IslandAPIRequestError() - - with pytest.raises(IslandCommunicationError): - control_channel.should_agent_stop() - - -def test_control_channel__should_agent_stop_raises_on_request_failed_error( - control_channel, island_api_client -): - island_api_client.should_agent_stop.side_effect = IslandAPIRequestFailedError() + island_api_client.should_agent_stop.side_effect = api_error() with pytest.raises(IslandCommunicationError): control_channel.should_agent_stop() From 2dc4871a7d7c2cd770d018f09fbc7478df13820d Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Mon, 19 Sep 2022 21:01:23 +0000 Subject: [PATCH 08/20] Agent: Pass IslandAPIClient to the ControlChannel --- monkey/infection_monkey/monkey.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 50ed3a3cb..403617721 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -120,6 +120,7 @@ class InfectionMonkey: self._control_client = ControlClient( server_address=server, island_api_client=self._island_api_client ) + self._control_channel = ControlChannel(server, GUID, self._island_api_client) # TODO Refactor the telemetry messengers to accept control client # and remove control_client_object @@ -195,7 +196,7 @@ class InfectionMonkey: run_aws_environment_check(self._telemetry_messenger) - should_stop = ControlChannel(self._control_client.server_address, GUID).should_agent_stop() + should_stop = self._control_channel.should_agent_stop() if should_stop: logger.info("The Monkey Island has instructed this agent to stop") return @@ -211,7 +212,6 @@ class InfectionMonkey: if firewall.is_enabled(): firewall.add_firewall_rule() - self._control_channel = ControlChannel(self._control_client.server_address, GUID) self._control_channel.register_agent(self._opts.parent) config = self._control_channel.get_config() From 8ebcd2ea3307711cd7475c7289fce017c4cf9bd6 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Tue, 20 Sep 2022 15:50:39 +0000 Subject: [PATCH 09/20] Island: Update register_agent errors --- .../http_island_api_client.py | 1 + .../island_api_client/i_island_api_client.py | 2 + .../test_http_island_api_client.py | 47 ++++++++++++++++++- 3 files changed, 48 insertions(+), 2 deletions(-) diff --git a/monkey/infection_monkey/island_api_client/http_island_api_client.py b/monkey/infection_monkey/island_api_client/http_island_api_client.py index 17e2f9284..f277ad9c0 100644 --- a/monkey/infection_monkey/island_api_client/http_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/http_island_api_client.py @@ -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 diff --git a/monkey/infection_monkey/island_api_client/i_island_api_client.py b/monkey/infection_monkey/island_api_client/i_island_api_client.py index 409a1b265..2f5edc8f4 100644 --- a/monkey/infection_monkey/island_api_client/i_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/i_island_api_client.py @@ -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 """ diff --git a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py index 1d0ac17fc..d4acc1c79 100644 --- a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py +++ b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py @@ -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) From 44d8dbeb5cc7d9f09c477909f5bd605aba1729b1 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Tue, 20 Sep 2022 15:57:19 +0000 Subject: [PATCH 10/20] Island: Get rid of server param in IslandAPIClient --- .../island_api_client/http_island_api_client.py | 14 +++++++------- .../island_api_client/i_island_api_client.py | 9 +++------ monkey/infection_monkey/master/control_channel.py | 8 +++----- 3 files changed, 13 insertions(+), 18 deletions(-) diff --git a/monkey/infection_monkey/island_api_client/http_island_api_client.py b/monkey/infection_monkey/island_api_client/http_island_api_client.py index f277ad9c0..bf6ba9355 100644 --- a/monkey/infection_monkey/island_api_client/http_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/http_island_api_client.py @@ -137,7 +137,7 @@ class HTTPIslandAPIClient(IIslandAPIClient): @handle_island_errors def register_agent(self, agent_registration_data: AgentRegistrationData): - url = f"https://{agent_registration_data.cc_server}/api/agents" + url = f"{self._api_url}/agents" response = requests.post( # noqa: DUO123 url, json=agent_registration_data.dict(simplify=True), @@ -148,8 +148,8 @@ class HTTPIslandAPIClient(IIslandAPIClient): @handle_island_errors @convert_json_error_to_island_api_error - def should_agent_stop(self, island_server: str, agent_id: str) -> bool: - url = f"https://{island_server}/api/monkey-control" f"/needs-to-stop/{agent_id}" + def should_agent_stop(self, agent_id: str) -> bool: + url = f"{self._api_url}/monkey-control/needs-to-stop/{agent_id}" response = requests.get( # noqa: DUO123 url, verify=False, @@ -162,9 +162,9 @@ class HTTPIslandAPIClient(IIslandAPIClient): @handle_island_errors @convert_json_error_to_island_api_error - def get_config(self, island_server: str) -> AgentConfiguration: + def get_config(self) -> AgentConfiguration: response = requests.get( # noqa: DUO123 - f"https://{island_server}/api/agent-configuration", + f"{self._api_url}/agent-configuration", verify=False, timeout=SHORT_REQUEST_TIMEOUT, ) @@ -178,9 +178,9 @@ class HTTPIslandAPIClient(IIslandAPIClient): @handle_island_errors @convert_json_error_to_island_api_error - def get_credentials_for_propagation(self, island_server: str) -> Sequence[Credentials]: + def get_credentials_for_propagation(self) -> Sequence[Credentials]: response = requests.get( # noqa: DUO123 - f"https://{island_server}/api/propagation-credentials", + f"{self._api_url}/propagation-credentials", verify=False, timeout=SHORT_REQUEST_TIMEOUT, ) diff --git a/monkey/infection_monkey/island_api_client/i_island_api_client.py b/monkey/infection_monkey/island_api_client/i_island_api_client.py index 2f5edc8f4..c77054e84 100644 --- a/monkey/infection_monkey/island_api_client/i_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/i_island_api_client.py @@ -110,11 +110,10 @@ class IIslandAPIClient(ABC): """ @abstractmethod - def should_agent_stop(self, island_server: str, agent_id: str) -> bool: + def should_agent_stop(self, agent_id: str) -> bool: """ Check with the island to see if the agent should stop - :param island_server: The server to query :param agent_id: The agent identifier for the agent to check :raises IslandAPIConnectionError: If the client could not connect to the island :raises IslandAPIRequestError: If there was a problem with the client request @@ -124,11 +123,10 @@ class IIslandAPIClient(ABC): """ @abstractmethod - def get_config(self, island_server: str) -> AgentConfiguration: + def get_config(self) -> AgentConfiguration: """ Get agent configuration from the island - :param island_server: The server to query :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 @@ -137,11 +135,10 @@ class IIslandAPIClient(ABC): """ @abstractmethod - def get_credentials_for_propagation(self, island_server: str) -> Sequence[Credentials]: + def get_credentials_for_propagation(self) -> Sequence[Credentials]: """ Get credentials from the island - :param island_server: The server to query :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 diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py index 6a972b6fa..7a2dd547a 100644 --- a/monkey/infection_monkey/master/control_channel.py +++ b/monkey/infection_monkey/master/control_channel.py @@ -67,14 +67,12 @@ class ControlChannel(IControlChannel): if not self._control_channel_server: logger.error("Agent should stop because it can't connect to the C&C server.") return True - return self._island_api_client.should_agent_stop( - self._control_channel_server, self._agent_id - ) + return self._island_api_client.should_agent_stop(self._agent_id) @handle_island_api_errors def get_config(self) -> AgentConfiguration: - return self._island_api_client.get_config(self._control_channel_server) + return self._island_api_client.get_config() @handle_island_api_errors def get_credentials_for_propagation(self) -> Sequence[Credentials]: - return self._island_api_client.get_credentials_for_propagation(self._control_channel_server) + return self._island_api_client.get_credentials_for_propagation() From fd08212763dddcba4701095bc7309cb5e8f45742 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Tue, 20 Sep 2022 16:30:24 +0000 Subject: [PATCH 11/20] UT: Add test for should_agent_stop --- .../http_island_api_client.py | 8 ++++ .../test_http_island_api_client.py | 45 +++++++++++++++++++ 2 files changed, 53 insertions(+) diff --git a/monkey/infection_monkey/island_api_client/http_island_api_client.py b/monkey/infection_monkey/island_api_client/http_island_api_client.py index bf6ba9355..d1e13778b 100644 --- a/monkey/infection_monkey/island_api_client/http_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/http_island_api_client.py @@ -35,6 +35,14 @@ def handle_island_errors(fn): def decorated(*args, **kwargs): try: return fn(*args, **kwargs) + except ( + IslandAPIConnectionError, + IslandAPIError, + IslandAPIRequestError, + IslandAPIRequestFailedError, + IslandAPITimeoutError, + ) as e: + raise e except (requests.exceptions.ConnectionError, requests.exceptions.TooManyRedirects) as err: raise IslandAPIConnectionError(err) except requests.exceptions.HTTPError as err: diff --git a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py index d4acc1c79..033936ba2 100644 --- a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py +++ b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py @@ -39,6 +39,7 @@ 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" ISLAND_REGISTER_AGENT_URI = f"https://{SERVER}/api/agents" +ISLAND_AGENT_STOP_URI = f"https://{SERVER}/api/monkey-control/needs-to-stop/{AGENT_ID}" class Event1(AbstractAgentEvent): @@ -318,3 +319,47 @@ def test_island_api_client_register_agent__status_code(status_code, expected_err with pytest.raises(expected_error): m.post(ISLAND_REGISTER_AGENT_URI, status_code=status_code) island_api_client.register_agent(AGENT_REGISTRATION) + + +@pytest.mark.parametrize( + "actual_error, expected_error", + [ + (requests.exceptions.ConnectionError, IslandAPIConnectionError), + (TimeoutError, IslandAPITimeoutError), + ], +) +def test_island_api_client__should_agent_stop(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.get(ISLAND_AGENT_STOP_URI, exc=actual_error) + island_api_client.should_agent_stop(AGENT_ID) + + +@pytest.mark.parametrize( + "status_code, expected_error", + [ + (401, IslandAPIRequestError), + (501, IslandAPIRequestFailedError), + ], +) +def test_island_api_client_should_agent_stop__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.get(ISLAND_AGENT_STOP_URI, status_code=status_code) + island_api_client.should_agent_stop(AGENT_ID) + + +def test_island_api_client_should_agent_stop__bad_json(): + with requests_mock.Mocker() as m: + m.get(ISLAND_URI) + island_api_client = HTTPIslandAPIClient(SERVER) + + with pytest.raises(IslandAPIRequestFailedError): + m.get(ISLAND_AGENT_STOP_URI, content=b"bad") + island_api_client.should_agent_stop(AGENT_ID) From 4c3a1ba89a48331275e4e27ef03c440ac09879ba Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Tue, 20 Sep 2022 16:33:48 +0000 Subject: [PATCH 12/20] UT: Add test for get_config --- .../test_http_island_api_client.py | 45 +++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py index 033936ba2..a013fcec7 100644 --- a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py +++ b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py @@ -40,6 +40,7 @@ ISLAND_GET_AGENT_BINARY_URI = f"https://{SERVER}/api/agent-binaries/{WINDOWS}" ISLAND_SEND_EVENTS_URI = f"https://{SERVER}/api/agent-events" 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" class Event1(AbstractAgentEvent): @@ -363,3 +364,47 @@ def test_island_api_client_should_agent_stop__bad_json(): with pytest.raises(IslandAPIRequestFailedError): m.get(ISLAND_AGENT_STOP_URI, content=b"bad") island_api_client.should_agent_stop(AGENT_ID) + + +@pytest.mark.parametrize( + "actual_error, expected_error", + [ + (requests.exceptions.ConnectionError, IslandAPIConnectionError), + (TimeoutError, IslandAPITimeoutError), + ], +) +def test_island_api_client__get_config(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.get(ISLAND_GET_CONFIG_URI, exc=actual_error) + island_api_client.get_config() + + +@pytest.mark.parametrize( + "status_code, expected_error", + [ + (401, IslandAPIRequestError), + (501, IslandAPIRequestFailedError), + ], +) +def test_island_api_client_get_config__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.get(ISLAND_GET_CONFIG_URI, status_code=status_code) + island_api_client.get_config() + + +def test_island_api_client_get_config__bad_json(): + with requests_mock.Mocker() as m: + m.get(ISLAND_URI) + island_api_client = HTTPIslandAPIClient(SERVER) + + with pytest.raises(IslandAPIRequestFailedError): + m.get(ISLAND_GET_CONFIG_URI, content=b"bad") + island_api_client.get_config() From 605e6ae4967ebc11e1c3adbb94298b43cdd94420 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Tue, 20 Sep 2022 16:36:33 +0000 Subject: [PATCH 13/20] UT: Add test for get_credentials_for_propagation --- .../test_http_island_api_client.py | 47 +++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py index a013fcec7..27d31610b 100644 --- a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py +++ b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py @@ -41,6 +41,7 @@ ISLAND_SEND_EVENTS_URI = f"https://{SERVER}/api/agent-events" 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" class Event1(AbstractAgentEvent): @@ -408,3 +409,49 @@ def test_island_api_client_get_config__bad_json(): with pytest.raises(IslandAPIRequestFailedError): m.get(ISLAND_GET_CONFIG_URI, content=b"bad") island_api_client.get_config() + + +@pytest.mark.parametrize( + "actual_error, expected_error", + [ + (requests.exceptions.ConnectionError, IslandAPIConnectionError), + (TimeoutError, IslandAPITimeoutError), + ], +) +def test_island_api_client__get_credentials_for_propagation(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.get(ISLAND_GET_PROPAGATION_CREDENTIALS_URI, exc=actual_error) + island_api_client.get_credentials_for_propagation() + + +@pytest.mark.parametrize( + "status_code, expected_error", + [ + (401, IslandAPIRequestError), + (501, IslandAPIRequestFailedError), + ], +) +def test_island_api_client_get_credentials_for_propagation__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.get(ISLAND_GET_PROPAGATION_CREDENTIALS_URI, status_code=status_code) + island_api_client.get_credentials_for_propagation() + + +def test_island_api_client_get_credentials_for_propagation__bad_json(): + with requests_mock.Mocker() as m: + m.get(ISLAND_URI) + island_api_client = HTTPIslandAPIClient(SERVER) + + with pytest.raises(IslandAPIRequestFailedError): + m.get(ISLAND_GET_PROPAGATION_CREDENTIALS_URI, content=b"bad") + island_api_client.get_credentials_for_propagation() From 2f5bcff631eb678890ea3ba2e0632d332f53658c Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Tue, 20 Sep 2022 18:04:16 +0000 Subject: [PATCH 14/20] UT: Use island_api_client fixture --- .../test_http_island_api_client.py | 50 +++++++++++-------- 1 file changed, 28 insertions(+), 22 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py index 27d31610b..03117b006 100644 --- a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py +++ b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py @@ -296,10 +296,10 @@ def test_island_api_client_send_events__status_code(island_api_client, status_co (TimeoutError, IslandAPITimeoutError), ], ) -def test_island_api_client__register_agent(actual_error, expected_error): +def test_island_api_client__register_agent(island_api_client, actual_error, expected_error): with requests_mock.Mocker() as m: m.get(ISLAND_URI) - island_api_client = HTTPIslandAPIClient(SERVER) + island_api_client.connect(SERVER) with pytest.raises(expected_error): m.post(ISLAND_REGISTER_AGENT_URI, exc=actual_error) @@ -313,10 +313,12 @@ def test_island_api_client__register_agent(actual_error, expected_error): (501, IslandAPIRequestFailedError), ], ) -def test_island_api_client_register_agent__status_code(status_code, expected_error): +def test_island_api_client_register_agent__status_code( + island_api_client, status_code, expected_error +): with requests_mock.Mocker() as m: m.get(ISLAND_URI) - island_api_client = HTTPIslandAPIClient(SERVER) + island_api_client.connect(SERVER) with pytest.raises(expected_error): m.post(ISLAND_REGISTER_AGENT_URI, status_code=status_code) @@ -330,10 +332,10 @@ def test_island_api_client_register_agent__status_code(status_code, expected_err (TimeoutError, IslandAPITimeoutError), ], ) -def test_island_api_client__should_agent_stop(actual_error, expected_error): +def test_island_api_client__should_agent_stop(island_api_client, actual_error, expected_error): with requests_mock.Mocker() as m: m.get(ISLAND_URI) - island_api_client = HTTPIslandAPIClient(SERVER) + island_api_client.connect(SERVER) with pytest.raises(expected_error): m.get(ISLAND_AGENT_STOP_URI, exc=actual_error) @@ -347,20 +349,22 @@ def test_island_api_client__should_agent_stop(actual_error, expected_error): (501, IslandAPIRequestFailedError), ], ) -def test_island_api_client_should_agent_stop__status_code(status_code, expected_error): +def test_island_api_client_should_agent_stop__status_code( + island_api_client, status_code, expected_error +): with requests_mock.Mocker() as m: m.get(ISLAND_URI) - island_api_client = HTTPIslandAPIClient(SERVER) + island_api_client.connect(SERVER) with pytest.raises(expected_error): m.get(ISLAND_AGENT_STOP_URI, status_code=status_code) island_api_client.should_agent_stop(AGENT_ID) -def test_island_api_client_should_agent_stop__bad_json(): +def test_island_api_client_should_agent_stop__bad_json(island_api_client): with requests_mock.Mocker() as m: m.get(ISLAND_URI) - island_api_client = HTTPIslandAPIClient(SERVER) + island_api_client.connect(SERVER) with pytest.raises(IslandAPIRequestFailedError): m.get(ISLAND_AGENT_STOP_URI, content=b"bad") @@ -374,10 +378,10 @@ def test_island_api_client_should_agent_stop__bad_json(): (TimeoutError, IslandAPITimeoutError), ], ) -def test_island_api_client__get_config(actual_error, expected_error): +def test_island_api_client__get_config(island_api_client, actual_error, expected_error): with requests_mock.Mocker() as m: m.get(ISLAND_URI) - island_api_client = HTTPIslandAPIClient(SERVER) + island_api_client.connect(SERVER) with pytest.raises(expected_error): m.get(ISLAND_GET_CONFIG_URI, exc=actual_error) @@ -391,20 +395,20 @@ def test_island_api_client__get_config(actual_error, expected_error): (501, IslandAPIRequestFailedError), ], ) -def test_island_api_client_get_config__status_code(status_code, expected_error): +def test_island_api_client_get_config__status_code(island_api_client, status_code, expected_error): with requests_mock.Mocker() as m: m.get(ISLAND_URI) - island_api_client = HTTPIslandAPIClient(SERVER) + island_api_client.connect(SERVER) with pytest.raises(expected_error): m.get(ISLAND_GET_CONFIG_URI, status_code=status_code) island_api_client.get_config() -def test_island_api_client_get_config__bad_json(): +def test_island_api_client_get_config__bad_json(island_api_client): with requests_mock.Mocker() as m: m.get(ISLAND_URI) - island_api_client = HTTPIslandAPIClient(SERVER) + island_api_client.connect(SERVER) with pytest.raises(IslandAPIRequestFailedError): m.get(ISLAND_GET_CONFIG_URI, content=b"bad") @@ -418,10 +422,12 @@ def test_island_api_client_get_config__bad_json(): (TimeoutError, IslandAPITimeoutError), ], ) -def test_island_api_client__get_credentials_for_propagation(actual_error, expected_error): +def test_island_api_client__get_credentials_for_propagation( + island_api_client, actual_error, expected_error +): with requests_mock.Mocker() as m: m.get(ISLAND_URI) - island_api_client = HTTPIslandAPIClient(SERVER) + island_api_client.connect(SERVER) with pytest.raises(expected_error): m.get(ISLAND_GET_PROPAGATION_CREDENTIALS_URI, exc=actual_error) @@ -436,21 +442,21 @@ def test_island_api_client__get_credentials_for_propagation(actual_error, expect ], ) def test_island_api_client_get_credentials_for_propagation__status_code( - status_code, expected_error + island_api_client, status_code, expected_error ): with requests_mock.Mocker() as m: m.get(ISLAND_URI) - island_api_client = HTTPIslandAPIClient(SERVER) + island_api_client.connect(SERVER) with pytest.raises(expected_error): m.get(ISLAND_GET_PROPAGATION_CREDENTIALS_URI, status_code=status_code) island_api_client.get_credentials_for_propagation() -def test_island_api_client_get_credentials_for_propagation__bad_json(): +def test_island_api_client_get_credentials_for_propagation__bad_json(island_api_client): with requests_mock.Mocker() as m: m.get(ISLAND_URI) - island_api_client = HTTPIslandAPIClient(SERVER) + island_api_client.connect(SERVER) with pytest.raises(IslandAPIRequestFailedError): m.get(ISLAND_GET_PROPAGATION_CREDENTIALS_URI, content=b"bad") From 1c6ca24a4755be5a241c3a3544965e83050c25b2 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Tue, 20 Sep 2022 18:21:55 +0000 Subject: [PATCH 15/20] Agent: Move register_agent out of ControlChannel --- monkey/infection_monkey/i_control_channel.py | 14 +------------ .../master/control_channel.py | 21 +------------------ monkey/infection_monkey/monkey.py | 18 ++++++++++++++-- .../master/test_control_channel.py | 15 ------------- 4 files changed, 18 insertions(+), 50 deletions(-) diff --git a/monkey/infection_monkey/i_control_channel.py b/monkey/infection_monkey/i_control_channel.py index 39075750a..25135231f 100644 --- a/monkey/infection_monkey/i_control_channel.py +++ b/monkey/infection_monkey/i_control_channel.py @@ -1,23 +1,11 @@ import abc -from typing import Optional, Sequence -from uuid import UUID +from typing import Sequence from common.agent_configuration import AgentConfiguration from common.credentials import Credentials class IControlChannel(metaclass=abc.ABCMeta): - @abc.abstractmethod - def register_agent(self, parent_id: Optional[UUID] = None): - """ - Registers this agent with the Island when this agent starts - - :param parent: The ID of the parent that spawned this agent, or None if this agent has no - parent - :raises IslandCommunicationError: If the agent cannot be successfully registered - """ - pass - @abc.abstractmethod def should_agent_stop(self) -> bool: """ diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py index 7a2dd547a..eddeb8090 100644 --- a/monkey/infection_monkey/master/control_channel.py +++ b/monkey/infection_monkey/master/control_channel.py @@ -1,14 +1,11 @@ import logging from functools import wraps -from typing import Optional, Sequence -from uuid import UUID +from typing import Sequence from urllib3 import disable_warnings -from common import AgentRegistrationData from common.agent_configuration import AgentConfiguration from common.credentials import Credentials -from common.network.network_utils import get_network_interfaces from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError from infection_monkey.island_api_client import ( IIslandAPIClient, @@ -17,8 +14,6 @@ from infection_monkey.island_api_client import ( IslandAPIRequestFailedError, IslandAPITimeoutError, ) -from infection_monkey.utils import agent_process -from infection_monkey.utils.ids import get_agent_id, get_machine_id disable_warnings() # noqa: DUO131 @@ -48,20 +43,6 @@ class ControlChannel(IControlChannel): self._control_channel_server = server self._island_api_client = api_client - @handle_island_api_errors - def register_agent(self, parent: Optional[UUID] = None): - agent_registration_data = AgentRegistrationData( - id=get_agent_id(), - machine_hardware_id=get_machine_id(), - start_time=agent_process.get_start_time(), - # parent_id=parent, - parent_id=None, # None for now, until we change GUID to UUID - cc_server=self._control_channel_server, - network_interfaces=get_network_interfaces(), - ) - - self._island_api_client.register_agent(agent_registration_data) - @handle_island_api_errors def should_agent_stop(self) -> bool: if not self._control_channel_server: diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 403617721..460c07592 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -14,6 +14,7 @@ from common.agent_event_serializers import ( register_common_agent_event_serializers, ) from common.agent_events import CredentialsStolenEvent +from common.agent_registration_data import AgentRegistrationData from common.event_queue import IAgentEventQueue, PyPubSubAgentEventQueue from common.network.network_utils import ( address_to_ip_port, @@ -88,9 +89,11 @@ from infection_monkey.telemetry.messengers.legacy_telemetry_messenger_adapter im LegacyTelemetryMessengerAdapter, ) from infection_monkey.telemetry.state_telem import StateTelem +from infection_monkey.utils import agent_process from infection_monkey.utils.aws_environment_check import run_aws_environment_check from infection_monkey.utils.environment import is_windows_os from infection_monkey.utils.file_utils import mark_file_for_deletion_on_windows +from infection_monkey.utils.ids import get_agent_id, get_machine_id from infection_monkey.utils.monkey_dir import ( create_monkey_dir, get_monkey_dir_path, @@ -121,6 +124,7 @@ class InfectionMonkey: server_address=server, island_api_client=self._island_api_client ) self._control_channel = ControlChannel(server, GUID, self._island_api_client) + self._register_agent() # TODO Refactor the telemetry messengers to accept control client # and remove control_client_object @@ -166,6 +170,18 @@ class InfectionMonkey: return server, island_api_client + def _register_agent(self, server: str): + agent_registration_data = AgentRegistrationData( + id=get_agent_id(), + machine_hardware_id=get_machine_id(), + start_time=agent_process.get_start_time(), + # parent_id=parent, + parent_id=None, # None for now, until we change GUID to UUID + cc_server=server, + network_interfaces=get_network_interfaces(), + ) + self._island_api_client.register_agent(agent_registration_data) + def _select_server( self, server_clients: Mapping[str, IIslandAPIClient] ) -> Tuple[Optional[str], Optional[IIslandAPIClient]]: @@ -212,8 +228,6 @@ class InfectionMonkey: if firewall.is_enabled(): firewall.add_firewall_rule() - self._control_channel.register_agent(self._opts.parent) - config = self._control_channel.get_config() relay_port = get_free_tcp_port() diff --git a/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py b/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py index 2521442c6..658635615 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py +++ b/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py @@ -33,21 +33,6 @@ def control_channel(island_api_client) -> ControlChannel: return ControlChannel(SERVER, AGENT_ID, island_api_client) -def test_control_channel__register_agent(control_channel, island_api_client): - control_channel.register_agent() - assert island_api_client.register_agent.called_once() - - -@pytest.mark.parametrize("api_error", [IslandAPIConnectionError, IslandAPITimeoutError]) -def test_control_channel__register_agent_raises_error( - control_channel, island_api_client, api_error -): - island_api_client.register_agent.side_effect = api_error() - - with pytest.raises(IslandCommunicationError): - control_channel.register_agent() - - def test_control_channel__should_agent_stop(control_channel, island_api_client): control_channel.should_agent_stop() assert island_api_client.should_agent_stop.called_once() From 417d167026086ff8d07614787b297a18b1ef6365 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Tue, 20 Sep 2022 18:27:17 +0000 Subject: [PATCH 16/20] Agent: Simplify error handling in HTTPIslandAPIClient --- .../island_api_client/http_island_api_client.py | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/monkey/infection_monkey/island_api_client/http_island_api_client.py b/monkey/infection_monkey/island_api_client/http_island_api_client.py index d1e13778b..819b774d0 100644 --- a/monkey/infection_monkey/island_api_client/http_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/http_island_api_client.py @@ -35,14 +35,8 @@ def handle_island_errors(fn): def decorated(*args, **kwargs): try: return fn(*args, **kwargs) - except ( - IslandAPIConnectionError, - IslandAPIError, - IslandAPIRequestError, - IslandAPIRequestFailedError, - IslandAPITimeoutError, - ) as e: - raise e + except IslandAPIError as err: + raise err except (requests.exceptions.ConnectionError, requests.exceptions.TooManyRedirects) as err: raise IslandAPIConnectionError(err) except requests.exceptions.HTTPError as err: @@ -67,8 +61,8 @@ def convert_json_error_to_island_api_error(fn): def wrapper(*args, **kwargs): try: fn(*args, **kwargs) - except json.JSONDecodeError as e: - raise IslandAPIRequestFailedError(e) + except json.JSONDecodeError as err: + raise IslandAPIRequestFailedError(err) return wrapper From 022630ddc119566f062d5af877de16f86843f1cf Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Tue, 20 Sep 2022 18:32:20 +0000 Subject: [PATCH 17/20] Agent: Simplify error handling in ControlChannel --- .../infection_monkey/master/control_channel.py | 17 +++-------------- 1 file changed, 3 insertions(+), 14 deletions(-) diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py index eddeb8090..3b8d64740 100644 --- a/monkey/infection_monkey/master/control_channel.py +++ b/monkey/infection_monkey/master/control_channel.py @@ -7,13 +7,7 @@ from urllib3 import disable_warnings from common.agent_configuration import AgentConfiguration from common.credentials import Credentials from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError -from infection_monkey.island_api_client import ( - IIslandAPIClient, - IslandAPIConnectionError, - IslandAPIRequestError, - IslandAPIRequestFailedError, - IslandAPITimeoutError, -) +from infection_monkey.island_api_client import IIslandAPIClient, IslandAPIError disable_warnings() # noqa: DUO131 @@ -26,13 +20,8 @@ def handle_island_api_errors(func): try: print(args) func(*args, **kwargs) - except ( - IslandAPIConnectionError, - IslandAPIRequestError, - IslandAPIRequestFailedError, - IslandAPITimeoutError, - ) as e: - raise IslandCommunicationError(e) + except IslandAPIError as err: + raise IslandCommunicationError(err) return wrapper From 2a1396c82da7615f2145ef64f1b6b28874eec525 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Tue, 20 Sep 2022 18:33:05 +0000 Subject: [PATCH 18/20] Agent: Use response.json() in HTTPIslandAPIClient --- .../island_api_client/http_island_api_client.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/monkey/infection_monkey/island_api_client/http_island_api_client.py b/monkey/infection_monkey/island_api_client/http_island_api_client.py index 819b774d0..467f97252 100644 --- a/monkey/infection_monkey/island_api_client/http_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/http_island_api_client.py @@ -159,8 +159,7 @@ class HTTPIslandAPIClient(IIslandAPIClient): ) response.raise_for_status() - json_response = json.loads(response.content.decode()) - return json_response["stop_agent"] + return response.json()["stop_agent"] @handle_island_errors @convert_json_error_to_island_api_error @@ -172,7 +171,7 @@ class HTTPIslandAPIClient(IIslandAPIClient): ) response.raise_for_status() - config_dict = json.loads(response.text) + config_dict = response.json() logger.debug(f"Received configuration:\n{pformat(config_dict)}") From 53dc34f895d80c60904ee9a304690114a4db5f9b Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Tue, 20 Sep 2022 18:34:04 +0000 Subject: [PATCH 19/20] Agent: Fix typo in IIslandAPIClient --- .../island_api_client/i_island_api_client.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/monkey/infection_monkey/island_api_client/i_island_api_client.py b/monkey/infection_monkey/island_api_client/i_island_api_client.py index c77054e84..3ec15237d 100644 --- a/monkey/infection_monkey/island_api_client/i_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/i_island_api_client.py @@ -1,11 +1,9 @@ from abc import ABC, abstractmethod from typing import Optional, Sequence -from common import OperatingSystem -from common.agent_events import AbstractAgentEvent - -from common import AgentRegistrationData +from common import AgentRegistrationData, OperatingSystem from common.agent_configuration import AgentConfiguration +from common.agent_events import AbstractAgentEvent from common.credentials import Credentials @@ -17,7 +15,7 @@ class IIslandAPIClient(ABC): @abstractmethod def connect(self, island_server: str): """ - Connectto the island's API + Connect to the island's API :param island_server: The socket address of the API :raises IslandAPIConnectionError: If the client cannot successfully connect to the island From 897d1e39ac6724432e111541cd4a2ac589f006df Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 20 Sep 2022 14:45:17 -0400 Subject: [PATCH 20/20] Agent: Remove debug print from handle_island_api_errors() --- monkey/infection_monkey/master/control_channel.py | 1 - 1 file changed, 1 deletion(-) diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py index 3b8d64740..39fe94d1b 100644 --- a/monkey/infection_monkey/master/control_channel.py +++ b/monkey/infection_monkey/master/control_channel.py @@ -18,7 +18,6 @@ def handle_island_api_errors(func): @wraps(func) def wrapper(*args, **kwargs): try: - print(args) func(*args, **kwargs) except IslandAPIError as err: raise IslandCommunicationError(err)