From 65f5dbeaaf9d6528e91406dc41ea0b1879fe83e9 Mon Sep 17 00:00:00 2001
From: Itay Mizeretz <itay.mizeretz@guardicore.com>
Date: Tue, 3 Oct 2017 15:47:50 +0300
Subject: [PATCH] Sleep only *between* life cycles

---
 chaos_monkey/config.py              |  2 ++
 chaos_monkey/example.conf           |  1 +
 chaos_monkey/monkey.py              | 36 ++++++++++++++++-------------
 monkey_island/cc/services/config.py |  6 +++++
 4 files changed, 29 insertions(+), 16 deletions(-)

diff --git a/chaos_monkey/config.py b/chaos_monkey/config.py
index 6b49d3bb3..c1cd618ef 100644
--- a/chaos_monkey/config.py
+++ b/chaos_monkey/config.py
@@ -173,6 +173,8 @@ class Configuration(object):
     # addresses of internet servers to ping and check if the monkey has internet acccess.
     internet_services = ["monkey.guardicore.com", "www.google.com"]
 
+    keep_tunnel_open_time = 60
+
     ###########################
     # scanners config
     ###########################
diff --git a/chaos_monkey/example.conf b/chaos_monkey/example.conf
index 285bffd11..b738cff75 100644
--- a/chaos_monkey/example.conf
+++ b/chaos_monkey/example.conf
@@ -6,6 +6,7 @@
      "monkey.guardicore.com",
      "www.google.com"
      ],
+    "keep_tunnel_open_time": 60,
     "range_class": "RelativeRange",
     "range_fixed": [
         ""
diff --git a/chaos_monkey/monkey.py b/chaos_monkey/monkey.py
index daabad0ee..426d121eb 100644
--- a/chaos_monkey/monkey.py
+++ b/chaos_monkey/monkey.py
@@ -1,17 +1,18 @@
-import sys
-import os
-import time
-import logging
-import tunnel
 import argparse
+import logging
+import os
 import subprocess
-from system_singleton import SystemSingleton
-from network.firewall import app as firewall
-from control import ControlClient
+import sys
+import time
+
+import tunnel
 from config import WormConfiguration
-from network.network_scanner import NetworkScanner
+from control import ControlClient
 from model import DELAY_DELETE_CMD
+from network.firewall import app as firewall
+from network.network_scanner import NetworkScanner
 from system_info import SystemInfoCollector
+from system_singleton import SystemSingleton
 
 __author__ = 'itamar'
 
@@ -101,7 +102,7 @@ class ChaosMonkey(object):
         else:
             LOG.debug("Running with depth: %d" % WormConfiguration.depth)
 
-        for _ in xrange(WormConfiguration.max_iterations):
+        for iteration_index in xrange(WormConfiguration.max_iterations):
             ControlClient.keepalive()
             ControlClient.load_control_config()
 
@@ -146,7 +147,6 @@ class ChaosMonkey(object):
                         LOG.debug("Skipping %r - exploitation failed before", machine)
                         continue
 
-
                 if monkey_tunnel:
                     monkey_tunnel.set_tunnel_for_host(machine)
                 if self._default_server:
@@ -196,8 +196,10 @@ class ChaosMonkey(object):
                 else:
                     self._fail_exploitation_machines.add(machine)
 
-            if not is_empty:
-                time.sleep(WormConfiguration.timeout_between_iterations)
+            if (not is_empty) and (WormConfiguration.max_iterations > iteration_index + 1):
+                time_to_sleep = WormConfiguration.timeout_between_iterations
+                LOG.info("Sleeping %d seconds before next life cycle iteration", time_to_sleep)
+                time.sleep(time_to_sleep)
 
         if self._keep_running and WormConfiguration.alive:
             LOG.info("Reached max iterations (%d)", WormConfiguration.max_iterations)
@@ -206,8 +208,10 @@ class ChaosMonkey(object):
 
         # if host was exploited, before continue to closing the tunnel ensure the exploited host had its chance to
         # connect to the tunnel
-        if last_exploit_time and (time.time() - last_exploit_time < 60):
-            time.sleep(time.time() - last_exploit_time)
+        if last_exploit_time and (time.time() - last_exploit_time < WormConfiguration.keep_tunnel_open_time):
+            time_to_sleep = WormConfiguration.keep_tunnel_open_time - (time.time() - last_exploit_time)
+            LOG.info("Sleeping %d seconds for exploited machines to connect to tunnel", time_to_sleep)
+            time.sleep(time_to_sleep)
 
         if monkey_tunnel:
             monkey_tunnel.stop()
@@ -242,7 +246,7 @@ class ChaosMonkey(object):
                                      close_fds=True, startupinfo=startupinfo)
                 else:
                     os.remove(sys.executable)
-            except Exception, exc:
+            except Exception as exc:
                 LOG.error("Exception in self delete: %s", exc)
 
         LOG.info("Monkey is shutting down")
diff --git a/monkey_island/cc/services/config.py b/monkey_island/cc/services/config.py
index 5e4d5abe0..6807d5d86 100644
--- a/monkey_island/cc/services/config.py
+++ b/monkey_island/cc/services/config.py
@@ -277,6 +277,12 @@ SCHEMA = {
                             "type": "string",
                             "default": "{2384ec59-0df8-4ab9-918c-843740924a28}",
                             "description": "The name of the mutex used to determine whether the monkey is already running"
+                        },
+                        "keep_tunnel_open_time": {
+                            "title": "Keep tunnel open time",
+                            "type": "integer",
+                            "default": 60,
+                            "description": "Time to keep tunnel open before going down since last exploit (in seconds)"
                         }
                     }
                 },