From ec42e4b3ef2c08a11c5fba3da2d2b84eae64d4a2 Mon Sep 17 00:00:00 2001
From: Alex <61382599+NeffIsBack@users.noreply.github.com>
Date: Tue, 15 Aug 2023 22:52:19 +0200
Subject: [PATCH] Add config options from default config if it is missing in
 the users config file

---
 cme/config.py     | 15 ++++++++++-----
 cme/data/cme.conf |  1 +
 cme/first_run.py  | 16 ----------------
 3 files changed, 11 insertions(+), 21 deletions(-)

diff --git a/cme/config.py b/cme/config.py
index 52dcaa4e..691eab87 100644
--- a/cme/config.py
+++ b/cme/config.py
@@ -17,12 +17,17 @@ if "CME" not in cme_config.sections():
     first_run_setup()
     cme_config.read(os.path.join(CME_PATH, "cme.conf"))
 
-for option in cme_default_config.options("CME"):
-    if option not in cme_config.options("CME"):
-        cme_logger.info("Adding missing option '{}' to cme.conf".format(option))
-        cme_config.set("CME", option, cme_default_config.get("CME", option))
+# Check if there are any missing options in the config file
+for section in cme_default_config.sections():
+    for option in cme_default_config.options(section):
+        if not cme_config.has_option(section, option):
+            cme_logger.display(f"Adding missing option '{option}' in config section '{section}' to cme.conf")
+            cme_config.set(section, option, cme_default_config.get(section, option))
 
-# These options have to exist in the default config file!!
+            with open(path_join(CME_PATH, "cme.conf"), "w") as config_file:
+                cme_config.write(config_file)
+
+#!!! THESE OPTIONS HAVE TO EXIST IN THE DEFAULT CONFIG FILE !!!
 cme_workspace = cme_config.get("CME", "workspace", fallback="default")
 pwned_label = cme_config.get("CME", "pwn3d_label", fallback="Pwn3d!")
 audit_mode = cme_config.get("CME", "audit_mode", fallback=False)
diff --git a/cme/data/cme.conf b/cme/data/cme.conf
index 4f6c21a0..7b75ff83 100755
--- a/cme/data/cme.conf
+++ b/cme/data/cme.conf
@@ -6,6 +6,7 @@ audit_mode = False
 reveal_chars_of_pwd = 0
 log_mode = False
 ignore_opsec = True
+host_info_colors = ["green", "red", "yellow", "cyan"]
 
 [BloodHound]
 bh_enabled = False
diff --git a/cme/first_run.py b/cme/first_run.py
index 3a6267a7..a55f0746 100755
--- a/cme/first_run.py
+++ b/cme/first_run.py
@@ -5,8 +5,6 @@ from os import mkdir
 from os.path import exists
 from os.path import join as path_join
 import shutil
-import configparser
-from configparser import NoSectionError, NoOptionError
 from cme.paths import CME_PATH, CONFIG_PATH, TMP_PATH, DATA_PATH
 from cme.cmedb import initialize_db
 from cme.logger import cme_logger
@@ -40,20 +38,6 @@ def first_run_setup(logger=cme_logger):
         logger.display("Copying default configuration file")
         default_path = path_join(DATA_PATH, "cme.conf")
         shutil.copy(default_path, CME_PATH)
-    else:
-        # This is just a quick check to make sure the config file isn't the old 3.x format
-        try:
-            config = configparser.ConfigParser()
-            config.read(CONFIG_PATH)
-            config.get("CME", "workspace")
-            config.get("CME", "pwn3d_label")
-            config.get("CME", "audit_mode")
-            config.get("BloodHound", "bh_enabled")
-            config.get("CME", "log_mode")
-        except (NoSectionError, NoOptionError):
-            logger.display("Old configuration file detected, replacing with new version")
-            default_path = path_join(DATA_PATH, "cme.conf")
-            shutil.copy(default_path, CME_PATH)
 
     # if not exists(CERT_PATH):
     #     logger.display('Generating SSL certificate')