Skip to content

Commit f607aa4

Browse files
committed
WIP
1 parent 30cfb3c commit f607aa4

File tree

16 files changed

+329
-232
lines changed

16 files changed

+329
-232
lines changed

poetry.lock

Lines changed: 12 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ pyparsing = "^3.1"
2727
pyyaml = "^6.0"
2828
requests = "^2.31"
2929
jinja2 = "^3.1"
30+
types-pyyaml = "^6.0.12.20240917"
3031

3132
[tool.poetry.group.dev.dependencies]
3233
black = "^24.4.2"

sigma/backends/test/backend.py

Lines changed: 43 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,19 @@
11
from collections import defaultdict
22
import re
3-
from typing import ClassVar, Dict, Optional, Pattern, Tuple
3+
from typing import Any, ClassVar, Dict, List, Optional, Pattern, Tuple, cast
44

55
from sigma.conversion.base import TextQueryBackend
66
from sigma.conversion.state import ConversionState
77
from sigma.pipelines.test import dummy_test_pipeline
88
from sigma.processing.pipeline import ProcessingItem, ProcessingPipeline
99
from sigma.processing.transformations import FieldMappingTransformation
10+
from sigma.rule.rule import SigmaRule
1011
from sigma.types import CompareOperators, SigmaCompareExpression
1112

1213

