Add yaml tags !include, !set_anchor and !load_anchor (#804)

* Add yaml tags !include, !set_anchor and !load_anchor
* Support yaml anchors for dicts, refactor and add more tests
* Add custom yaml tag documentation

Author: ppcad

CHANGELOG.md

@@ -5,6 +5,8 @@
 ### Features
 
 * add `replacer` processor to replace substrings in fields using a syntax similar to the `dissector`
+* add custom yaml tag `!include PATH_TO_YAML_FILE` that allows to include other yaml files. 
+* add custom yaml tags `!set_anchor ANCHOR_NAME` and `!load_anchor ANCHOR_NAME` that allow to use anchors across documents inside a file/stream.
 
 ### Improvements
 

doc/source/configuration/index.rst

@@ -20,3 +20,4 @@ Configuration
    rules
    getter
    metrics
+   yaml_tags

doc/source/configuration/yaml_tags.rst

@@ -0,0 +1,9 @@
+.. _yaml_tags:
+
+YAML Tags
+=========
+
+
+.. automodule:: logprep.util.tag_yaml_loader
+    :members:
+    :noindex:

logprep/run_logprep.py

@@ -20,12 +20,14 @@
 from logprep.util.logging import LogprepMPQueueListener, logqueue
 from logprep.util.pseudo.commands import depseudonymize, generate_keys, pseudonymize
 from logprep.util.rule_dry_runner import DryRunner
+from logprep.util.tag_yaml_loader import init_yaml_loader_tags
 
 warnings.simplefilter("always", DeprecationWarning)
 logging.captureWarnings(True)
 logging.config.dictConfig(DEFAULT_LOG_CONFIG)
 logger = logging.getLogger("logprep")
 EPILOG_STR = "Check out our docs at https://logprep.readthedocs.io/en/latest/"
+init_yaml_loader_tags("safe", "rt")
 
 
 def _print_version(config: "Configuration") -> None:

logprep/util/tag_yaml_loader.py

@@ -0,0 +1,231 @@
+"""
+Logprep supports the custom YAML tags :code:`!include`, :code:`!set_anchor` and
+:code:`!load_anchor`.
+Those can be used inside any YAML file that is loaded by Logprep.
+
+Include tags
+^^^^^^^^^^^^
+
+The tag :code:`!include PATH_TO_YAML_FILE` loads a single YAML document from a local file path and
+inserts it in its place.
+
+Included files can't contain an :code:`!include` tag.
+
+Example:
+
+..  code-block:: yaml
+    :linenos:
+    :caption: Include tag
+
+    filter: to_resolve
+    generic_resolver:
+        field_mapping:
+            to_resolve: resolved
+        resolve_list: !include /path/to/resolve_list.yml
+
+Anchor tags
+^^^^^^^^^^^
+
+Anchor tags work similar to regular YAML anchors, but are valid for all documents inside a file or
+stream.
+Tags are set with :code:`!set_anchor(:[0-9])?` and loaded with :code:`!load_anchor(:[0-9])?`.
+Ten anchors can be active inside a single file or stream.
+`!set_anchor` and :code:`!load_anchor` are shorthands for :code:`!set_anchor:0` and
+:code:`!load_anchor:0`.
+
+`!include` and :code:`!set_anchor` can't be nested inside :code:`!set_anchor`.
+
+Examples:
+
+..  code-block:: yaml
+    :linenos:
+    :caption: Anchor tag without shorthand
+
+    filter: to_resolve
+    generic_resolver:
+        field_mapping:
+            to_resolve: resolved
+        resolve_list: !set_anchor:1
+            one: foo
+            two: bar
+    ---
+    filter: to_resolve_2
+    generic_resolver:
+        field_mapping:
+            to_resolve_2: resolved
+        resolve_list: !load_anchor:1
+
+..  code-block:: yaml
+    :linenos:
+    :caption: Anchor tag with shorthand
+
+    filter: to_resolve
+    generic_resolver:
+        field_mapping:
+            to_resolve: resolved
+        resolve_list: !set_anchor
+            one: foo
+            two: bar
+    ---
+    filter: to_resolve_2
+    generic_resolver:
+        field_mapping:
+            to_resolve_2: resolved
+        resolve_list: !load_anchor
+"""
+
+import os.path
+from typing import Set, Callable, Any
+
+from ruamel.yaml import YAML, Node, BaseConstructor
+
+
+def init_yaml_loader_tags(*loader_types: str) -> None:
+    """Add custom tags !include, !set_anchor and !load_anchor to the specified loader types.
+
+    Must be initialized before yaml files have been loaded to take effect.
+
+    Parameters
+    ----------
+    *loader_types : str
+        Types of loaders for which tags will be initialized (i.e. "safe" or "rt").
+    """
+
+    def include(_yaml: YAML) -> Callable[[BaseConstructor, Node], Any]:
+        """Includes the contents of a yaml file specified by the !include tag.
+
+        Parameters
+        ----------
+        _yaml : YAML
+            Used to load the yaml file that will be included.
+
+        Returns
+        -------
+        Yaml data where the !include tag has been replaced by the content of the include file.
+        """
+
+        def _include(_: BaseConstructor, node: Node) -> Any:
+            if not isinstance(node.value, (str, os.PathLike)):
+                raise ValueError(f"'{node.value}' is not a file path")
+            if not os.path.isfile(node.value):
+                raise FileNotFoundError(node.value)
+            with open(node.value, "r", encoding="utf-8") as yaml_file:
+                try:
+                    data = _yaml.load(yaml_file)
+                except AttributeError as error:
+                    raise ValueError(f"'{node.tag} {node.value}' could not be loaded") from error
+                if data is None:
+                    raise ValueError(f"'{node.value}' is empty")
+                return data
+
+        return _include
+
+    def set_anchor(
+        _yaml: YAML, _anchors: dict[str, Any], _last_buffer: Set[str]
+    ) -> Callable[[BaseConstructor, Node], Any]:
+        """Sets a global anchor if the '!set_anchor'tag is used, which is valid within a file.
+
+        Setting it for a node with children stores the children inside the anchor.
+        Setting it for a scalar node stores the value of that node inside the anchor.
+
+        Parameters
+        ----------
+        _yaml : YAML
+            Used to load the yaml file that will be included.
+        _anchors : dict[str, Any]
+            The dict where all anchors are stored.
+        _last_buffer : Set[str]
+            Used to check if a different file/stream has been loaded.
+
+        Returns
+        -------
+        The loaded yaml data without any modifications.
+        """
+
+        def _set_anchor(constructor: BaseConstructor, node: Node) -> Any:
+            clear_anchors_if_buffer_changed(constructor, _anchors, _last_buffer)
+
+            anchor_name = get_anchor_name(node)
+            _anchors[anchor_name] = _extract_anchor_value(constructor, node)
+            return _anchors[anchor_name]
+
+        def _extract_anchor_value(constructor: BaseConstructor, node: Node) -> Any:
+            lines = constructor.loader.reader.buffer.splitlines()
+            anchor_value_lines = lines[node.start_mark.line : node.end_mark.line + 1]
+            anchor_value_lines[0] = anchor_value_lines[0][node.start_mark.column + len(node.tag) :]
+            anchor_value_lines[-1] = anchor_value_lines[-1][: node.end_mark.column]
+            anchor_value = "\n".join(anchor_value_lines)
+            try:
+                data = _yaml.load(anchor_value)
+            except AttributeError as error:
+                _, _, value = "\\n".join(lines).partition(node.tag)
+                raise ValueError(f"'{node.tag}{value}' could not be loaded") from error
+            if data is None:
+                raise ValueError(f"'{lines[node.start_mark.line]}' is en empty anchor")
+            return data
+
+        return _set_anchor
+
+    def load_anchor(
+        _anchors: dict[str, Any], _last_buffer: Set[str]
+    ) -> Callable[[BaseConstructor, Node], Any]:
+        """Loads a global anchor if the '!load_anchor'tag is used, which is valid within a file.
+
+        Parameters
+        ----------
+        _anchors : dict[str, Any]
+            The dict where all anchors are stored.
+        _last_buffer : Set[str]
+            Used to check if a different file/stream has been loaded.
+
+        Returns
+        -------
+        Yaml data where the !load_anchor tag has been replaced by the content of the anchor.
+        """
+
+        def _load_anchor(constructor: BaseConstructor, node: Node) -> Any:
+            clear_anchors_if_buffer_changed(constructor, _anchors, _last_buffer)
+
+            anchor_name = get_anchor_name(node)
+            try:
+                return _anchors[anchor_name]
+            except KeyError as error:
+                raise ValueError(
+                    f"'{node.value}' is not a defined anchor within this yaml stream"
+                ) from error
+
+        return _load_anchor
+
+    def clear_anchors_if_buffer_changed(
+        constructor: BaseConstructor, _anchors: dict[str, Any], _last_buffer: Set[str]
+    ) -> None:
+        if constructor.loader.reader.buffer not in _last_buffer:
+            _last_buffer.clear()
+            _anchors.clear()
+            _last_buffer.add(constructor.loader.reader.buffer)
+
+    def get_anchor_name(node: Node) -> str:
+        anchor_name: str
+        _, _, anchor_name = node.tag.partition(":")
+        if anchor_name == "":
+            anchor_name = "0"
+        anchor_name = anchor_name.strip()
+        return anchor_name
+
+    for loader_type in loader_types:
+        yaml = YAML(pure=True, typ=loader_type)
+
+        yaml.constructor.add_constructor("!include", include(yaml))
+
+        last_buffer: Set[str] = set()
+        anchors: dict[str, Any] = {}
+        yaml.constructor.add_constructor("!set_anchor", set_anchor(yaml, anchors, last_buffer))
+        yaml.constructor.add_constructor("!load_anchor", load_anchor(anchors, last_buffer))
+
+        for num in range(10):
+            yaml.constructor.add_constructor(
+                f"!set_anchor:{num}", set_anchor(yaml, anchors, last_buffer)
+            )
+            yaml.constructor.add_constructor(
+                f"!load_anchor:{num}", load_anchor(anchors, last_buffer)
+            )

tests/unit/test_run_logprep.py

@@ -3,6 +3,7 @@
 # pylint: disable=attribute-defined-outside-init
 import logging
 import sys
+import tempfile
 from importlib.metadata import version
 from pathlib import Path
 from unittest import mock
@@ -11,6 +12,7 @@
 import requests
 import responses
 from click.testing import CliRunner
+from ruamel.yaml import YAML
 
 from logprep import run_logprep
 from logprep.run_logprep import cli
@@ -525,3 +527,24 @@ def test_pseudonymize_depseudonymize_with_mode(self, mode, tmp_path):
         )
         assert result.exit_code == 0
         assert result.output.strip() == "string"
+
+
+class TestYamlLoaderTags:
+    def test_yaml_loader_include_tag_initialized(self, tmp_path):
+        yaml = YAML(pure=True, typ="safe")
+        path_to_file_to_include = self._write_to_yaml_file("this: was included", tmp_path)
+        yml_with_tag = f"""
+        foo:
+            bar: !include {path_to_file_to_include}
+        """
+        yaml_file = self._write_to_yaml_file(yml_with_tag, tmp_path)
+        with open(yaml_file, "r", encoding="utf-8") as file:
+            loaded = yaml.load(file)
+        assert loaded["foo"]["bar"] == {"this": "was included"}
+
+    @staticmethod
+    def _write_to_yaml_file(file_content: str, target_directory: Path):
+        rule_file = tempfile.mktemp(dir=target_directory, suffix=".yml")
+        with open(rule_file, "w", encoding="utf-8") as file:
+            file.write(file_content)
+        return rule_file