From ccfc41fc2f35f4f7e8ccff7c855711b4d8a8efac Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Wed, 21 Sep 2022 17:54:33 +0000 Subject: [PATCH] Agent: Fix mypy issues in IslandAPIClient --- .../island_api_client/http_island_api_client.py | 8 ++++---- .../island_api_client/i_island_api_client.py | 5 ++--- 2 files changed, 6 insertions(+), 7 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 a0bebc5a0..3072a54b7 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 @@ -94,7 +94,7 @@ class HTTPIslandAPIClient(IIslandAPIClient): return response.content @handle_island_errors - def get_agent_binary(self, operating_system: OperatingSystem): + def get_agent_binary(self, operating_system: OperatingSystem) -> bytes: os_name = operating_system.value response = requests.get( # noqa: DUO123 f"{self._api_url}/agent-binaries/{os_name}", @@ -106,7 +106,7 @@ class HTTPIslandAPIClient(IIslandAPIClient): return response.content @handle_island_errors - def send_events(self, events: Sequence[JSONSerializable]): + def send_events(self, events: Sequence[AbstractAgentEvent]): response = requests.post( # noqa: DUO123 f"{self._api_url}/agent-events", json=self._serialize_events(events), @@ -132,9 +132,9 @@ class HTTPIslandAPIClient(IIslandAPIClient): class HTTPIslandAPIClientFactory(AbstractIslandAPIClientFactory): def __init__( self, - agent_event_serializer_registry: AgentEventSerializerRegistry = None, + agent_event_serializer_registry: AgentEventSerializerRegistry, ): self._agent_event_serializer_registry = agent_event_serializer_registry - def create_island_api_client(self): + def create_island_api_client(self) -> IIslandAPIClient: return HTTPIslandAPIClient(self._agent_event_serializer_registry) 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..006a2103f 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,5 +1,5 @@ from abc import ABC, abstractmethod -from typing import Optional, Sequence +from typing import Sequence from common import OperatingSystem from common.agent_events import AbstractAgentEvent @@ -60,7 +60,7 @@ class IIslandAPIClient(ABC): """ @abstractmethod - def get_agent_binary(self, operating_system: OperatingSystem) -> Optional[bytes]: + def get_agent_binary(self, operating_system: OperatingSystem) -> bytes: """ Get an agent binary for the given OS from the island @@ -74,7 +74,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