Merge pull request #185 from fireeye/fix-168

main: don't crash when the .viv cannot be saved
This commit is contained in:
Willi Ballenthin
2020-07-23 17:13:57 -06:00
committed by GitHub
+21 -9
View File
@@ -210,7 +210,7 @@ def is_supported_file_type(sample):
SHELLCODE_BASE = 0x690000
def get_shellcode_vw(sample, arch="auto"):
def get_shellcode_vw(sample, arch="auto", should_save=True):
"""
Return shellcode workspace using explicit arch or via auto detect
"""
@@ -222,12 +222,17 @@ def get_shellcode_vw(sample, arch="auto"):
# choose arch with most functions, idea by Jay G.
vw_cands = []
for arch in ["i386", "amd64"]:
vw_cands.append(viv_utils.getShellcodeWorkspace(sample_bytes, arch, base=SHELLCODE_BASE))
vw_cands.append(
viv_utils.getShellcodeWorkspace(sample_bytes, arch, base=SHELLCODE_BASE, should_save=should_save)
)
if not vw_cands:
raise ValueError("could not generate vivisect workspace")
vw = max(vw_cands, key=lambda vw: len(vw.getFunctions()))
else:
vw = viv_utils.getShellcodeWorkspace(sample_bytes, arch, base=SHELLCODE_BASE)
vw = viv_utils.getShellcodeWorkspace(sample_bytes, arch, base=SHELLCODE_BASE, should_save=should_save)
vw.setMeta("StorageName", "%s.viv" % sample)
return vw
@@ -246,20 +251,20 @@ class UnsupportedFormatError(ValueError):
pass
def get_workspace(path, format):
def get_workspace(path, format, should_save=True):
import viv_utils
logger.debug("generating vivisect workspace for: %s", path)
if format == "auto":
if not is_supported_file_type(path):
raise UnsupportedFormatError()
vw = viv_utils.getWorkspace(path)
vw = viv_utils.getWorkspace(path, should_save=should_save)
elif format == "pe":
vw = viv_utils.getWorkspace(path)
vw = viv_utils.getWorkspace(path, should_save=should_save)
elif format == "sc32":
vw = get_shellcode_vw(path, arch="i386")
vw = get_shellcode_vw(path, arch="i386", should_save=should_save)
elif format == "sc64":
vw = get_shellcode_vw(path, arch="amd64")
vw = get_shellcode_vw(path, arch="amd64", should_save=should_save)
logger.debug("%s", get_meta_str(vw))
return vw
@@ -267,7 +272,14 @@ def get_workspace(path, format):
def get_extractor_py2(path, format):
import capa.features.extractors.viv
vw = get_workspace(path, format)
vw = get_workspace(path, format, should_save=False)
try:
vw.saveWorkspace()
except IOError:
# see #168 for discussion around how to handle non-writable directories
logger.info("source directory is not writable, won't save intermediate workspace")
return capa.features.extractors.viv.VivisectFeatureExtractor(vw, path)