Update code

BC-SECURITY · Nov 25, 2024 · 708ff79 · 708ff79
1 parent 564fd2e
commit 708ff79
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 14 deletions.
diff --git a/empire/config_manager.py b/empire/config_manager.py
@@ -44,22 +44,20 @@ def check_config_permission(config_dict: dict, config_type: str):
     # Define paths to check based on config type
     if config_type == "server":
         paths_to_check = {
-            ("api", "cert_path"): config_dict["api"]["cert_path"],
-            ("database", "sqlite", "location"): config_dict["database"]["sqlite"][
-                "location"
-            ],
-            ("starkiller", "directory"): config_dict["starkiller"]["directory"],
-            ("logging", "directory"): config_dict["logging"]["directory"],
-            ("debug", "last_task", "file"): config_dict["debug"]["last_task"]["file"],
-            ("directories", "downloads"): config_dict["directories"].get("downloads"),
+            ("api", "cert_path"): config_dict.get("api", {}).get("cert_path"),
+            ("database", "sqlite", "location"): config_dict.get("database", {}).get("sqlite", {}).get("location"),
+            ("starkiller", "directory"): config_dict.get("starkiller", {}).get("directory"),
+            ("logging", "directory"): config_dict.get("logging", {}).get("directory"),
+            ("debug", "last_task", "file"): config_dict.get("debug", {}).get("last_task", {}).get("file"),
+            ("directories", "downloads"): config_dict.get("directories", {}).get("downloads"),
         }
         config_path = CONFIG_SERVER_PATH  # Use the server config path
 
     elif config_type == "client":
         paths_to_check = {
-            ("logging", "directory"): config_dict["logging"]["directory"],
-            ("directories", "downloads"): config_dict["directories"].get("downloads"),
-            ("directories", "generated-stagers"): config_dict["directories"].get(
+            ("logging", "directory"): config_dict.get("logging", {}).get("directory"),
+            ("directories", "downloads"): config_dict.get("directories", {}).get("downloads"),
+            ("directories", "generated-stagers"): config_dict.get("directories", {}).get(
                 "generated-stagers"
             ),
         }
@@ -70,6 +68,8 @@ def check_config_permission(config_dict: dict, config_type: str):
 
     # Check permissions and update paths as needed
     for keys, dir_path in paths_to_check.items():
+        if dir_path is None:
+            continue
         if not os.access(dir_path, os.W_OK):
             log.info(
                 "No write permission for %s. Switching to fallback directory.", dir_path

diff --git a/empire/server/core/config.py b/empire/server/core/config.py
@@ -174,9 +174,6 @@ def set_yaml(location: str):
     if loaded_config:
         config_dict = loaded_config
 
-#tmp_config_dict = EmpireConfig(config_dict).model_dump()
-#print(tmp_config_dict)
 # Change this function to return back the updated dict after it writes out the file
 config_dict = config_manager.check_config_permission(config_dict, "server")
-print(config_dict)
 empire_config = EmpireConfig(config_dict)