feat: Add custom_param to SigmaValidator class

SigmaHQ · Jul 21, 2024 · 3030276 · 3030276
1 parent 09fe2ab
commit 3030276
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 5 deletions.
diff --git a/sigma/validation.py b/sigma/validation.py
@@ -1,5 +1,5 @@
 from collections import defaultdict
-from typing import Callable, DefaultDict, Dict, Iterable, Iterator, List, Set, Type
+from typing import Callable, DefaultDict, Dict, Iterable, Iterator, Optional, List, Set, Type
 from uuid import UUID
 from sigma.collection import SigmaCollection
 from sigma.exceptions import SigmaConfigurationError
@@ -20,17 +20,22 @@ class SigmaValidator:
 
     validators: Set[SigmaRuleValidator]
     exclusions: DefaultDict[UUID, Set[Type[SigmaRuleValidator]]]
+    custom_param: Optional[Dict]
 
     def __init__(
         self,
         validators: Iterable[Type[SigmaRuleValidator]],
         exclusions: Dict[UUID, Set[SigmaRuleValidator]] = dict(),
+        custom_param: Optional[Dict] = None,
     ):
         self.validators = {validator() for validator in validators}
         self.exclusions = defaultdict(set, exclusions)
+        self.custom_param = custom_param
 
     @classmethod
-    def from_dict(cls, d: Dict, validators: Dict[str, SigmaRuleValidator]) -> "SigmaValidator":
+    def from_dict(
+        cls, d: Dict, validators: Dict[str, SigmaRuleValidator], custom_param: Optional[Dict] = None
+    ) -> "SigmaValidator":
         """
         Instantiate SigmaValidator from dict definition. The dict should have the following
         elements:
@@ -44,6 +49,8 @@ def from_dict(cls, d: Dict, validators: Dict[str, SigmaRuleValidator]) -> "Sigma
         :type d: Dict
         :param validators: Mapping from string identifiers to validator classes.
         :type validators: Dict[str, SigmaRuleValidator]
+        :param custom_param: Optinal custom parameter
+        :type custom_param: Dict
         :return: Instantiated SigmaValidator
         :rtype: SigmaValidator
         """
@@ -84,13 +91,16 @@ def from_dict(cls, d: Dict, validators: Dict[str, SigmaRuleValidator]) -> "Sigma
         except KeyError as e:
             raise SigmaConfigurationError(f"Unknown validator '{ e.args[0] }'")
 
-        return cls(validator_classes, exclusions)
+        return cls(validator_classes, exclusions, custom_param)
 
     @classmethod
     def from_yaml(
-        cls, validator_config: str, validators: Dict[str, SigmaRuleValidator]
+        cls,
+        validator_config: str,
+        validators: Dict[str, SigmaRuleValidator],
+        custom_param: Optional[Dict] = None,
     ) -> "SigmaValidator":
-        return cls.from_dict(yaml.safe_load(validator_config), validators)
+        return cls.from_dict(yaml.safe_load(validator_config), validators, custom_param)
 
     def validate_rule(self, rule: SigmaRule) -> List[SigmaValidationIssue]:
         """

diff --git a/tests/test_validation.py b/tests/test_validation.py
@@ -74,6 +74,7 @@ def test_sigmavalidator_from_dict(validators):
             },
         },
         validators,
+        {"filename_min": 10, "filename_max": 90},
     )
     assert DanglingDetectionValidator in (v.__class__ for v in validator.validators)
     assert TLPv1TagValidator not in (v.__class__ for v in validator.validators)
@@ -85,6 +86,7 @@ def test_sigmavalidator_from_dict(validators):
             NumberAsStringValidator,
         },
     }
+    assert validator.custom_param["filename_min"] == 10
 
 
 def test_sigmavalidator_from_yaml(validators):