1314
class TextQueryTestBackend(TextQueryBackend):
14-
name: str = "Test backend"
15-
formats: Dict[str, str] = {
15+
name: ClassVar[str] = "Test backend"
16+
formats: ClassVar[Dict[str, str]] = {
1617
"default": "Default format",
1718
"test": "Dummy test format",
1819
"state": "Test format that obtains information from state",
@@ -29,15 +30,15 @@ class TextQueryTestBackend(TextQueryBackend):
2930
eq_token: ClassVar[str] = "="
3031

3132
field_quote: ClassVar[str] = "'"
32-
field_quote_pattern: ClassVar[Pattern] = re.compile("^\\w+$")
33+
field_quote_pattern: ClassVar[Pattern[str]] = re.compile("^\\w+$")
3334

3435
str_quote: ClassVar[str] = '"'
3536
escape_char: ClassVar[str] = "\\"
3637
wildcard_multi: ClassVar[str] = "*"
3738
wildcard_single: ClassVar[str] = "?"
3839
add_escaped: ClassVar[str] = ":"
3940
filter_chars: ClassVar[str] = "&"
40-
bool_values: ClassVar[Dict[bool, str]] = {
41+
bool_values: ClassVar[Dict[bool, Optional[str]]] = {
4142
True: "1",
4243
False: "0",
4344
}
@@ -52,7 +53,7 @@ class TextQueryTestBackend(TextQueryBackend):
5253

5354
re_expression: ClassVar[str] = "{field}=/{regex}/"
5455
re_escape_char: ClassVar[str] = "\\"
55-
re_escape: ClassVar[Tuple[str]] = ("/", "bar")
56+
re_escape: ClassVar[List[str]] = ["/", "bar"]
5657

5758
case_sensitive_match_expression = "{field} casematch {value}"
5859
case_sensitive_startswith_expression: ClassVar[str] = "{field} startswith_cased {value}"
@@ -115,8 +116,12 @@ class TextQueryTestBackend(TextQueryBackend):
115116
"test": "Test correlation method",
116117
}
117118
default_correlation_method: ClassVar[str] = "test"
118-
default_correlation_query: ClassVar[str] = {"test": "{search}\n{aggregate}\n{condition}"}
119-
temporal_correlation_query: ClassVar[str] = {"test": "{search}\n\n{aggregate}\n\n{condition}"}
119+
default_correlation_query: ClassVar[Dict[str, str]] = {
120+
"test": "{search}\n{aggregate}\n{condition}"
121+
}
122+
temporal_correlation_query: ClassVar[Dict[str, str]] = {
123+
"test": "{search}\n\n{aggregate}\n\n{condition}"
124+
}
120125

121126
correlation_search_single_rule_expression: ClassVar[str] = "{query}"
122127
correlation_search_multi_rule_expression: ClassVar[str] = "{queries}"
@@ -169,33 +174,39 @@ def __init__(
169174
processing_pipeline: Optional[ProcessingPipeline] = None,
170175
collect_errors: bool = False,
171176
testparam: Optional[str] = None,
172-
**kwargs,
177+
**kwargs: Dict[str, Any],
173178
):
174179
super().__init__(processing_pipeline, collect_errors, **kwargs)
175180
self.testparam = testparam
176181

177-
def finalize_query_test(self, rule, query, index, state):
178-
return "[ " + self.finalize_query_default(rule, query, index, state) + " ]"
182+
def finalize_query_test(
183+
self, rule: SigmaRule, query: str, index: int, state: ConversionState
184+
) -> str:
185+
return "[ " + cast(str, self.finalize_query_default(rule, query, index, state)) + " ]"
179186

180-
def finalize_output_test(self, queries):
181-
return self.finalize_output_default(queries)
187+
def finalize_output_test(self, queries: List[str]) -> str:
188+
return cast(str, self.finalize_output_default(queries))
182189

183-
def finalize_query_state(self, rule, query, index, state: ConversionState):
190+
def finalize_query_state(
191+
self, rule: SigmaRule, query: str, index: int, state: ConversionState
192+
) -> str:
184193
return (
185194
"index="
186-
+ state.processing_state.get("index", "default")
195+
+ cast(str, state.processing_state.get("index", "default"))
187196
+ " ("
188-
+ self.finalize_query_default(rule, query, index, state)
197+
+ cast(str, self.finalize_query_default(rule, query, index, state))
189198
+ ")"
190199
)
191200

192-
def finalize_output_state(self, queries):
193-
return self.finalize_output_default(queries)
201+
def finalize_output_state(self, queries: List[str]) -> str:
202+
return cast(str, self.finalize_output_default(queries))
194203

195-
def finalize_query_list_of_dict(self, rule, query, index, state):
196-
return self.finalize_query_default(rule, query, index, state)
204+
def finalize_query_list_of_dict(
205+
self, rule: SigmaRule, query: str, index: int, state: ConversionState
206+
) -> str:
207+
return cast(str, self.finalize_query_default(rule, query, index, state))
197208

198-
def finalize_output_list_of_dict(self, queries):
209+
def finalize_output_list_of_dict(self, queries: List[str]) -> List[Dict[str, Optional[str]]]:
199210
return [
200211
(
201212
{"query": query, "test": self.testparam}
@@ -205,18 +216,22 @@ def finalize_output_list_of_dict(self, queries):
205216
for query in self.finalize_output_default(queries)
206217
]
207218

208-
def finalize_query_bytes(self, rule, query, index, state):
209-
return self.finalize_query_default(rule, query, index, state)
219+
def finalize_query_bytes(
220+
self, rule: SigmaRule, query: str, index: int, state: ConversionState
221+
) -> str:
222+
return cast(str, self.finalize_query_default(rule, query, index, state))
210223

211-
def finalize_output_bytes(self, queries):
224+
def finalize_output_bytes(self, queries: List[str]) -> bytes:
212225
return bytes("\x00".join(self.finalize_output_default(queries)), "utf-8")
213226

214-
def finalize_query_str(self, rule, query, index, state):
215-
return self.finalize_query_default(rule, query, index, state)
227+
def finalize_query_str(
228+
self, rule: SigmaRule, query: str, index: int, state: ConversionState
229+
) -> str:
230+
return cast(str, self.finalize_query_default(rule, query, index, state))
216231

217-
def finalize_output_str(self, queries):
232+
def finalize_output_str(self, queries: List[str]) -> str:
218233
return "\n".join(self.finalize_output_default(queries))
219234

220235

221236
class MandatoryPipelineTestBackend(TextQueryTestBackend):
222-
requires_pipeline: bool = True
237+
requires_pipeline: ClassVar[bool] = True

sigma/collection.py

Lines changed: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
from dataclasses import dataclass, field
22
from functools import reduce
33
from pathlib import Path
4-
from typing import Callable, Dict, Iterable, List, Optional, Union, IO
4+
from typing import Callable, Dict, Iterable, List, Optional, Union, IO, cast
55
from uuid import UUID
66

77
import yaml
@@ -15,13 +15,16 @@
1515
)
1616
from sigma.rule import SigmaRule, SigmaRuleBase
1717
from sigma.filters import SigmaFilter
18+
from typing import TypeVar, Union
19+
20+
NestedDict = Dict[str, Union[str, int, float, bool, None, "NestedDict"]]
1821

1922

2023
@dataclass
2124
class SigmaCollection:
2225
"""Collection of Sigma rules"""
2326

24-
rules: List[SigmaRuleBase]
27+
rules: List[Union[SigmaRule, SigmaCorrelationRule]]
2528
errors: List[SigmaError] = field(default_factory=list)
2629
ids_to_rules: Dict[UUID, SigmaRuleBase] = field(
2730
init=False, repr=False, hash=False, compare=False
@@ -30,7 +33,7 @@ class SigmaCollection:
3033
init=False, repr=False, hash=False, compare=False
3134
)
3235

33-
def __post_init__(self):
36+
def __post_init__(self) -> None:
3437
"""
3538
Map rule identifiers to rules and resolve rule references in correlation rules.
3639
"""
@@ -42,7 +45,7 @@ def __post_init__(self):
4245
if rule.name is not None:
4346
self.names_to_rules[rule.name] = rule
4447

45-
def resolve_rule_references(self):
48+
def resolve_rule_references(self) -> None:
4649
"""
4750
Resolve rule references in correlation rules to the actual rule objects and sort the rules
4851
by reference order (rules that are referenced by other rules come first).
@@ -55,12 +58,21 @@ def resolve_rule_references(self):
5558
rule.resolve_rule_references(self)
5659

5760
# Extract all filters from the rules
58-
filters: List[SigmaFilter] = [rule for rule in self.rules if isinstance(rule, SigmaFilter)]
61+
filters: List[SigmaFilter] = [
62+
cast(SigmaFilter, rule) for rule in self.rules if isinstance(rule, SigmaFilter)
63+
]
5964
self.rules = [rule for rule in self.rules if not isinstance(rule, SigmaFilter)]
6065

6166
# Apply filters on each rule and replace the rule with the filtered rule
6267
self.rules = (
63-
[reduce(lambda r, f: f.apply_on_rule(r), filters, rule) for rule in self.rules]
68+
[
69+
reduce(
70+
lambda r, f: f.apply_on_rule(r) if isinstance(r, SigmaRule) else r,
71+
filters,
72+
rule,
73+
)
74+
for rule in self.rules
75+
]
6476
if filters
6577
else self.rules
6678
)
@@ -71,7 +83,7 @@ def resolve_rule_references(self):
7183
@classmethod
7284
def from_dicts(
7385
cls,
74-
rules: List[dict],
86+
rules: List[NestedDict],
7587
collect_errors: bool = False,
7688
source: Optional[SigmaRuleLocation] = None,
7789
) -> "SigmaCollection":
@@ -84,7 +96,7 @@ def from_dicts(
8496
errors = []
8597
parsed_rules = list()
8698
prev_rule = None
87-
global_rule = dict()
99+
global_rule: NestedDict = dict()
88100

89101
for i, rule in zip(range(1, len(rules) + 1), rules):
90102
if isinstance(

sigma/conversion/base.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -148,7 +148,7 @@ def __init__(
148148
self,
149149
processing_pipeline: Optional[ProcessingPipeline] = None,
150150
collect_errors: bool = False,
151-
**backend_options: Dict,
151+
**backend_options: Dict[str, Any],
152152
):
153153
self.processing_pipeline = processing_pipeline
154154
self.errors = list()
@@ -846,7 +846,7 @@ class variables. If this is not sufficient, the respective methods can be implem
846846
re_escape_char: ClassVar[Optional[str]] = (
847847
None # Character used for escaping in regular expressions
848848
)
849-
re_escape: ClassVar[Tuple[str]] = () # List of strings that are escaped
849+
re_escape: ClassVar[List[str]] = [] # List of strings that are escaped
850850
re_escape_escape_char: bool = True # If True, the escape character is also escaped
851851
re_flag_prefix: bool = (
852852
True # If True, the flags are prepended as (?x) group at the beginning of the regular expression, e.g. (?i). If this is not supported by the target, it should be set to False.

0 commit comments

Comments
 (0)