forked from p15670423/monkey
Agent: Find server and send control relay message to all other servers
This commit is contained in:
parent
65226d5a9c
commit
e539495545
|
@ -42,7 +42,7 @@ from infection_monkey.master.control_channel import ControlChannel
|
|||
from infection_monkey.model import VictimHostFactory
|
||||
from infection_monkey.network.firewall import app as firewall
|
||||
from infection_monkey.network.info import get_network_interfaces
|
||||
from infection_monkey.network.relay.utils import find_server
|
||||
from infection_monkey.network.relay.utils import find_server, send_relay_control_message
|
||||
from infection_monkey.network_scanning.elasticsearch_fingerprinter import ElasticSearchFingerprinter
|
||||
from infection_monkey.network_scanning.http_fingerprinter import HTTPFingerprinter
|
||||
from infection_monkey.network_scanning.mssql_fingerprinter import MSSQLFingerprinter
|
||||
|
@ -98,9 +98,9 @@ class InfectionMonkey:
|
|||
self._opts = self._get_arguments(args)
|
||||
|
||||
# TODO: Revisit variable names
|
||||
server = find_server(self._opts.servers)
|
||||
self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(server)
|
||||
self._control_client = ControlClient(server_address=server)
|
||||
self._get_server()
|
||||
self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(self.server)
|
||||
self._control_client = ControlClient(server_address=self.server)
|
||||
|
||||
# TODO Refactor the telemetry messengers to accept control client
|
||||
# and remove control_client_object
|
||||
|
@ -122,6 +122,11 @@ class InfectionMonkey:
|
|||
|
||||
return opts
|
||||
|
||||
def _get_server(self):
|
||||
servers_iterator = (s for s in self._opts.servers)
|
||||
self.server = find_server(servers_iterator)
|
||||
send_relay_control_message(servers_iterator)
|
||||
|
||||
@staticmethod
|
||||
def _log_arguments(args):
|
||||
arg_string = " ".join([f"{key}: {value}" for key, value in vars(args).items()])
|
||||
|
@ -159,14 +164,13 @@ class InfectionMonkey:
|
|||
logger.debug(f"Default server set to: {self._control_client.server_address}")
|
||||
else:
|
||||
raise Exception(
|
||||
f"Failed to connect to the island via "
|
||||
f"any known server address: {self._opts.servers}"
|
||||
f"Failed to connect to the island via " f"any known server address: {self.server}"
|
||||
)
|
||||
|
||||
self._control_client.wakeup(parent=self._opts.parent)
|
||||
|
||||
def _current_server_is_set(self) -> bool:
|
||||
if find_server(servers=self._opts.servers):
|
||||
if self.server:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
|
Loading…
Reference in New Issue