From c35652c8c88ca23f7eeabb886aaf5039c0097e4c Mon Sep 17 00:00:00 2001 From: Mika Ayenson Date: Thu, 4 Apr 2024 15:50:48 -0500 Subject: [PATCH] [Bug] Add explicit format preserver (#3566) --- detection_rules/rule_formatter.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/detection_rules/rule_formatter.py b/detection_rules/rule_formatter.py index 5d8516768..c1fa33ac0 100644 --- a/detection_rules/rule_formatter.py +++ b/detection_rules/rule_formatter.py @@ -74,6 +74,31 @@ class NonformattedField(str): """Non-formatting class.""" +def preserve_formatting_for_fields(data: OrderedDict, fields_to_preserve: list) -> OrderedDict: + """Preserve formatting for specified nested fields in an action.""" + + def apply_preservation(target: OrderedDict, keys: list) -> None: + """Apply NonformattedField preservation based on keys path.""" + for key in keys[:-1]: + # Iterate to the key, diving into nested dictionaries + if key in target and isinstance(target[key], dict): + target = target[key] + else: + # Cannot preserve formatting for missing or non-dict intermediate + return + + final_key = keys[-1] + if final_key in target: + # Apply NonformattedField to the target field if it exists + target[final_key] = NonformattedField(target[final_key]) + + for field_path in fields_to_preserve: + keys = field_path.split('.') + apply_preservation(data, keys) + + return data + + class RuleTomlEncoder(toml.TomlEncoder): """Generate a pretty form of toml.""" @@ -186,6 +211,11 @@ def toml_write(rule_contents, outfile=None): for k in sorted(list(_contents)): v = _contents.pop(k) + if k == 'actions': + # explicitly preserve formatting for message field in actions + preserved_fields = ["params.message"] + v = [preserve_formatting_for_fields(action, preserved_fields) for action in v] + if isinstance(v, dict): bottom[k] = OrderedDict(sorted(v.items())) elif isinstance(v, list):