This commit is contained in:
William Ballenthin
2021-09-04 15:59:38 -06:00
parent 2ea4dc9d7e
commit 39e4e47763
5 changed files with 13 additions and 9 deletions

View File

@@ -58,7 +58,7 @@ class FeatureExtractor:
fetch the preferred load address at which the sample was analyzed. fetch the preferred load address at which the sample was analyzed.
""" """
raise NotImplemented raise NotImplemented
@abc.abstractmethod @abc.abstractmethod
def extract_global_features(self) -> Iterator[Tuple[Feature, int]]: def extract_global_features(self) -> Iterator[Tuple[Feature, int]]:
""" """
@@ -74,7 +74,7 @@ class FeatureExtractor:
Tuple[Feature, int]: feature and its location Tuple[Feature, int]: feature and its location
""" """
raise NotImplemented raise NotImplemented
@abc.abstractmethod @abc.abstractmethod
def extract_file_features(self) -> Iterator[Tuple[Feature, int]]: def extract_file_features(self) -> Iterator[Tuple[Feature, int]]:
""" """
@@ -272,12 +272,12 @@ class NullFeatureExtractor(FeatureExtractor):
def get_base_address(self): def get_base_address(self):
return self.features["base address"] return self.features["base address"]
def extract_global_features(self): def extract_global_features(self):
for p in self.features.get("global features", []): for p in self.features.get("global features", []):
va, feature = p va, feature = p
yield feature, va yield feature, va
def extract_file_features(self): def extract_file_features(self):
for p in self.features.get("file features", []): for p in self.features.get("file features", []):
va, feature = p va, feature = p

View File

@@ -150,7 +150,7 @@ FILE_HANDLERS = (
extract_file_function_names, extract_file_function_names,
extract_file_format, extract_file_format,
) )
def extract_global_features(pe, buf): def extract_global_features(pe, buf):
""" """

View File

@@ -192,7 +192,7 @@ def loads(s):
va = int(va, 0x10) va = int(va, 0x10)
feature = deserialize_feature(feature[:2]) feature = deserialize_feature(feature[:2])
features["global features"].append((va, feature)) features["global features"].append((va, feature))
for feature in doc.get("scopes", {}).get("file", []): for feature in doc.get("scopes", {}).get("file", []):
va, loc = feature[2:] va, loc = feature[2:]
va = int(va, 0x10) va = int(va, 0x10)

View File

@@ -88,12 +88,16 @@ def find_function_capabilities(ruleset: RuleSet, extractor: FeatureExtractor, f:
# - basic blocks # - basic blocks
bb_features = collections.defaultdict(set) bb_features = collections.defaultdict(set)
for feature, va in itertools.chain(extractor.extract_basic_block_features(f, bb), extractor.extract_global_features()): for feature, va in itertools.chain(
extractor.extract_basic_block_features(f, bb), extractor.extract_global_features()
):
bb_features[feature].add(va) bb_features[feature].add(va)
function_features[feature].add(va) function_features[feature].add(va)
for insn in extractor.get_instructions(f, bb): for insn in extractor.get_instructions(f, bb):
for feature, va in itertools.chain(extractor.extract_insn_features(f, bb, insn), extractor.extract_global_features()): for feature, va in itertools.chain(
extractor.extract_insn_features(f, bb, insn), extractor.extract_global_features()
):
bb_features[feature].add(va) bb_features[feature].add(va)
function_features[feature].add(va) function_features[feature].add(va)

View File

@@ -139,7 +139,7 @@ def extract_global_features(extractor):
for feature, va in extractor.extract_global_features(): for feature, va in extractor.extract_global_features():
features[feature].add(va) features[feature].add(va)
return features return features
@lru_cache() @lru_cache()
def extract_file_features(extractor): def extract_file_features(extractor):