Merge pull request #2241 from guardicore/2216-modify-agent-command-line-arguments

Agent: Modify command line arguments to accept list of servers

CHANGELOG.md

@@ -59,6 +59,8 @@ Changelog](https://keepachangelog.com/en/1.0.0/).
 - `/api/registration` endpoint to `/api/register`. #2105
 - `/api/file-upload` endpoit to `/api/pba/upload`. #2154
 - Improved the speed of ransomware encryption by 2-3x. #2123
+- "-s/--server" to "-s/--servers". #2216
+- "-s/--servers" accepts list of servers separated by comma. #2216
 
 ### Removed
 - VSFTPD exploiter. #1533
@@ -103,6 +105,7 @@ Changelog](https://keepachangelog.com/en/1.0.0/).
 - "/api/configuration/import" endpoint. #2002
 - "/api/configuration/export" endpoint. #2002
 - "/api/island-configuration" endpoint. #2003
+- "-t/--tunnel" from agent command line arguments. #2216
 
 ### Fixed
 - A bug in network map page that caused delay of telemetry log loading. #1545

monkey/infection_monkey/dropper.py

@@ -44,8 +44,7 @@ class MonkeyDrops(object):
     def __init__(self, args):
         arg_parser = argparse.ArgumentParser()
         arg_parser.add_argument("-p", "--parent")
-        arg_parser.add_argument("-t", "--tunnel")
-        arg_parser.add_argument("-s", "--server")
+        arg_parser.add_argument("-s", "--servers", type=lambda arg: arg.strip().split(","))
         arg_parser.add_argument("-d", "--depth", type=positive_int, default=0)
         arg_parser.add_argument("-l", "--location")
         arg_parser.add_argument("-vp", "--vulnerable-port")
@@ -132,8 +131,7 @@ class MonkeyDrops(object):
 
         monkey_options = build_monkey_commandline_explicitly(
             parent=self.opts.parent,
-            tunnel=self.opts.tunnel,
-            server=self.opts.server,
+            servers=self.opts.servers,
             depth=self.opts.depth,
             location=None,
         )

monkey/infection_monkey/monkey.py

@@ -95,8 +95,8 @@ class InfectionMonkey:
         logger.info("Monkey is initializing...")
         self._singleton = SystemSingleton()
         self._opts = self._get_arguments(args)
-        self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(self._opts.server)
-        self._control_client = ControlClient(self._opts.server)
+        self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(self._opts.servers)
+        self._control_client = ControlClient(self._opts.servers)
         # TODO Refactor the telemetry messengers to accept control client
         # and remove control_client_object
         ControlClient.control_client_object = self._control_client
@@ -110,8 +110,7 @@ class InfectionMonkey:
     def _get_arguments(args):
         arg_parser = argparse.ArgumentParser()
         arg_parser.add_argument("-p", "--parent")
-        arg_parser.add_argument("-t", "--tunnel")
-        arg_parser.add_argument("-s", "--server")
+        arg_parser.add_argument("-s", "--servers", type=lambda arg: arg.strip().split(","))
         arg_parser.add_argument("-d", "--depth", type=positive_int, default=0)
         opts = arg_parser.parse_args(args)
         InfectionMonkey._log_arguments(opts)
@@ -154,12 +153,14 @@ class InfectionMonkey:
         if self._current_server_is_set():
             logger.debug(f"Default server set to: {self._control_client.server_address}")
         else:
-            raise Exception(f"Monkey couldn't find server with {self._opts.tunnel} default tunnel.")
+            raise Exception(
+                f"Failed to connect to the island via any known server address: {self._opts.servers}"
+            )
 
         self._control_client.wakeup(parent=self._opts.parent)
 
     def _current_server_is_set(self) -> bool:
-        if self._control_client.find_server(default_tunnel=self._opts.tunnel):
+        if self._control_client.find_server(default_tunnel=self._opts.servers):
             return True
 
         return False