Merge pull request #1479 from blacklanternsecurity/fix-missing-config

Restore secrets.yml functionality

bbot/core/config/files.py

@@ -2,7 +2,6 @@
 from pathlib import Path
 from omegaconf import OmegaConf
 
-from ..helpers.misc import mkdir
 from ...logger import log_to_stderr
 from ...errors import ConfigLoadError
 
@@ -15,26 +14,11 @@ class BBOTConfigFiles:
     config_dir = (Path.home() / ".config" / "bbot").resolve()
     defaults_filename = (bbot_code_dir / "defaults.yml").resolve()
     config_filename = (config_dir / "bbot.yml").resolve()
+    secrets_filename = (config_dir / "secrets.yml").resolve()
 
     def __init__(self, core):
         self.core = core
 
-    def ensure_config_file(self):
-        mkdir(self.config_dir)
-
-        comment_notice = (
-            "# NOTICE: THESE ENTRIES ARE COMMENTED BY DEFAULT\n"
-            + "# Please be sure to uncomment when inserting API keys, etc.\n"
-        )
-
-        # ensure bbot.yml
-        if not self.config_filename.exists():
-            log_to_stderr(f"Creating BBOT config at {self.config_filename}")
-            yaml = OmegaConf.to_yaml(self.core.default_config)
-            yaml = comment_notice + "\n".join(f"# {line}" for line in yaml.splitlines())
-            with open(str(self.config_filename), "w") as f:
-                f.write(yaml)
-
     def _get_config(self, filename, name="config"):
         filename = Path(filename).resolve()
         try:
@@ -49,8 +33,10 @@ def _get_config(self, filename, name="config"):
             return OmegaConf.create()
 
     def get_custom_config(self):
-        self.ensure_config_file()
-        return self._get_config(self.config_filename, name="config")
+        return OmegaConf.merge(
+            self._get_config(self.config_filename, name="config"),
+            self._get_config(self.secrets_filename, name="secrets"),
+        )
 
     def get_default_config(self):
         return self._get_config(self.defaults_filename, name="defaults")

bbot/core/core.py

@@ -102,13 +102,17 @@ def custom_config(self):
         # we temporarily clear out the config so it can be refreshed if/when custom_config changes
         self._config = None
         if self._custom_config is None:
-            self._custom_config = self.files_config.get_custom_config()
+            self.custom_config = self.files_config.get_custom_config()
         return self._custom_config
 
     @custom_config.setter
     def custom_config(self, value):
         # we temporarily clear out the config so it can be refreshed if/when custom_config changes
         self._config = None
+        # ensure the modules key is always a dictionary
+        modules_entry = value.get("modules", None)
+        if modules_entry is not None and not OmegaConf.is_dict(modules_entry):
+            value["modules"] = {}
         self._custom_config = value
 
     def merge_custom(self, config):

bbot/core/helpers/misc.py

@@ -1,5 +1,6 @@
 import os
 import sys
+import copy
 import json
 import random
 import string
@@ -2733,3 +2734,71 @@ def truncate_filename(file_path, max_length=255):
 
     new_path = directory / (truncated_stem + suffix)
     return new_path
