From d61c080848d96ff8981041e9cd96e4acc58d7975 Mon Sep 17 00:00:00 2001
From: VakarisZ <vakarisz@yahoo.com>
Date: Wed, 20 Mar 2019 14:02:53 +0200
Subject: [PATCH] Small refactor, values to island's config added.

---
 monkey/infection_monkey/monkey.py                |  2 +-
 monkey/infection_monkey/utils.py                 | 16 ++++++++--------
 .../monkey_island/cc/services/config_schema.py   | 14 +++++++++++++-
 3 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py
index fe7558714..f09425f25 100644
--- a/monkey/infection_monkey/monkey.py
+++ b/monkey/infection_monkey/monkey.py
@@ -219,7 +219,6 @@ class InfectionMonkey(object):
             self._singleton.unlock()
 
         InfectionMonkey.self_delete()
-        utils.remove_monkey_dir()
         LOG.info("Monkey is shutting down")
 
     @staticmethod
@@ -234,6 +233,7 @@ class InfectionMonkey(object):
         if WormConfiguration.self_delete_in_cleanup \
                 and -1 == sys.executable.find('python'):
             try:
+                utils.remove_monkey_dir()
                 if "win32" == sys.platform:
                     from _subprocess import SW_HIDE, STARTF_USESHOWWINDOW, CREATE_NEW_CONSOLE
                     startupinfo = subprocess.STARTUPINFO()
diff --git a/monkey/infection_monkey/utils.py b/monkey/infection_monkey/utils.py
index ccb6010c7..741d7c950 100644
--- a/monkey/infection_monkey/utils.py
+++ b/monkey/infection_monkey/utils.py
@@ -42,19 +42,19 @@ def create_monkey_dir():
     """
     Creates directory for monkey and related files
     """
-    if is_windows_os():
-        if not os.path.exists(WormConfiguration.monkey_dir_windows):
-            os.mkdir(WormConfiguration.monkey_dir_windows)
-    else:
-        if not os.path.exists(WormConfiguration.monkey_log_path_linux):
-            os.mkdir(WormConfiguration.monkey_dir_linux)
+    if not os.path.exists(get_monkey_dir_path()):
+        os.mkdir(get_monkey_dir_path())
 
 
 def remove_monkey_dir():
     """
     Removes monkey's root directory
     """
+    shutil.rmtree(get_monkey_dir_path(), ignore_errors=True)
+
+
+def get_monkey_dir_path():
     if is_windows_os():
-        shutil.rmtree(WormConfiguration.monkey_dir_windows, ignore_errors=True)
+        return WormConfiguration.monkey_dir_windows
     else:
-        shutil.rmtree(WormConfiguration.monkey_dir_linux, ignore_errors=True)
+        return WormConfiguration.monkey_dir_linux
diff --git a/monkey/monkey_island/cc/services/config_schema.py b/monkey/monkey_island/cc/services/config_schema.py
index 8d99540bf..cbcc6ba0a 100644
--- a/monkey/monkey_island/cc/services/config_schema.py
+++ b/monkey/monkey_island/cc/services/config_schema.py
@@ -423,7 +423,19 @@ SCHEMA = {
                             "type": "integer",
                             "default": 60,
                             "description": "Time to keep tunnel open before going down after last exploit (in seconds)"
-                        }
+                        },
+                        "monkey_dir_windows": {
+                            "title": "Monkey's windows directory",
+                            "type": "string",
+                            "default": r"C:\Windows\temp\monkey_dir",
+                            "description": "Directory containing all monkey files on windows"
+                        },
+                        "monkey_dir_linux": {
+                            "title": "Monkey's linux directory",
+                            "type": "string",
+                            "default": "/tmp/monkey_dir",
+                            "description": "Directory containing all monkey files on linux"
+                        },
                     }
                 },
                 "classes": {