*: py3 compat

capa/features/freeze.py

@@ -101,7 +101,9 @@ def dumps(extractor):
             for feature, va in extractor.extract_basic_block_features(f, bb):
                 ret["scopes"]["basic block"].append(serialize_feature(feature) + (hex(va), (hex(f), hex(bb),)))
 
-            for insn, insnva in sorted([(insn, int(insn)) for insn in extractor.get_instructions(f, bb)]):
+            for insnva, insn in sorted(
+                [(insn.__int__(), insn) for insn in extractor.get_instructions(f, bb)], key=lambda p: p[0]
+            ):
                 ret["functions"][hex(f)][hex(bb)].append(hex(insnva))
 
                 for feature, va in extractor.extract_insn_features(f, bb, insn):
@@ -245,12 +247,7 @@ def main(argv=None):
         logging.basicConfig(level=logging.INFO)
         logging.getLogger().setLevel(logging.INFO)
 
-    vw = capa.main.get_workspace(args.sample, args.format)
-
-    # don't import this at top level to support ida/py3 backend
-    import capa.features.extractors.viv
-
-    extractor = capa.features.extractors.viv.VivisectFeatureExtractor(vw, args.sample)
+    extractor = capa.main.get_extractor(args.sample, args.format)
     with open(args.output, "wb") as f:
         f.write(dump(extractor))
 

tests/fixtures.py

@@ -7,6 +7,7 @@
 # See the License for the specific language governing permissions and limitations under the License.
 
 import os
+import sys
 import os.path
 import collections
 
@@ -40,6 +41,16 @@ def get_viv_extractor(path):
     return capa.features.extractors.viv.VivisectFeatureExtractor(vw, path)
 
 
+@lru_cache
+def get_lancelot_extractor(path):
+    import capa.features.extractors.lancelot
+
+    with open(path, "rb") as f:
+        buf = f.read()
+
+    return capa.features.extractors.lancelot.LancelotFeatureExtractor(buf)
+
+
 @lru_cache()
 def extract_file_features(extractor):
     features = collections.defaultdict(set)
@@ -386,9 +397,10 @@ def do_test_feature_count(get_extractor, sample, scope, feature, expected):
 
 
 def get_extractor(path):
-    # decide here which extractor to load for tests.
-    # maybe check which python version we've loaded or if we're in IDA.
-    extractor = get_viv_extractor(path)
+    if sys.version_info >= (3, 0):
+        extractor = get_lancelot_extractor(path)
+    else:
+        extractor = get_viv_extractor(path)
 
     # overload the extractor so that the fixture exposes `extractor.path`
     setattr(extractor, "path", path)

tests/test_lancelot_features.py

@@ -10,16 +10,6 @@
 from fixtures import *
 
 
-@lru_cache
-def get_lancelot_extractor(path):
-    import capa.features.extractors.lancelot
-
-    with open(path, "rb") as f:
-        buf = f.read()
-
-    return capa.features.extractors.lancelot.LancelotFeatureExtractor(buf)
-
-
 @parametrize(
     "sample,scope,feature,expected", FEATURE_PRESENCE_TESTS, indirect=["sample", "scope"],
 )

tests/test_main.py

@@ -14,7 +14,6 @@ import capa.main
 import capa.rules
 import capa.engine
 import capa.features
-import capa.features.extractors.viv
 from capa.engine import *