diff --git a/capa/features/extractors/ida/__init__.py b/capa/features/extractors/ida/__init__.py index d1ca5d31..0d36423f 100644 --- a/capa/features/extractors/ida/__init__.py +++ b/capa/features/extractors/ida/__init__.py @@ -64,6 +64,12 @@ class IdaFeatureExtractor(FeatureExtractor): setattr(f, "ctx", ctx) yield add_ea_int_cast(f) + @staticmethod + def get_function(self, ea): + f = idaapi.get_func(ea) + setattr(f, "ctx", {}) + return add_ea_int_cast(f) + def extract_function_features(self, f): for (feature, ea) in capa.features.extractors.ida.function.extract_features(f): yield feature, ea diff --git a/tests/test_ida_features.py b/tests/test_ida_features.py index a5774588..e84b094d 100644 --- a/tests/test_ida_features.py +++ b/tests/test_ida_features.py @@ -75,10 +75,7 @@ def extract_basic_block_features(f, bb): @pytest.mark.skip(reason="IDA Pro tests must be run within IDA") def test_api_features(): - # have to import import this inline so pytest doesn't bail outside of IDA - import idaapi - - f = idaapi.get_func(0x403BAC) + f = get_extractor().get_function(0x403BAC) features = extract_function_features(f) assert capa.features.insn.API("advapi32.CryptAcquireContextW") in features assert capa.features.insn.API("advapi32.CryptAcquireContext") in features @@ -94,9 +91,7 @@ def test_api_features(): @pytest.mark.skip(reason="IDA Pro tests must be run within IDA") def test_string_features(): - import idaapi - - f = idaapi.get_func(0x40105D) + f = get_extractor().get_function(0x40105D) features = extract_function_features(f) assert capa.features.String("SCardControl") in features assert capa.features.String("SCardTransmit") in features @@ -107,9 +102,7 @@ def test_string_features(): @pytest.mark.skip(reason="IDA Pro tests must be run within IDA") def test_byte_features(): - import idaapi - - f = idaapi.get_func(0x40105D) + f = get_extractor().get_function(0x40105D) features = extract_function_features(f) wanted = capa.features.Bytes("SCardControl".encode("utf-16le")) # use `==` rather than `is` because the result is not `True` but a truthy value. @@ -118,9 +111,7 @@ def test_byte_features(): @pytest.mark.skip(reason="IDA Pro tests must be run within IDA") def test_number_features(): - import idaapi - - f = idaapi.get_func(0x40105D) + f = get_extractor().get_function(0x40105D) features = extract_function_features(f) assert capa.features.insn.Number(0xFF) in features assert capa.features.insn.Number(0x3136B0) in features @@ -131,9 +122,7 @@ def test_number_features(): @pytest.mark.skip(reason="IDA Pro tests must be run within IDA") def test_number_arch_features(): - import idaapi - - f = idaapi.get_func(0x40105D) + f = get_extractor().get_function(0x40105D) features = extract_function_features(f) assert capa.features.insn.Number(0xFF) in features assert capa.features.insn.Number(0xFF, arch=ARCH_X32) in features @@ -142,9 +131,7 @@ def test_number_arch_features(): @pytest.mark.skip(reason="IDA Pro tests must be run within IDA") def test_offset_features(): - import idaapi - - f = idaapi.get_func(0x40105D) + f = get_extractor().get_function(0x40105D) features = extract_function_features(f) assert capa.features.insn.Offset(0x0) in features assert capa.features.insn.Offset(0x4) in features @@ -156,7 +143,7 @@ def test_offset_features(): # this function has the following negative offsets # movzx ecx, byte ptr [eax-1] # movzx eax, byte ptr [eax-2] - f = idaapi.get_func(0x4011FB) + f = get_extractor().get_function(0x4011FB) features = extract_function_features(f) assert capa.features.insn.Offset(-0x1) in features assert capa.features.insn.Offset(-0x2) in features @@ -164,9 +151,7 @@ def test_offset_features(): @pytest.mark.skip(reason="IDA Pro tests must be run within IDA") def test_offset_arch_features(mimikatz): - import idaapi - - f = idaapi.get_func(0x40105D) + f = get_extractor().get_function(0x40105D) features = extract_function_features(f) assert capa.features.insn.Offset(0x0) in features assert capa.features.insn.Offset(0x0, arch=ARCH_X32) in features @@ -175,18 +160,14 @@ def test_offset_arch_features(mimikatz): @pytest.mark.skip(reason="IDA Pro tests must be run within IDA") def test_nzxor_features(): - import idaapi - - f = idaapi.get_func(0x410DFC) + f = get_extractor().get_function(0x410DFC) features = extract_function_features(f) assert capa.features.Characteristic("nzxor") in features # 0x0410F0B @pytest.mark.skip(reason="IDA Pro tests must be run within IDA") def test_mnemonic_features(): - import idaapi - - f = idaapi.get_func(0x40105D) + f = get_extractor().get_function(0x40105D) features = extract_function_features(f) assert capa.features.insn.Mnemonic("push") in features assert capa.features.insn.Mnemonic("movzx") in features @@ -206,10 +187,9 @@ def test_file_section_name_features(): @pytest.mark.skip(reason="IDA Pro tests must be run within IDA") def test_tight_loop_features(): - import idaapi - extractor = get_extractor() - f = idaapi.get_func(0x402EC4) + + f = extractor.get_function(0x402EC4) for bb in extractor.get_basic_blocks(f): if bb.__int__() != 0x402F8E: continue @@ -220,10 +200,9 @@ def test_tight_loop_features(): @pytest.mark.skip(reason="IDA Pro tests must be run within IDA") def test_tight_loop_bb_features(): - import idaapi - extractor = get_extractor() - f = idaapi.get_func(0x402EC4) + + f = extractor.get_function(0x402EC4) for bb in extractor.get_basic_blocks(f): if bb.__int__() != 0x402F8E: continue @@ -245,46 +224,38 @@ def test_file_import_name_features(): @pytest.mark.skip(reason="IDA Pro tests must be run within IDA") def test_stackstring_features(): - import idaapi - - f = idaapi.get_func(0x4556E5) + f = get_extractor().get_function(0x4556E5) features = extract_function_features(f) assert capa.features.Characteristic("stack string") in features @pytest.mark.skip(reason="IDA Pro tests must be run within IDA") def test_switch_features(): - import idaapi - - f = idaapi.get_func(0x409411) + f = get_extractor().get_function(0x409411) features = extract_function_features(f) assert capa.features.Characteristic("switch") in features - f = idaapi.get_func(0x409393) + f = get_extractor().get_function(0x409393) features = extract_function_features(f) assert capa.features.Characteristic("switch") not in features @pytest.mark.skip(reason="IDA Pro tests must be run within IDA") def test_function_calls_to(): - import idaapi - # this function is used in a function pointer - f = idaapi.get_func(0x4011FB) + f = get_extractor().get_function(0x4011FB) features = extract_function_features(f) assert capa.features.Characteristic("calls to") not in features # __FindPESection is called once - f = idaapi.get_func(0x470360) + f = get_extractor().get_function(0x470360) features = extract_function_features(f) assert len(features[capa.features.Characteristic("calls to")]) == 1 @pytest.mark.skip(reason="IDA Pro tests must be run within IDA") def test_function_calls_from(): - import idaapi - - f = idaapi.get_func(0x4011FB) + f = get_extractor().get_function(0x4011FB) features = extract_function_features(f) assert capa.features.Characteristic("calls from") in features assert len(features[capa.features.Characteristic("calls from")]) == 3 @@ -292,9 +263,7 @@ def test_function_calls_from(): @pytest.mark.skip(reason="IDA Pro tests must be run within IDA") def test_basic_block_count(): - import idaapi - - f = idaapi.get_func(0x4011FB) + f = get_extractor().get_function(0x4011FB) features = extract_function_features(f) assert len(features[capa.features.basicblock.BasicBlock()]) == 15