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

Agent: Fix type hints mypy found in relay/utils.py

This commit is contained in:
vakarisz 2022-09-22 17:33:00 +03:00
parent 9749984640
commit 3d7e9be150
1 changed files with 8 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
monkey/infection_monkey/network/relay/utils.py
View File

@ -2,7 +2,7 @@ import logging
import socket import socket
from contextlib import suppress from contextlib import suppress
from ipaddress import IPv4Address from ipaddress import IPv4Address
from typing import Dict, Iterable, Iterator, Mapping, MutableMapping, Optional, Tuple from typing import Dict, Iterable, Iterator, Optional
from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT
from common.network.network_utils import address_to_ip_port from common.network.network_utils import address_to_ip_port
@ -26,13 +26,15 @@ logger = logging.getLogger(__name__)
# practical purposes. Revisit this if it's not. # practical purposes. Revisit this if it's not.
NUM_FIND_SERVER_WORKERS = 32 NUM_FIND_SERVER_WORKERS = 32
IslandAPISearchResults = Dict[str, Optional[IIslandAPIClient]]
def find_available_island_apis( def find_available_island_apis(
servers: Iterable[str], island_api_client_factory: AbstractIslandAPIClientFactory servers: Iterable[str], island_api_client_factory: AbstractIslandAPIClientFactory
) -> Mapping[str, Optional[IIslandAPIClient]]: ) -> IslandAPISearchResults:
server_list = list(servers) server_list = list(servers)
server_iterator = ThreadSafeIterator(server_list.__iter__()) server_iterator = ThreadSafeIterator(server_list.__iter__())
server_results: Dict[str, Tuple[bool, IIslandAPIClient]] = {} server_results: IslandAPISearchResults = {}
run_worker_threads( run_worker_threads(
_find_island_server, _find_island_server,
@ -46,17 +48,17 @@ def find_available_island_apis(
def _find_island_server( def _find_island_server(
servers: Iterator[str], servers: Iterator[str],
server_status: MutableMapping[str, Optional[IIslandAPIClient]], server_results: IslandAPISearchResults,
island_api_client_factory: AbstractIslandAPIClientFactory, island_api_client_factory: AbstractIslandAPIClientFactory,
): ):
with suppress(StopIteration): with suppress(StopIteration):
server = next(servers) server = next(servers)
server_status[server] = _check_if_island_server(server, island_api_client_factory) server_results[server] = _check_if_island_server(server, island_api_client_factory)
def _check_if_island_server( def _check_if_island_server(
server: str, island_api_client_factory: AbstractIslandAPIClientFactory server: str, island_api_client_factory: AbstractIslandAPIClientFactory
) -> IIslandAPIClient: ) -> Optional[IIslandAPIClient]:
logger.debug(f"Trying to connect to server: {server}") logger.debug(f"Trying to connect to server: {server}")
try: try: