add yaml representer for the Scope class, as well as other bugfixes

This commit is contained in:
Yacine Elhamer
2023-08-26 18:11:35 +02:00
parent e9a9b3a6b6
commit 49adecb25c
3 changed files with 11 additions and 1 deletions

View File

@@ -59,6 +59,9 @@ class NullStaticFeatureExtractor(StaticFeatureExtractor):
def get_base_address(self):
return self.base_address
def get_sample_hashes(self) -> SampleHashes:
return self.sample_hashes
def extract_global_features(self):
for feature in self.global_features:
yield feature, NO_ADDRESS
@@ -121,6 +124,9 @@ class NullDynamicFeatureExtractor(DynamicFeatureExtractor):
for feature in self.global_features:
yield feature, NO_ADDRESS
def get_sample_hashes(self) -> SampleHashes:
return self.sample_hashes
def extract_file_features(self):
for address, feature in self.file_features:
yield feature, address

View File

@@ -86,6 +86,10 @@ class Scope(str, Enum):
# not used to validate rules.
GLOBAL = "global"
@classmethod
def to_yaml(cls, representer, node):
return representer.represent_str(f"{node.value}")
# these literals are used to check if the flavor
# of a rule is correct.
@@ -979,6 +983,7 @@ class Rule:
# we use the ruamel.yaml parser because it supports roundtripping of documents with comments.
y = ruamel.yaml.YAML(typ="rt")
y.register_class(Scope)
# use block mode, not inline json-like mode
y.default_flow_style = False

View File

@@ -263,7 +263,6 @@ def assert_round_trip(rd: rdoc.ResultDocument):
pytest.param("a076114_rd"),
pytest.param("pma0101_rd"),
pytest.param("dotnet_1c444e_rd"),
pytest.param(""),
],
)
def test_round_trip(request, rd_file):