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

Common: Make SocketAddress hashable

This commit is contained in:
Kekoa Kaaikala 2022-09-26 13:27:08 +00:00 committed by Shreya Malviya
parent c4804f06a9
commit 110542eeb8
4 changed files with 31 additions and 68 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
monkey/common/types.py
View File

@ -41,5 +41,8 @@ class SocketAddress(InfectionMonkeyBaseModel):
raise ValueError("SocketAddress requires a port")
return SocketAddress(ip=IPv4Address(ip), port=int(port))
def __hash__(self):
return hash(str(self))
def __str__(self):
return f"{self.ip}:{self.port}"

17
monkey/infection_monkey/monkey.py
View File

@ -115,7 +115,6 @@ class InfectionMonkey:
self._singleton = SystemSingleton()
self._opts = self._get_arguments(args)
self._server_strings = [str(s) for s in self._opts.servers]
self._agent_event_serializer_registry = self._setup_agent_event_serializers()
@ -155,7 +154,7 @@ class InfectionMonkey:
# TODO: By the time we finish 2292, _connect_to_island_api() may not need to return `server`
def _connect_to_island_api(self) -> Tuple[Optional[str], Optional[IIslandAPIClient]]:
logger.debug(f"Trying to wake up with servers: {', '.join(self._server_strings)}")
logger.debug(f"Trying to wake up with servers: {', '.join(map(str, self._opts.servers))}")
server_clients = find_available_island_apis(
self._opts.servers, HTTPIslandAPIClientFactory(self._agent_event_serializer_registry)
)
@ -166,13 +165,14 @@ class InfectionMonkey:
logger.info(f"Successfully connected to the island via {server}")
else:
raise Exception(
f"Failed to connect to the island via any known servers: {self._server_strings}"
"Failed to connect to the island via any known servers: "
f"[{', '.join(map(str, self._opts.servers))}]"
)
# NOTE: Since we pass the address for each of our interfaces to the exploited
# machines, is it possible for a machine to unintentionally unregister itself from the
# relay if it is able to connect to the relay over multiple interfaces?
servers_to_close = (s for s in self._server_strings if s != server and server_clients[s])
servers_to_close = (s for s in self._opts.servers if s != server and server_clients[s])
send_remove_from_waitlist_control_message_to_relays(servers_to_close)
return server, island_api_client
@ -192,9 +192,9 @@ class InfectionMonkey:
def _select_server(
self, server_clients: IslandAPISearchResults
) -> Tuple[Optional[SocketAddress], Optional[IIslandAPIClient]]:
for result in server_clients:
if result.client is not None:
return result.server, result.client
for server in self._opts.servers:
if server_clients[server] is not None:
return server, server_clients[server]
return None, None
@ -261,8 +261,9 @@ class InfectionMonkey:
return agent_event_serializer_registry
def _build_server_list(self, relay_port: int):
my_servers = [str(s) for s in self._opts.servers]
relay_servers = [f"{ip}:{relay_port}" for ip in get_my_ip_addresses()]
return self._server_strings + relay_servers
return my_servers + relay_servers
def _build_master(self, relay_port: int):
servers = self._build_server_list(relay_port)

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

