Skip to content

Commit

Permalink
add yaml representer for the Scope class, as well as other bugfixes
Browse files Browse the repository at this point in the history
  • Loading branch information
yelhamer committed Aug 26, 2023
1 parent e9a9b3a commit 49adecb
Show file tree
Hide file tree
Showing 3 changed files with 11 additions and 1 deletion.
6 changes: 6 additions & 0 deletions capa/features/extractors/null.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -121,6 +124,9 @@ def extract_global_features(self):
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
Expand Down
5 changes: 5 additions & 0 deletions capa/rules/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -979,6 +983,7 @@ def _get_ruamel_yaml_parser():

# 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
Expand Down
1 change: 0 additions & 1 deletion tests/test_result_document.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down

0 comments on commit 49adecb

Please sign in to comment.