+
+
+def filter_dict(d, *key_names, fuzzy=False, exclude_keys=None, _prev_key=None):
+    """
+    Recursively filter a dictionary based on key names.
+
+    Args:
+        d (dict): The input dictionary.
+        *key_names: Names of keys to filter for.
+        fuzzy (bool): Whether to perform fuzzy matching on keys.
+        exclude_keys (list, None): List of keys to be excluded from the final dict.
+        _prev_key (str, None): For internal recursive use; the previous key in the hierarchy.
+
+    Returns:
+        dict: A dictionary containing only the keys specified in key_names.
+
+    Examples:
+        >>> filter_dict({"key1": "test", "key2": "asdf"}, "key2")
+        {"key2": "asdf"}
+        >>> filter_dict({"key1": "test", "key2": {"key3": "asdf"}}, "key1", "key3", exclude_keys="key2")
+        {'key1': 'test'}
+    """
+    if exclude_keys is None:
+        exclude_keys = []
+    if isinstance(exclude_keys, str):
+        exclude_keys = [exclude_keys]
+    ret = {}
+    if isinstance(d, dict):
+        for key in d:
+            if key in key_names or (fuzzy and any(k in key for k in key_names)):
+                if not any(k in exclude_keys for k in [key, _prev_key]):
+                    ret[key] = copy.deepcopy(d[key])
+            elif isinstance(d[key], list) or isinstance(d[key], dict):
+                child = filter_dict(d[key], *key_names, fuzzy=fuzzy, _prev_key=key, exclude_keys=exclude_keys)
+                if child:
+                    ret[key] = child
+    return ret
+
+
+def clean_dict(d, *key_names, fuzzy=False, exclude_keys=None, _prev_key=None):
+    """
+    Recursively clean unwanted keys from a dictionary.
+    Useful for removing secrets from a config.
+
+    Args:
+        d (dict): The input dictionary.
+        *key_names: Names of keys to remove.
+        fuzzy (bool): Whether to perform fuzzy matching on keys.
+        exclude_keys (list, None): List of keys to be excluded from removal.
+        _prev_key (str, None): For internal recursive use; the previous key in the hierarchy.
+
+    Returns:
+        dict: A dictionary cleaned of the keys specified in key_names.
+
+    """
+    if exclude_keys is None:
+        exclude_keys = []
+    if isinstance(exclude_keys, str):
+        exclude_keys = [exclude_keys]
+    d = copy.deepcopy(d)
+    if isinstance(d, dict):
+        for key, val in list(d.items()):
+            if key in key_names or (fuzzy and any(k in key for k in key_names)):
+                if _prev_key not in exclude_keys:
+                    d.pop(key)
+            else:
+                d[key] = clean_dict(val, *key_names, fuzzy=fuzzy, _prev_key=key, exclude_keys=exclude_keys)
+    return d

bbot/core/modules.py

@@ -18,7 +18,17 @@
 
 from .flags import flag_descriptions
 from .shared_deps import SHARED_DEPS
-from .helpers.misc import list_files, sha1, search_dict_by_key, search_format_dict, make_table, os_platform, mkdir
+from .helpers.misc import (
+    list_files,
+    sha1,
+    search_dict_by_key,
+    filter_dict,
+    clean_dict,
+    search_format_dict,
+    make_table,
+    os_platform,
+    mkdir,
+)
 
 
 log = logging.getLogger("bbot.module_loader")
@@ -680,5 +690,49 @@ def filter_modules(self, modules=None, mod_type=None):
         module_list.sort(key=lambda x: x[-1]["type"], reverse=True)
         return module_list
 
+    def ensure_config_files(self):
+
+        secrets_strings = ["api_key", "username", "password", "token", "secret", "_id"]
+        exclude_keys = ["modules"]
+
+        files = self.core.files_config
+        mkdir(files.config_dir)
+
+        comment_notice = (
+            "# NOTICE: THESE ENTRIES ARE COMMENTED BY DEFAULT\n"
+            + "# Please be sure to uncomment when inserting API keys, etc.\n"
+        )
+
+        config_obj = secrets_only_config = OmegaConf.to_object(self.core.default_config)
+
+        # ensure bbot.yml
+        if not files.config_filename.exists():
+            log_to_stderr(f"Creating BBOT config at {files.config_filename}")
+            no_secrets_config = clean_dict(
+                config_obj,
+                *secrets_strings,
+                fuzzy=True,
+                exclude_keys=exclude_keys,
+            )
+            yaml = OmegaConf.to_yaml(no_secrets_config)
+            yaml = comment_notice + "\n".join(f"# {line}" for line in yaml.splitlines())
+            with open(str(files.config_filename), "w") as f:
+                f.write(yaml)
+
+        # ensure secrets.yml
+        if not files.secrets_filename.exists():
+            log_to_stderr(f"Creating BBOT secrets at {files.secrets_filename}")
+            secrets_only_config = filter_dict(
+                config_obj,
+                *secrets_strings,
+                fuzzy=True,
+                exclude_keys=exclude_keys,
+            )
+            yaml = OmegaConf.to_yaml(secrets_only_config)
+            yaml = comment_notice + "\n".join(f"# {line}" for line in yaml.splitlines())
+            with open(str(files.secrets_filename), "w") as f:
+                f.write(yaml)
+            files.secrets_filename.chmod(0o600)
+
 
 MODULE_LOADER = ModuleLoader()

bbot/scanner/preset/preset.py

@@ -545,6 +545,7 @@ def module_loader(self):
             from bbot.core.modules import MODULE_LOADER
 
             self._module_loader = MODULE_LOADER
