From 065a560a83ffde7f7af9b5b58547363fbaeaa7be Mon Sep 17 00:00:00 2001 From: TheTechromancer Date: Thu, 20 Jun 2024 10:05:28 -0400 Subject: [PATCH 1/4] restore secrets.yml functionality --- bbot/core/config/files.py | 19 +------- bbot/core/helpers/misc.py | 69 +++++++++++++++++++++++++++ bbot/core/modules.py | 56 +++++++++++++++++++++- bbot/scanner/preset/preset.py | 1 + bbot/test/test_step_1/test_helpers.py | 67 ++++++++++++++++++++++++++ 5 files changed, 193 insertions(+), 19 deletions(-) diff --git a/bbot/core/config/files.py b/bbot/core/config/files.py index 0f05c0b50..d9cc7644d 100644 --- a/bbot/core/config/files.py +++ b/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,7 +33,6 @@ 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") def get_default_config(self): diff --git a/bbot/core/helpers/misc.py b/bbot/core/helpers/misc.py index 9dedd0d28..129422ae0 100644 --- a/bbot/core/helpers/misc.py +++ b/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 diff --git a/bbot/core/modules.py b/bbot/core/modules.py index c5eb8f902..084c5c76f 100644 --- a/bbot/core/modules.py +++ b/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() diff --git a/bbot/scanner/preset/preset.py b/bbot/scanner/preset/preset.py index 599ae8591..83e92cafc 100644 --- a/bbot/scanner/preset/preset.py +++ b/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 diff --git a/bbot/test/test_step_1/test_helpers.py b/bbot/test/test_step_1/test_helpers.py index 65af36c45..53e2c4222 100644 --- a/bbot/test/test_step_1/test_helpers.py +++ b/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"]] From 26992f88165881fd2e02e4f76ae00a46646e7021 Mon Sep 17 00:00:00 2001 From: TheTechromancer Date: Thu, 20 Jun 2024 12:41:13 -0400 Subject: [PATCH 2/4] modules oopsie safeguard, include secrets.yml in config --- bbot/core/config/files.py | 6 +++++- bbot/core/core.py | 6 +++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/bbot/core/config/files.py b/bbot/core/config/files.py index d9cc7644d..21745f932 100644 --- a/bbot/core/config/files.py +++ b/bbot/core/config/files.py @@ -33,7 +33,11 @@ def _get_config(self, filename, name="config"): return OmegaConf.create() def get_custom_config(self): - return self._get_config(self.config_filename, name="config") + return OmegaConf.merge( + default_config, + 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") diff --git a/bbot/core/core.py b/bbot/core/core.py index 0a789c0ac..75684eee3 100644 --- a/bbot/core/core.py +++ b/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 not OmegaConf.is_dict(modules_entry): + value["modules"] = {} self._custom_config = value def merge_custom(self, config): From e00724bb4b33b1a42ac9959baebae08ce11f906d Mon Sep 17 00:00:00 2001 From: TheTechromancer Date: Thu, 20 Jun 2024 12:42:09 -0400 Subject: [PATCH 3/4] fix undefined variable --- bbot/core/config/files.py | 1 - 1 file changed, 1 deletion(-) diff --git a/bbot/core/config/files.py b/bbot/core/config/files.py index 21745f932..c66e92116 100644 --- a/bbot/core/config/files.py +++ b/bbot/core/config/files.py @@ -34,7 +34,6 @@ def _get_config(self, filename, name="config"): def get_custom_config(self): return OmegaConf.merge( - default_config, self._get_config(self.config_filename, name="config"), self._get_config(self.secrets_filename, name="secrets"), ) From 4efc7b05967e4c2bb4964630e1660ca30db94fbb Mon Sep 17 00:00:00 2001 From: TheTechromancer Date: Thu, 20 Jun 2024 18:01:29 -0400 Subject: [PATCH 4/4] fix tests --- bbot/core/core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bbot/core/core.py b/bbot/core/core.py index 75684eee3..16b02a6bf 100644 --- a/bbot/core/core.py +++ b/bbot/core/core.py @@ -111,7 +111,7 @@ def custom_config(self, value): self._config = None # ensure the modules key is always a dictionary modules_entry = value.get("modules", None) - if not OmegaConf.is_dict(modules_entry): + if modules_entry is not None and not OmegaConf.is_dict(modules_entry): value["modules"] = {} self._custom_config = value