diff --git a/capa/features/freeze.py b/capa/features/freeze.py index 37148957..bf4f0b66 100644 --- a/capa/features/freeze.py +++ b/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)) diff --git a/tests/fixtures.py b/tests/fixtures.py index 50a766d8..a1812b15 100644 --- a/tests/fixtures.py +++ b/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) diff --git a/tests/test_main.py b/tests/test_main.py index e9cb3d00..9125c123 100644 --- a/tests/test_main.py +++ b/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 *