Merge pull request #1339 from blacklanternsecurity/preset-cache

Preset cache

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):

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(
             """