Storing final conversion state into rule object

SigmaHQ · Mar 12, 2024 · 2b3a1a6 · 2b3a1a6
1 parent e9fdde0
commit 2b3a1a6
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 8 deletions.
diff --git a/sigma/conversion/base.py b/sigma/conversion/base.py
@@ -211,6 +211,7 @@ def convert_rule(self, rule: SigmaRule, output_format: Optional[str] = None) ->
                 for index, query in enumerate(queries)
             ]
             rule.set_conversion_result(finalized_queries)
+            rule.set_conversion_states(states)
             if rule._output:
                 return finalized_queries
             else:

diff --git a/sigma/rule.py b/sigma/rule.py
@@ -680,6 +680,9 @@ class SigmaRuleBase:
     _conversion_result: Optional[List[Any]] = field(
         init=False, default=None, repr=False, compare=False
     )
+    _conversion_states: Optional[List["sigma.conversion.state.ConversionState"]] = field(
+        init=False, default=None, repr=False, compare=False
+    )
     _output: bool = field(init=False, default=True, repr=False, compare=False)
 
     def __post_init__(self):
@@ -980,6 +983,19 @@ def get_conversion_result(self) -> List[Any]:
             )
         return self._conversion_result
 
+    def set_conversion_states(self, state: List["sigma.conversion.state.ConversionState"]):
+        """Set conversion state."""
+        self._conversion_states = state
+
+    def get_conversion_state(self) -> List["sigma.conversion.state.ConversionState"]:
+        """Get conversion state."""
+        if self._conversion_states is None:
+            raise sigma_exceptions.SigmaConversionError(
+                self,
+                "Conversion state not available",
+            )
+        return self._conversion_states
+
     def disable_output(self):
         """Disable output of rule."""
         self._output = False

diff --git a/tests/test_conversion_base.py b/tests/test_conversion_base.py
@@ -2,6 +2,7 @@
 from sigma.backends.test import TextQueryTestBackend
 from sigma.collection import SigmaCollection
 from sigma.conversion.base import TextQueryBackend
+from sigma.conversion.state import ConversionState
 from sigma.processing.conditions import IncludeFieldCondition
 from sigma.processing.finalization import ConcatenateQueriesFinalizer
 from sigma.processing.pipeline import ProcessingPipeline, ProcessingItem, QueryPostprocessingItem
@@ -2218,10 +2219,8 @@ def test_convert_list_cidr_wildcard_asterisk(test_backend, monkeypatch):
 
 
 def test_convert_state(test_backend):
-    assert (
-        test_backend.convert(
-            SigmaCollection.from_yaml(
-                """
+    rules = SigmaCollection.from_yaml(
+        """
             title: Test
             status: test
             logsource:
@@ -2232,12 +2231,22 @@ def test_convert_state(test_backend):
                     fieldA: value
                 condition: sel
         """
-            ),
-            "state",
-        )
-        == ['index=test (mappedA="value")']
     )
 
+    assert test_backend.convert(
+        rules,
+        "state",
+    ) == ['index=test (mappedA="value")']
+    assert rules[0].get_conversion_state() == [
+        ConversionState(
+            processing_state={
+                "index": "test",
+                "data_source": "state_source",
+                "output": "state_output",
+            }
+        )
+    ]
+
 
 def test_convert_query_expression(monkeypatch, test_backend: TextQueryTestBackend):
     monkeypatch.setattr(