Repository: mandiant/capa
File: capa/features/extractors/ida/insn.py (Python)

Code before:
        if not f or not (f.flags & idaapi.FUNC_THUNK):
            break

    if info:
        yield info


def extract_insn_api_features(fh: FunctionHandle, bbh: BBHandle, ih: InsnHandle) -> Iterator[Tuple[Feature, Address]]:
    """
    parse instruction API features

    example:
       call dword [0x00473038]
    """
    insn: idaapi.insn_t = ih.inner

    if insn.get_canon_mnem() not in ("call", "jmp"):
        return

    # check calls to imported functions
    for api in check_for_api_call(insn, get_imports(fh.ctx)):
        # tuple (<module>, <function>, <ordinal>)
        for name in capa.features.extractors.helpers.generate_symbols(api[0], api[1]):
            yield API(name), ih.address
        return

    # check calls to extern functions
    for api in check_for_api_call(insn, get_externs(fh.ctx)):
        # tuple (<module>, <function>, <ordinal>)
        yield API(api[1]), ih.address
        return

    # extract dynamically resolved APIs stored in renamed globals (renamed for example using `renimp.idc`)
    if insn.Op1.type == ida_ua.o_mem:
        op_addr = insn.Op1.addr
        op_name = idaapi.get_name(op_addr)
        # check that the memory address has a name and a type
        if (not op_name.startswith("off_")) and idc.get_type(op_addr):
            # Remove suffix used in repeated names, for example _0 in VirtualFree_0
            match = re.match(r"(.+)_\d+", op_name)
            if match:
                op_name = match.group(1)
            yield API(op_name), ih.address
            if capa.features.extractors.helpers.is_aw_function(op_name):
                yield API(op_name[:-1]), ih.address

    # extract IDA/FLIRT recognized API functions
    targets = tuple(idautils.CodeRefsFrom(insn.ea, False))
    if not targets:
        return


Review comment (refactor):
imports are also calls to operands of type `o_mem`, which means the added code will extract it again. A `call` instruction should call only one function, so I think we can stop if a call to an import/extern is extracted. Let me know if there are any concerns with it.

Code after:
        if not f or not (f.flags & idaapi.FUNC_THUNK):
            break

    return info


def extract_insn_api_features(fh: FunctionHandle, bbh: BBHandle, ih: InsnHandle) -> Iterator[Tuple[Feature, Address]]:
    """
    parse instruction API features

    example:
       call dword [0x00473038]
    """
    insn: idaapi.insn_t = ih.inner

    if insn.get_canon_mnem() not in ("call", "jmp"):
        return

    # check call to imported functions
    api = check_for_api_call(insn, get_imports(fh.ctx))
    if api:
        # tuple (<module>, <function>, <ordinal>)
        for name in capa.features.extractors.helpers.generate_symbols(api[0], api[1]):
            yield API(name), ih.address
        # a call instruction should only call one function, stop if a call to an import is extracted
        return

    # check call to extern functions
    api = check_for_api_call(insn, get_externs(fh.ctx))
    if api:
        # tuple (<module>, <function>, <ordinal>)
        yield API(api[1]), ih.address
        # a call instruction should only call one function, stop if a call to an extern is extracted
        return

    # extract dynamically resolved APIs stored in renamed globals (renamed for example using `renimp.idc`)
    # examples: `CreateProcessA`, `HttpSendRequestA`
    if insn.Op1.type == ida_ua.o_mem:
        op_addr = insn.Op1.addr
        op_name = idaapi.get_name(op_addr)
        # when renaming a global using an API name, IDA assigns it the function type
        # ensure we do not extract something wrong by checking that the address has a name and a type
        # we could check that the type is a function definition, but that complicates the code
        if (not op_name.startswith("off_")) and idc.get_type(op_addr):
            # Remove suffix used in repeated names, for example _0 in VirtualFree_0
            match = re.match(r"(.+)_\d+", op_name)
            if match:
                op_name = match.group(1)
            # the global name does not include the DLL name, so we can't extract it
            for name in capa.features.extractors.helpers.generate_symbols("", op_name):
                yield API(name), ih.address