From 37bc47c77259abac600a8eb36736ae464c2bf412 Mon Sep 17 00:00:00 2001
From: William Ballenthin <william.ballenthin@fireeye.com>
Date: Wed, 11 Aug 2021 14:41:11 -0600
Subject: [PATCH] extractors: viv: extract from bytes not file path

---
 capa/features/extractors/viv/extractor.py |  4 ++-
 capa/features/extractors/viv/file.py      | 30 +++++++++--------------
 2 files changed, 15 insertions(+), 19 deletions(-)

diff --git a/capa/features/extractors/viv/extractor.py b/capa/features/extractors/viv/extractor.py
index 15f3973a..05fec5a3 100644
--- a/capa/features/extractors/viv/extractor.py
+++ b/capa/features/extractors/viv/extractor.py
@@ -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):
diff --git a/capa/features/extractors/viv/file.py b/capa/features/extractors/viv/file.py
index 00eb57ed..9ec9d672 100644
--- a/capa/features/extractors/viv/file.py
+++ b/capa/features/extractors/viv/file.py
@@ -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