extractors: viv: extract from bytes not file path

This commit is contained in:
William Ballenthin
2021-08-11 14:41:11 -06:00
parent baaa8ba2c1
commit 37bc47c772
2 changed files with 15 additions and 19 deletions

View File

@@ -37,13 +37,15 @@ class VivisectFeatureExtractor(FeatureExtractor):
super(VivisectFeatureExtractor, self).__init__()
self.vw = vw
self.path = path
with open(self.path, "rb") as f:
self.buf = f.read()
def get_base_address(self):
# assume there is only one file loaded into the vw
return list(self.vw.filemeta.values())[0]["imagebase"]
def extract_file_features(self):
for feature, va in capa.features.extractors.viv.file.extract_features(self.vw, self.path):
for feature, va in capa.features.extractors.viv.file.extract_features(self.vw, self.buf):
yield feature, va
def get_functions(self):

View File

@@ -17,20 +17,17 @@ from capa.features.file import Export, Import, Section, FunctionName
from capa.features.common import String, Characteristic
def extract_file_embedded_pe(vw, file_path):
with open(file_path, "rb") as f:
fbytes = f.read()
for offset, i in pe_carve.carve(fbytes, 1):
def extract_file_embedded_pe(vw, buf):
for offset, i in pe_carve.carve(buf, 1):
yield Characteristic("embedded pe"), offset
def extract_file_export_names(vw, file_path):
def extract_file_export_names(vw, buf):
for va, etype, name, _ in vw.getExports():
yield Export(name), va
def extract_file_import_names(vw, file_path):
def extract_file_import_names(vw, buf):
"""
extract imported function names
1. imports by ordinal:
@@ -64,26 +61,23 @@ def is_viv_ord_impname(impname: str) -> bool:
return True
def extract_file_section_names(vw, file_path):
def extract_file_section_names(vw, buf):
for va, _, segname, _ in vw.getSegments():
yield Section(segname), va
def extract_file_strings(vw, file_path):
def extract_file_strings(vw, buf):
"""
extract ASCII and UTF-16 LE strings from file
"""
with open(file_path, "rb") as f:
b = f.read()
for s in capa.features.extractors.strings.extract_ascii_strings(b):
for s in capa.features.extractors.strings.extract_ascii_strings(buf):
yield String(s.s), s.offset
for s in capa.features.extractors.strings.extract_unicode_strings(b):
for s in capa.features.extractors.strings.extract_unicode_strings(buf):
yield String(s.s), s.offset
def extract_file_function_names(vw, file_path):
def extract_file_function_names(vw, buf):
"""
extract the names of statically-linked library functions.
"""
@@ -93,20 +87,20 @@ def extract_file_function_names(vw, file_path):
yield FunctionName(name), va
def extract_features(vw, file_path):
def extract_features(vw, buf: bytes):
"""
extract file features from given workspace
args:
vw (vivisect.VivWorkspace): the vivisect workspace
file_path: path to the input file
buf: the raw input file bytes
yields:
Tuple[Feature, VA]: a feature and its location.
"""
for file_handler in FILE_HANDLERS:
for feature, va in file_handler(vw, file_path):
for feature, va in file_handler(vw, buf):
yield feature, va