From 355a75feeff0e52c8b67932bb69b6c8bccff9570 Mon Sep 17 00:00:00 2001
From: Itay Mizeretz <itay.mizeretz@guardicore.com>
Date: Wed, 28 Feb 2018 16:21:44 +0200
Subject: [PATCH] seperate the wakeup and server lookup processes

---
 infection_monkey/control.py | 57 ++++++++++++++++++++++---------------
 infection_monkey/monkey.py  |  3 +-
 2 files changed, 36 insertions(+), 24 deletions(-)

diff --git a/infection_monkey/control.py b/infection_monkey/control.py
index e7fb4cebb..8d14766c6 100644
--- a/infection_monkey/control.py
+++ b/infection_monkey/control.py
@@ -24,10 +24,10 @@ class ControlClient(object):
     proxies = {}
 
     @staticmethod
-    def wakeup(parent=None, default_tunnel=None, has_internet_access=None):
-        LOG.debug("Trying to wake up with Monkey Island servers list: %r" % WormConfiguration.command_servers)
-        if parent or default_tunnel:
-            LOG.debug("parent: %s, default_tunnel: %s" % (parent, default_tunnel))
+    def wakeup(parent=None, has_internet_access=None):
+        if parent:
+            LOG.debug("parent: %s" % (parent,))
+
         hostname = gethostname()
         if not parent:
             parent = GUID
@@ -35,31 +35,43 @@ class ControlClient(object):
         if has_internet_access is None:
             has_internet_access = check_internet_access(WormConfiguration.internet_services)
 
+        monkey = {'guid': GUID,
+                  'hostname': hostname,
+                  'ip_addresses': local_ips(),
+                  'description': " ".join(platform.uname()),
+                  'internet_access': has_internet_access,
+                  'config': WormConfiguration.as_dict(),
+                  'parent': parent}
+
+        if ControlClient.proxies:
+            monkey['tunnel'] = ControlClient.proxies.get('https')
+
+        requests.post("https://%s/api/monkey" % (WormConfiguration.current_server,),
+                      data=json.dumps(monkey),
+                      headers={'content-type': 'application/json'},
+                      verify=False,
+                      proxies=ControlClient.proxies,
+                      timeout=20)
+
+    @staticmethod
+    def find_server(default_tunnel=None):
+        LOG.debug("Trying to wake up with Monkey Island servers list: %r" % WormConfiguration.command_servers)
+        if default_tunnel:
+            LOG.debug("default_tunnel: %s" % (default_tunnel,))
+
         for server in WormConfiguration.command_servers:
             try:
                 WormConfiguration.current_server = server
 
-                monkey = {'guid': GUID,
-                          'hostname': hostname,
-                          'ip_addresses': local_ips(),
-                          'description': " ".join(platform.uname()),
-                          'internet_access': has_internet_access,
-                          'config': WormConfiguration.as_dict(),
-                          'parent': parent}
-
-                if ControlClient.proxies:
-                    monkey['tunnel'] = ControlClient.proxies.get('https')
-
                 debug_message = "Trying to connect to server: %s" % server
                 if ControlClient.proxies:
                     debug_message += " through proxies: %s" % ControlClient.proxies
                 LOG.debug(debug_message)
-                reply = requests.post("https://%s/api/monkey" % (server,),
-                                      data=json.dumps(monkey),
-                                      headers={'content-type': 'application/json'},
-                                      verify=False,
-                                      proxies=ControlClient.proxies,
-                                      timeout=20)
+                # TODO: use different api call to check connectivity.
+                requests.get("https://%s/api/monkey" % (server,),
+                             verify=False,
+                             proxies=ControlClient.proxies)
+
                 break
 
             except Exception as exc:
@@ -74,7 +86,7 @@ class ControlClient(object):
                     proxy_address, proxy_port = proxy_find
                     LOG.info("Found tunnel at %s:%s" % (proxy_address, proxy_port))
                     ControlClient.proxies['https'] = 'https://%s:%s' % (proxy_address, proxy_port)
-                    ControlClient.wakeup(parent=parent, has_internet_access=has_internet_access)
+                    ControlClient.find_server()
                 else:
                     LOG.info("No tunnel found")
 
@@ -234,7 +246,6 @@ class ControlClient(object):
                                   data=json.dumps(host_dict),
                                   headers={'content-type': 'application/json'},
                                   verify=False, proxies=ControlClient.proxies)
-
             if 200 == reply.status_code:
                 result_json = reply.json()
                 filename = result_json.get('filename')
diff --git a/infection_monkey/monkey.py b/infection_monkey/monkey.py
index b05d94fd4..19a456cff 100644
--- a/infection_monkey/monkey.py
+++ b/infection_monkey/monkey.py
@@ -76,7 +76,8 @@ class InfectionMonkey(object):
 
         if firewall.is_enabled():
             firewall.add_firewall_rule()
-        ControlClient.wakeup(parent=self._parent, default_tunnel=self._default_tunnel)
+        ControlClient.find_server(default_tunnel=self._default_tunnel)
+        ControlClient.wakeup(parent=self._parent)
         ControlClient.load_control_config()
 
         if not WormConfiguration.alive: