From e509105229ebdd60f30f1a9ff5d92d4ed30b54ae Mon Sep 17 00:00:00 2001
From: William Ballenthin <william.ballenthin@fireeye.com>
Date: Thu, 23 Jul 2020 10:49:14 -0600
Subject: [PATCH 1/2] main: don't crash when the .viv cannot be saved

closes #168
---
 capa/main.py | 27 ++++++++++++++++++---------
 1 file changed, 18 insertions(+), 9 deletions(-)

diff --git a/capa/main.py b/capa/main.py
index 4e5455e9..d6f38d0d 100644
--- a/capa/main.py
+++ b/capa/main.py
@@ -212,7 +212,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
     """
@@ -224,12 +224,14 @@ 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)
     return vw
 
 
@@ -248,20 +250,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
 
@@ -269,7 +271,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)
 
 

From 7665581c967a7eea630f05407e26377f3cdaf39d Mon Sep 17 00:00:00 2001
From: William Ballenthin <william.ballenthin@fireeye.com>
Date: Thu, 23 Jul 2020 17:12:22 -0600
Subject: [PATCH 2/2] main: set shellcode save path

---
 capa/main.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/capa/main.py b/capa/main.py
index d6f38d0d..f085f6a2 100644
--- a/capa/main.py
+++ b/capa/main.py
@@ -232,6 +232,9 @@ def get_shellcode_vw(sample, arch="auto", should_save=True):
         vw = max(vw_cands, key=lambda vw: len(vw.getFunctions()))
     else:
         vw = viv_utils.getShellcodeWorkspace(sample_bytes, arch, base=SHELLCODE_BASE, should_save=should_save)
+
+    vw.setMeta("StorageName", "%s.viv" % sample)
+
     return vw