From 44d8dbeb5cc7d9f09c477909f5bd605aba1729b1 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Tue, 20 Sep 2022 15:57:19 +0000 Subject: [PATCH] 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()