From 4b356c2cf4d30b97efa7421af18f388d86f1f7a5 Mon Sep 17 00:00:00 2001
From: VakarisZ <vakarisz@yahoo.com>
Date: Wed, 21 Aug 2019 11:54:38 +0300
Subject: [PATCH] Fixes bug where monkey crashes if default server flag is not
 passed

---
 monkey/infection_monkey/monkey.py | 27 ++++++++++++++++++---------
 1 file changed, 18 insertions(+), 9 deletions(-)

diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py
index ce5ab2093..692e278fb 100644
--- a/monkey/infection_monkey/monkey.py
+++ b/monkey/infection_monkey/monkey.py
@@ -67,10 +67,7 @@ class InfectionMonkey(object):
         self._parent = self._opts.parent
         self._default_tunnel = self._opts.tunnel
         self._default_server = self._opts.server
-        try:
-            self._default_server_port = self._default_server.split(':')[1]
-        except KeyError:
-            self._default_server_port = ''
+
         if self._opts.depth:
             WormConfiguration._depth_from_commandline = True
         self._keep_running = True
@@ -87,9 +84,10 @@ class InfectionMonkey(object):
     def start(self):
         LOG.info("Monkey is running...")
 
-        if not ControlClient.find_server(default_tunnel=self._default_tunnel):
-            LOG.info("Monkey couldn't find server. Going down.")
+        # Sets island's IP and port for monkey to communicate to
+        if not self.set_default_server():
             return
+        self.set_default_port()
 
         # Create a dir for monkey files if there isn't one
         utils.create_monkey_dir()
@@ -116,9 +114,6 @@ class InfectionMonkey(object):
             monkey_tunnel.start()
 
         StateTelem(False).send()
-
-        self._default_server = WormConfiguration.current_server
-        LOG.debug("default server: %s" % self._default_server)
         TunnelTelem().send()
 
         if WormConfiguration.collect_system_info:
@@ -329,3 +324,17 @@ class InfectionMonkey(object):
             self._keep_running = False
 
             LOG.info("Max exploited victims reached (%d)", WormConfiguration.victims_max_exploit)
+
+    def set_default_port(self):
+        try:
+            self._default_server_port = self._default_server.split(':')[1]
+        except KeyError:
+            self._default_server_port = ''
+
+    def set_default_server(self):
+        if not ControlClient.find_server(default_tunnel=self._default_tunnel):
+            LOG.info("Monkey couldn't find server. Going down.")
+            return False
+        self._default_server = WormConfiguration.current_server
+        LOG.debug("default server set to: %s" % self._default_server)
+        return True