diff --git a/capa/features/extractors/ida/insn.py b/capa/features/extractors/ida/insn.py index 225f2613..d06d01e9 100644 --- a/capa/features/extractors/ida/insn.py +++ b/capa/features/extractors/ida/insn.py @@ -98,8 +98,11 @@ def extract_insn_number_features(f, bb, insn): # .text:00401145 add esp, 0Ch return - for op in capa.features.extractors.ida.helpers.get_insn_ops(insn, target_ops=(idaapi.o_imm,)): - const = capa.features.extractors.ida.helpers.mask_op_val(op) + for op in capa.features.extractors.ida.helpers.get_insn_ops(insn, target_ops=(idaapi.o_imm, idaapi.o_mem)): + if op.type == idaapi.o_imm: + const = capa.features.extractors.ida.helpers.mask_op_val(op) + else: + const = op.addr if not idaapi.is_mapped(const): yield Number(const), insn.ea yield Number(const, arch=get_arch(f.ctx)), insn.ea diff --git a/capa/features/extractors/viv/insn.py b/capa/features/extractors/viv/insn.py index a00cfeb7..5aa8e3e1 100644 --- a/capa/features/extractors/viv/insn.py +++ b/capa/features/extractors/viv/insn.py @@ -128,10 +128,13 @@ def extract_insn_number_features(f, bb, insn): # push 3136B0h ; dwControlCode for oper in insn.opers: # this is for both x32 and x64 - if not isinstance(oper, envi.archs.i386.disasm.i386ImmOper): + if not isinstance(oper, (envi.archs.i386.disasm.i386ImmOper, envi.archs.i386.disasm.i386ImmMemOper)): continue - v = oper.getOperValue(oper) + if isinstance(oper, envi.archs.i386.disasm.i386ImmOper): + v = oper.getOperValue(oper) + else: + v = oper.getOperAddr(oper) if f.vw.probeMemory(v, 1, envi.memory.MM_READ): # this is a valid address diff --git a/tests/fixtures.py b/tests/fixtures.py index 1ea5b193..a4a59f82 100644 --- a/tests/fixtures.py +++ b/tests/fixtures.py @@ -83,3 +83,9 @@ def sample_39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41(): def sample_499c2a85f6e8142c3f48d4251c9c7cd6_raw32(): path = os.path.join(CD, "data", "499c2a85f6e8142c3f48d4251c9c7cd6.raw32") return Sample(viv_utils.getShellcodeWorkspace(path), path) + + +@pytest.fixture +def sample_al_khaser_x86(): + path = os.path.join(CD, "data", "al-khaser_x86.exe_") + return Sample(viv_utils.getWorkspace(path), path) diff --git a/tests/test_viv_features.py b/tests/test_viv_features.py index ce480a7e..4e612bdf 100644 --- a/tests/test_viv_features.py +++ b/tests/test_viv_features.py @@ -126,6 +126,11 @@ def test_number_arch_features(mimikatz): assert capa.features.insn.Number(0xFF, arch=ARCH_X64) not in features +def test_unmapped_immediate_memory_reference_features(sample_al_khaser_x86): + features = extract_function_features(viv_utils.Function(sample_al_khaser_x86.vw, 0x41AAB4)) + assert capa.features.insn.Number(0x7FFE02D4) in features + + def test_offset_features(mimikatz): features = extract_function_features(viv_utils.Function(mimikatz.vw, 0x40105D)) assert capa.features.insn.Offset(0x0) in features