Merge pull request #372 from SigmaHQ/fix-mypy-issues

Fixed typing issues

Author: thomaspatzke

pyproject.toml

@@ -7,7 +7,6 @@ authors = [
   { name = "Thomas Patzke", email = "thomas@patzke.org" },
 ]
 readme = "README.md"
-repository = "https://github.com/SigmaHQ/pySigma"
 classifiers = [
     "Development Status :: 4 - Beta",
     "Intended Audience :: Developers",

sigma/collection.py

@@ -1,7 +1,7 @@
 from dataclasses import InitVar, dataclass, field
 from functools import reduce
 from pathlib import Path
-from typing import Any, Callable, Dict, Iterable, List, Optional, Union, IO, TYPE_CHECKING
+from typing import Any, Callable, Dict, Iterable, List, Optional, Union, IO, TYPE_CHECKING, cast
 from uuid import UUID
 
 import yaml
@@ -88,7 +88,9 @@ def resolve_rule_references(self) -> None:
                 rule.resolve_rule_references(self)
 
         # Extract all filters from the rules
-        filters: List[SigmaFilter] = [rule for rule in self.rules if isinstance(rule, SigmaFilter)]
+        filters: List[SigmaFilter] = [
+            cast(SigmaFilter, rule) for rule in self.rules if isinstance(rule, SigmaFilter)
+        ]
         self.rules = [rule for rule in self.rules if not isinstance(rule, SigmaFilter)]
 
         # Apply filters on each rule and replace the rule with the filtered rule
@@ -126,34 +128,33 @@ def from_dicts(
             if isinstance(
                 rule, SigmaRule
             ):  # Included rules are already parsed, skip collection action processing
-                parsed_rule = rule
-                parsed_rules.append(parsed_rule)
-                parsed_rule.source = source
+                parsed_rules.append(rule)
+                rule.source = source
             else:
                 action = rule.get("action")
                 if action is None:  # no action defined
                     if "correlation" in rule:  # correlation rule - no global rule merge
-                        parsed_rule = SigmaCorrelationRule.from_dict(
+                        parsed_correlation_rule = SigmaCorrelationRule.from_dict(
                             rule,
                             collect_errors,
                             source,
                         )
-                        parsed_rules.append(parsed_rule)
-                        errors.extend(parsed_rule.errors)  # Propagate errors from rule
+                        parsed_rules.append(parsed_correlation_rule)
+                        errors.extend(parsed_correlation_rule.errors)  # Propagate errors from rule
                     elif "filter" in rule:  # correlation rule - no global rule merge
-                        parsed_rule = SigmaFilter.from_dict(
+                        parsed_filter_rule = SigmaFilter.from_dict(
                             rule,
                             collect_errors,
                             source,
                         )
-                        parsed_rules.append(parsed_rule)
-                        errors.extend(parsed_rule.errors)  # Propagate errors from rule
+                        parsed_rules.append(parsed_filter_rule)
+                        errors.extend(parsed_filter_rule.errors)  # Propagate errors from rule
                     else:  # merge with global rule and parse as simple rule
-                        parsed_rule = SigmaRule.from_dict(
+                        parsed_merged_rule = SigmaRule.from_dict(
                             deep_dict_update(rule, global_rule), collect_errors, source
                         )
-                        parsed_rules.append(parsed_rule)
-                        errors.extend(parsed_rule.errors)  # Propagate errors from rule
+                        parsed_rules.append(parsed_merged_rule)
+                        errors.extend(parsed_merged_rule.errors)  # Propagate errors from rule
                         prev_rule = rule
                 elif action == "global":  # set global rule template
                     del rule["action"]
@@ -245,6 +246,7 @@ def load_ruleset(
         :param recursion_pattern: Pattern used to recurse into directories, defaults to ``**/*.yml``.
 
         :return: :class:`SigmaCollection` of all sigma rules contained in given paths.
+
         """
         if not isinstance(inputs, Iterable) or isinstance(inputs, str):
             raise TypeError(

sigma/conversion/base.py

@@ -1984,7 +1984,7 @@ def convert_correlation_rule_from_template(
         rule: SigmaCorrelationRule,
         correlation_type: SigmaCorrelationTypeLiteral,
         method: str,
-    ) -> str:
+    ) -> List[str]:
         template = (
             getattr(self, f"{correlation_type}_correlation_query") or self.default_correlation_query
         )
@@ -2202,7 +2202,7 @@ def convert_correlation_aggregation_fields_from_template(
         referenced_rules: List[SigmaRuleReference],
         group_by: Optional[List[str]],
         method: str,
-    ):
+    ) -> str:
         if self.correlation_fields_expression is None:
             return ""
         else:
@@ -2215,9 +2215,13 @@ def convert_correlation_aggregation_fields_from_template(
             # Include fields from the correlation rule
             for fld in referenced_rules_fields + correlation_rule_fields:
                 # Exclude groupby fields and keep only unique fields (remove duplicates)
-                if fld not in group_by and fld not in all_fields:
+                if (group_by is None or fld not in group_by) and fld not in all_fields:
                     all_fields.append(fld)
-            if len(all_fields) == 0:  # if no fields
+            if (
+                len(all_fields) == 0
+                or self.correlation_fields_field_expression is None
+                or self.correlation_fields_field_expression_joiner is None
+            ):
                 return ""
             return self.correlation_fields_expression[method].format(
                 fields=self.correlation_fields_field_expression_joiner[method].join(

sigma/correlations.py

@@ -423,15 +423,17 @@ def resolve_rule_references(self, rule_collection: "SigmaCollection") -> None:
 
         self.aliases.resolve_rule_references(rule_collection)
 
-    def flatten_rules(self, include_correlations: bool = True) -> List[SigmaRule]:
+    def flatten_rules(
+        self, include_correlations: bool = True
+    ) -> List[Union[SigmaRule, "SigmaCorrelationRule"]]:
         """
         Flattens the rules in the correlation rule and returns a list of Sigma rules. If include_correlations
         is set to False, only the Sigma rules are returned, excluding nested correlation rules.
 
         Returns:
             List of Sigma rules.
         """
-        rules = []
+        rules: List[Union[SigmaRule, "SigmaCorrelationRule"]] = []
         for rule_ref in self.rules:
             rule = rule_ref.rule
             if isinstance(rule, SigmaCorrelationRule):

sigma/modifiers.py

@@ -411,7 +411,7 @@ class SigmaNotEqualModifier(SigmaCompareModifier):
     op: ClassVar[CompareOperators] = CompareOperators.NEQ
 
 
-class SigmaFieldReferenceModifier(SigmaValueModifier):
+class SigmaFieldReferenceModifier(SigmaValueModifier[SigmaString, SigmaFieldReference]):
     """Modifiers a plain string into the field reference type."""
 
     def modify(self, val: SigmaString) -> SigmaFieldReference:
@@ -434,7 +434,11 @@ def modify(self, val: SigmaBool) -> SigmaExists:
         return SigmaExists(val.boolean)
 
 
-class SigmaExpandModifier(SigmaValueModifier[SigmaString, SigmaString]):
+class SigmaExpandModifier(
+    SigmaValueModifier[
+        Union[SigmaString, SigmaRegularExpression], Union[SigmaString, SigmaRegularExpression]
+    ]
+):
     """
     Modifier for expansion of placeholders in values. It replaces placeholder strings (%something%)
     with stub objects that are later expanded to one or multiple strings or replaced with some SIEM

sigma/processing/resolver.py

@@ -7,6 +7,7 @@
 )
 from sigma.processing.pipeline import ProcessingPipeline
 from typing import Dict, Iterable, List, Optional, Tuple, Union, cast, Callable
+from collections import namedtuple
 
 
 @dataclass
@@ -81,11 +82,13 @@ def resolve(
         compatibility check for the usage of the specified backend with the pipeline.
         """
 
-        def resolve_path(spec):
+        PipelineInfo = namedtuple("PipelineInfo", ["pipeline", "priority", "path"])
+
+        def resolve_path(spec: str) -> PipelineInfo:
             pipeline = self.resolve_pipeline(spec, target)
-            return {"pipeline": pipeline, "priority": pipeline.priority, "path": spec}
+            return PipelineInfo(pipeline=pipeline, priority=pipeline.priority, path=spec)
 
-        def resolve_spec(pipelines, spec):
+        def resolve_spec(pipelines: List[PipelineInfo], spec: str) -> List[PipelineInfo]:
             spec_path = Path(spec.rstrip("/*"))
             if spec_path.is_dir():
                 pipelines.extend([resolve_path(str(path)) for path in spec_path.glob("**/*.yml")])
@@ -94,11 +97,9 @@ def resolve_spec(pipelines, spec):
 
             return pipelines
 
-        pipelines = reduce(resolve_spec, pipeline_specs, [])
+        pipelines: List[PipelineInfo] = reduce(resolve_spec, pipeline_specs, [])
 
         return (
-            sum(
-                [p["pipeline"] for p in sorted(pipelines, key=lambda p: (p["priority"], p["path"]))]
-            )
+            sum([p.pipeline for p in sorted(pipelines, key=lambda p: (p.priority, p.path))])
             or ProcessingPipeline()
         )

sigma/processing/transformations/__init__.py

@@ -69,8 +69,8 @@
 }
 
 __all__ = [
-    #    "Transformation",
-    #    "PreprocessingTransformation",
+    "Transformation",
+    "PreprocessingTransformation",
     "FieldMappingTransformation",
     "FieldPrefixMappingTransformation",
     "FieldFunctionTransformation",

sigma/processing/transformations/failure.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass
 from typing import Set, Union
 from sigma.correlations import SigmaCorrelationRule
-from sigma.exceptions import SigmaTransformationError, SigmaTypeError
+from sigma.exceptions import SigmaTransformationError
 from sigma.processing.transformations.base import (
     PreprocessingTransformation,
     DetectionItemTransformation,
@@ -72,32 +72,35 @@ def _get_fields_from_detection(self, detection: SigmaDetection) -> Set[str]:
 
         return field_names
 
-    def apply(self, rule: SigmaRule) -> None:
+    def apply(self, rule: Union[SigmaRule, SigmaCorrelationRule]) -> None:
         super().apply(rule)
-
-        pipeline: "sigma.processing.pipeline.ProcessingPipeline" = self._pipeline
-        field_mappings: "sigma.processing.tracking.FieldMappingTracking" = pipeline.field_mappings
-
-        # Get all field names used in the rule (after any transformations have been applied)
-        all_fields = self._get_all_field_names(rule)
-
-        # Check which original fields from the rule were not explicitly mapped
-        # We need to check the target_fields reverse mapping to see which original fields
-        # are represented by the current field names
-        unmapped_fields = []
-
-        for field in all_fields:
-            # Check if this field is in the target_fields (meaning it was mapped from an original field)
-            # or if it's in the field_mappings keys (meaning it was an original field that was mapped)
-            is_mapped = field in field_mappings or field in field_mappings.target_fields
-            if not is_mapped:
-                unmapped_fields.append(field)
-
-        # Raise error if there are unmapped fields
-        if unmapped_fields:
-            unmapped_fields_str = ", ".join(
-                unmapped_fields
-            )  # Create a comma-separated list of unmapped fields
-            raise SigmaTransformationError(
-                f"The following fields are not mapped: {unmapped_fields_str}", source=rule.source
-            )
+        if isinstance(rule, SigmaRule):
+            pipeline = self._pipeline
+            if pipeline is None:
+                raise SigmaTransformationError("Pipeline is not set for the transformation.")
+            field_mappings = pipeline.field_mappings
+
+            # Get all field names used in the rule (after any transformations have been applied)
+            all_fields = self._get_all_field_names(rule)
+
+            # Check which original fields from the rule were not explicitly mapped
+            # We need to check the target_fields reverse mapping to see which original fields
+            # are represented by the current field names
+            unmapped_fields = []
+
+            for field in all_fields:
+                # Check if this field is in the target_fields (meaning it was mapped from an original field)
+                # or if it's in the field_mappings keys (meaning it was an original field that was mapped)
+                is_mapped = field in field_mappings or field in field_mappings.target_fields
+                if not is_mapped:
+                    unmapped_fields.append(field)
+
+            # Raise error if there are unmapped fields
+            if unmapped_fields:
+                unmapped_fields_str = ", ".join(
+                    unmapped_fields
+                )  # Create a comma-separated list of unmapped fields
+                raise SigmaTransformationError(
+                    f"The following fields are not mapped: {unmapped_fields_str}",
+                    source=rule.source,
+                )

sigma/processing/transformations/placeholder.py

@@ -54,19 +54,23 @@ def __post_init__(self) -> None:
         self.check_exclusivity()
         return super().__post_init__()
 
-    def apply_value(
-        self, field: str, val: Union[SigmaString, SigmaRegularExpression]
-    ) -> Union[
-        SigmaString, Iterable[SigmaString], SigmaRegularExpression, Iterable[SigmaRegularExpression]
+    def apply_value(self, field: Optional[str], val: SigmaType) -> Union[
+        None,
+        SigmaString,
+        Iterable[SigmaString],
+        SigmaRegularExpression,
+        Iterable[SigmaRegularExpression],
     ]:
-        if val.contains_placeholder(self.include, self.exclude):
+        if isinstance(val, (SigmaString, SigmaRegularExpression)) and val.contains_placeholder(
+            self.include, self.exclude
+        ):
             return val.replace_placeholders(self.placeholder_replacements_base)
         else:
             return None
 
     def placeholder_replacements_base(
         self, p: Placeholder
-    ) -> Iterable[Union[str, SpecialChars, Placeholder, SigmaString]]:
+    ) -> Iterator[Union[str, SpecialChars, Placeholder, SigmaString]]:
         """
         Base placeholder replacement callback. Calls real callback if placeholder is included or not excluded,
         else it passes the placeholder back to caller.

sigma/types.py

@@ -478,7 +478,7 @@ def contains_placeholder(
     def replace_placeholders(
         self,
         callback: Callable[
-            [Placeholder], Iterable[Union[str, SpecialChars, Placeholder, "SigmaString"]]
+            [Placeholder], Iterator[Union[str, SpecialChars, Placeholder, "SigmaString"]]
         ],
     ) -> List["SigmaString"]:
         """
@@ -690,15 +690,6 @@ def __init__(self, timestamp_part: TimestampPart, number: int):
         super().__init__(number)
 
 
-class SigmaTimestampPart(SigmaNumber):
-
-    timestamp_part: TimestampPart
-
-    def __init__(self, timestamp_part: TimestampPart, number: int):
-        self.timestamp_part = timestamp_part
-        super().__init__(number)
-
-
 @dataclass
 class SigmaBool(SigmaType):
     """Boolean value type"""
@@ -831,7 +822,10 @@ def insert_placeholders(self) -> "SigmaRegularExpression":
         return self
 
     def replace_placeholders(
-        self, callback: Callable[[Placeholder], Iterator[Union[str, SpecialChars, Placeholder]]]
+        self,
+        callback: Callable[
+            [Placeholder], Iterator[Union[str, SpecialChars, Placeholder, "SigmaString"]]
+        ],
     ) -> List["SigmaRegularExpression"]:
         """
         Replace all occurrences of string part matching regular expression with placeholder.