From 177605aaf8bc06c6dfb08ba9c12b29ea420f426d Mon Sep 17 00:00:00 2001
From: William Ballenthin <william.ballenthin@fireeye.com>
Date: Mon, 15 Mar 2021 13:38:29 -0600
Subject: [PATCH] flirt: enable only on py3, ignore otherwise

---
 capa/features/extractors/viv/__init__.py | 11 ++++++++---
 capa/features/extractors/viv/insn.py     | 13 ++++++++-----
 capa/main.py                             |  3 ++-
 3 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/capa/features/extractors/viv/__init__.py b/capa/features/extractors/viv/__init__.py
index f41b4b3c..9ac8fe3f 100644
--- a/capa/features/extractors/viv/__init__.py
+++ b/capa/features/extractors/viv/__init__.py
@@ -5,12 +5,13 @@
 # Unless required by applicable law or agreed to in writing, software distributed under the License
 #  is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and limitations under the License.
-
+import sys
 import types
 import logging
 
 import viv_utils
-import viv_utils.flirt
+if sys.version_info >= (3, 0):
+    import viv_utils.flirt
 
 import capa.features.extractors
 import capa.features.extractors.viv.file
@@ -84,7 +85,11 @@ class VivisectFeatureExtractor(FeatureExtractor):
             yield feature, va
 
     def is_library_function(self, va):
-        return viv_utils.flirt.is_library_function(self.vw, va)
+        if sys.version_info >= (3, 0):
+            import viv_utils.flirt
+            return viv_utils.flirt.is_library_function(self.vw, va)
+        else:
+            return False
 
     def get_function_name(self, va):
         return viv_utils.get_function_name(self.vw, va)
diff --git a/capa/features/extractors/viv/insn.py b/capa/features/extractors/viv/insn.py
index 3679a46c..4e78a3da 100644
--- a/capa/features/extractors/viv/insn.py
+++ b/capa/features/extractors/viv/insn.py
@@ -5,10 +5,12 @@
 # Unless required by applicable law or agreed to in writing, software distributed under the License
 #  is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and limitations under the License.
+import sys
 
 import viv_utils
 import envi.memory
-import viv_utils.flirt
+if sys.version_info >= (3, 0):
+    import viv_utils.flirt
 import envi.archs.i386.disasm
 
 import capa.features.extractors.viv
@@ -113,10 +115,11 @@ def extract_insn_api_features(f, bb, insn):
         if not target:
             return
 
-        if viv_utils.flirt.is_library_function(f.vw, target):
-            name = viv_utils.get_function_name(f.vw, target)
-            yield API(name), insn.va
-            return
+        if sys.version_info >= (3, 0):
+            if viv_utils.flirt.is_library_function(f.vw, target):
+                name = viv_utils.get_function_name(f.vw, target)
+                yield API(name), insn.va
+                return
 
         for _ in range(THUNK_CHAIN_DEPTH_DELTA):
             if target in imports:
diff --git a/capa/main.py b/capa/main.py
index f0da7393..84310005 100644
--- a/capa/main.py
+++ b/capa/main.py
@@ -380,7 +380,8 @@ def get_workspace(path, format, sigpaths):
     else:
         raise ValueError("unexpected format: " + format)
 
-    register_flirt_signature_analyzers(vw, sigpaths)
+    if sys.version_info >= (3, 0):
+        register_flirt_signature_analyzers(vw, sigpaths)
 
     vw.analyze()