@ -1,8 +1,7 @@
import logging
import socket
from contextlib import suppress
from dataclasses import dataclass
from typing import Dict, Iterable, Iterator, List, Optional, Tuple
from typing import Dict, Iterable, Iterator, Optional
from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT
from common.types import SocketAddress
@ -27,42 +26,34 @@ logger = logging.getLogger(__name__)
NUM_FIND_SERVER_WORKERS = 32
@dataclass
class IslandAPISearchResult:
server: SocketAddress
client: Optional[IIslandAPIClient]
IslandAPISearchResults = List[IslandAPISearchResult]
IslandAPISearchResults = Dict[SocketAddress, Optional[IIslandAPIClient]]
def find_available_island_apis(
servers: Iterable[SocketAddress], island_api_client_factory: AbstractIslandAPIClientFactory
) -> IslandAPISearchResults:
server_list = list(servers)
server_iterator = ThreadSafeIterator(enumerate(server_list.__iter__()))
results: Dict[int, IslandAPISearchResult] = {}
server_iterator = ThreadSafeIterator(server_list.__iter__())
server_results: IslandAPISearchResults = {}
run_worker_threads(
_find_island_server,
"FindIslandServer",
args=(server_iterator, results, island_api_client_factory),
args=(server_iterator, server_results, island_api_client_factory),
num_workers=NUM_FIND_SERVER_WORKERS,
)
return [results[i] for i in sorted(results.keys())]
return server_results
def _find_island_server(
servers: Iterator[Tuple[int, SocketAddress]],
server_results: Dict[int, IslandAPISearchResult],
servers: Iterator[SocketAddress],
server_results: IslandAPISearchResults,
island_api_client_factory: AbstractIslandAPIClientFactory,
):
with suppress(StopIteration):
index, server = next(servers)
server_results[index] = IslandAPISearchResult(
server, _check_if_island_server(server, island_api_client_factory)
)
server = next(servers)
server_results[server] = _check_if_island_server(server, island_api_client_factory)
def _check_if_island_server(

50
monkey/tests/unit_tests/infection_monkey/network/relay/test_utils.py
View File

@ -1,5 +1,3 @@
from typing import Callable, Optional
import pytest
import requests_mock
@ -10,7 +8,7 @@ from infection_monkey.island_api_client import (
IIslandAPIClient,
IslandAPIConnectionError,
)
from infection_monkey.network.relay.utils import IslandAPISearchResult, find_available_island_apis
from infection_monkey.network.relay.utils import find_available_island_apis
SERVER_1 = SocketAddress(ip="1.1.1.1", port=12312)
SERVER_2 = SocketAddress(ip="2.2.2.2", port=4321)
@ -48,41 +46,11 @@ def test_find_available_island_apis(
assert len(available_apis) == len(server_response_pairs)
for result in available_apis:
if result.server in expected_available_servers:
assert result.client is not None
for server, island_api_client in available_apis.items():
if server in expected_available_servers:
assert island_api_client is not None
else:
assert result.client is None
def test_find_available_island_apis__preserves_input_order(island_api_client_factory):
available_servers = [SERVER_2, SERVER_3]
with requests_mock.Mocker() as mock:
mock.get(f"https://{SERVER_1}/api?action=is-up", exc=IslandAPIConnectionError)
for server in available_servers:
mock.get(f"https://{server}/api?action=is-up", text="")
available_apis = find_available_island_apis(servers, island_api_client_factory)
for index in range(len(servers)):
assert available_apis[index].server == servers[index]
def _is_none(value) -> bool:
return value is None
def _is_island_client(value) -> bool:
return isinstance(value, IIslandAPIClient)
def _assert_server_and_predicate(
result: IslandAPISearchResult,
server: SocketAddress,
predicate: Callable[[Optional[IIslandAPIClient]], bool],
):
assert result.server == server
assert predicate(result.client)
assert island_api_client is None
def test_find_available_island_apis__multiple_successes(island_api_client_factory):
@ -94,7 +62,7 @@ def test_find_available_island_apis__multiple_successes(island_api_client_factor
available_apis = find_available_island_apis(servers, island_api_client_factory)
_assert_server_and_predicate(available_apis[0], SERVER_1, _is_none)
_assert_server_and_predicate(available_apis[1], SERVER_2, _is_island_client)
_assert_server_and_predicate(available_apis[2], SERVER_3, _is_island_client)
_assert_server_and_predicate(available_apis[3], SERVER_4, _is_none)
assert available_apis[SERVER_1] is None
assert available_apis[SERVER_4] is None
for server in available_servers:
assert isinstance(available_apis[server], IIslandAPIClient)