+            self._module_loader.ensure_config_files()
 
         return self._module_loader
 

bbot/test/test_step_1/test_helpers.py

@@ -254,6 +254,73 @@ async def test_helpers_misc(helpers, scan, bbot_scanner, bbot_httpserver):
     assert replaced["asdf"][1][500] == True
     assert replaced["asdf"][0]["wat"]["here"] == "asdf!"
 
+    filtered_dict = helpers.filter_dict(
+        {"modules": {"c99": {"api_key": "1234", "filterme": "asdf"}, "ipneighbor": {"test": "test"}}}, "api_key"
+    )
+    assert "api_key" in filtered_dict["modules"]["c99"]
+    assert "filterme" not in filtered_dict["modules"]["c99"]
+    assert "ipneighbor" not in filtered_dict["modules"]
+
+    filtered_dict2 = helpers.filter_dict(
+        {"modules": {"c99": {"api_key": "1234", "filterme": "asdf"}, "ipneighbor": {"test": "test"}}}, "c99"
+    )
+    assert "api_key" in filtered_dict2["modules"]["c99"]
+    assert "filterme" in filtered_dict2["modules"]["c99"]
+    assert "ipneighbor" not in filtered_dict2["modules"]
+
+    filtered_dict3 = helpers.filter_dict(
+        {"modules": {"c99": {"api_key": "1234", "filterme": "asdf"}, "ipneighbor": {"test": "test"}}},
+        "key",
+        fuzzy=True,
+    )
+    assert "api_key" in filtered_dict3["modules"]["c99"]
+    assert "filterme" not in filtered_dict3["modules"]["c99"]
+    assert "ipneighbor" not in filtered_dict3["modules"]
+
+    filtered_dict4 = helpers.filter_dict(
+        {"modules": {"secrets_db": {"api_key": "1234"}, "ipneighbor": {"secret": "test", "asdf": "1234"}}},
+        "secret",
+        fuzzy=True,
+        exclude_keys="modules",
+    )
+    assert not "secrets_db" in filtered_dict4["modules"]
+    assert "ipneighbor" in filtered_dict4["modules"]
+    assert "secret" in filtered_dict4["modules"]["ipneighbor"]
+    assert "asdf" not in filtered_dict4["modules"]["ipneighbor"]
+
+    cleaned_dict = helpers.clean_dict(
+        {"modules": {"c99": {"api_key": "1234", "filterme": "asdf"}, "ipneighbor": {"test": "test"}}}, "api_key"
+    )
+    assert "api_key" not in cleaned_dict["modules"]["c99"]
+    assert "filterme" in cleaned_dict["modules"]["c99"]
+    assert "ipneighbor" in cleaned_dict["modules"]
+
+    cleaned_dict2 = helpers.clean_dict(
+        {"modules": {"c99": {"api_key": "1234", "filterme": "asdf"}, "ipneighbor": {"test": "test"}}}, "c99"
+    )
+    assert "c99" not in cleaned_dict2["modules"]
+    assert "ipneighbor" in cleaned_dict2["modules"]
+
+    cleaned_dict3 = helpers.clean_dict(
+        {"modules": {"c99": {"api_key": "1234", "filterme": "asdf"}, "ipneighbor": {"test": "test"}}},
+        "key",
+        fuzzy=True,
+    )
+    assert "api_key" not in cleaned_dict3["modules"]["c99"]
+    assert "filterme" in cleaned_dict3["modules"]["c99"]
+    assert "ipneighbor" in cleaned_dict3["modules"]
+
+    cleaned_dict4 = helpers.clean_dict(
+        {"modules": {"secrets_db": {"api_key": "1234"}, "ipneighbor": {"secret": "test", "asdf": "1234"}}},
+        "secret",
+        fuzzy=True,
+        exclude_keys="modules",
+    )
+    assert "secrets_db" in cleaned_dict4["modules"]
+    assert "ipneighbor" in cleaned_dict4["modules"]
+    assert "secret" not in cleaned_dict4["modules"]["ipneighbor"]
+    assert "asdf" in cleaned_dict4["modules"]["ipneighbor"]
+
     assert helpers.split_list([1, 2, 3, 4, 5]) == [[1, 2], [3, 4, 5]]
     assert list(helpers.grouper("ABCDEFG", 3)) == [["A", "B", "C"], ["D", "E", "F"], ["G"]]