diff --git a/bbot/scanner/preset/preset.py b/bbot/scanner/preset/preset.py index 83e92cafc..f28e2e78e 100644 --- a/bbot/scanner/preset/preset.py +++ b/bbot/scanner/preset/preset.py @@ -17,6 +17,9 @@ log = logging.getLogger("bbot.presets") +_preset_cache = dict() + + # cache default presets to prevent having to reload from disk DEFAULT_PRESETS = None @@ -620,7 +623,11 @@ def from_dict(cls, preset_dict, name=None, _exclude=None, _log=False): def include_preset(self, filename): """ - Load a preset from a yaml file and merge it into this one + Load a preset from a yaml file and merge it into this one. + + If the full path is not specified, BBOT will look in all the usual places for it. + + The file extension is optional. Args: filename (Path): The preset YAML file to merge @@ -635,26 +642,30 @@ def include_preset(self, filename): @classmethod def from_yaml_file(cls, filename, _exclude=None, _log=False): """ - Create a preset from a YAML file. If the full path is not specified, BBOT will look in all the usual places for it. - - The file extension is optional. + Load a preset from a YAML file. """ filename = Path(filename).resolve() - if _exclude is None: - _exclude = set() - if _exclude is not None and filename in _exclude: - log.debug(f"Not loading {filename} because it was already loaded {_exclude}") - return False - log.debug(f"Loading {filename} because it's not in excluded list ({_exclude})") - _exclude = set(_exclude) - _exclude.add(filename) try: - yaml_str = open(filename).read() - except FileNotFoundError: - raise PresetNotFoundError(f'Could not find preset at "{filename}" - file does not exist') - preset = cls.from_dict(omegaconf.OmegaConf.create(yaml_str), name=filename.stem, _exclude=_exclude, _log=_log) - preset._yaml_str = yaml_str - return preset + return _preset_cache[filename] + except KeyError: + if _exclude is None: + _exclude = set() + if _exclude is not None and filename in _exclude: + log.debug(f"Not loading {filename} because it was already loaded {_exclude}") + return False + log.debug(f"Loading {filename} because it's not in excluded list ({_exclude})") + _exclude = set(_exclude) + _exclude.add(filename) + try: + yaml_str = open(filename).read() + except FileNotFoundError: + raise PresetNotFoundError(f'Could not find preset at "{filename}" - file does not exist') + preset = cls.from_dict( + omegaconf.OmegaConf.create(yaml_str), name=filename.stem, _exclude=_exclude, _log=_log + ) + preset._yaml_str = yaml_str + _preset_cache[filename] = preset + return preset @classmethod def from_yaml_string(cls, yaml_preset): diff --git a/bbot/test/test_step_1/test_presets.py b/bbot/test/test_step_1/test_presets.py index afa0e3ded..ee6387175 100644 --- a/bbot/test/test_step_1/test_presets.py +++ b/bbot/test/test_step_1/test_presets.py @@ -143,6 +143,30 @@ def test_preset_yaml(clean_default_config): assert yaml_string_2 == yaml_string_1 +def test_preset_cache(): + preset_file = bbot_test_dir / "test_preset.yml" + yaml_string = """ +flags: + - subdomain-enum + +exclude_flags: + - aggressive + - slow +""" + with open(preset_file, "w") as f: + f.write(yaml_string) + + preset = Preset.from_yaml_file(preset_file) + assert "subdomain-enum" in preset.flags + assert "aggressive" in preset.exclude_flags + assert "slow" in preset.exclude_flags + from bbot.scanner.preset.preset import _preset_cache + + assert preset_file in _preset_cache + + preset_file.unlink() + + def test_preset_scope(): blank_preset = Preset() @@ -605,16 +629,16 @@ def test_preset_conditions(): mkdir(custom_preset_dir_1) mkdir(custom_preset_dir_2) - preset_file_1 = custom_preset_dir_1 / "preset1.yml" + preset_file_1 = custom_preset_dir_1 / "preset_condition_1.yml" with open(preset_file_1, "w") as f: f.write( """ include: - - preset2 + - preset_condition_2 """ ) - preset_file_2 = custom_preset_dir_2 / "preset2.yml" + preset_file_2 = custom_preset_dir_2 / "preset_condition_2.yml" with open(preset_file_2, "w") as f: f.write( """