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

Code before:
import idc
import idaapi
import ida_nalt
import idautils
import ida_bytes
import ida_segment

from capa.features.address import AbsoluteVirtualAddress
from capa.features.extractors.base_extractor import FunctionHandle


def find_byte_sequence(start: int, end: int, seq: bytes) -> Iterator[int]:
    """yield all ea of a given byte sequence

    args:
        start: min virtual address
        end: max virtual address
        seq: bytes to search e.g. b"\x01\x03"
    """
    patterns = ida_bytes.compiled_binpat_vec_t()
    encoding = ida_nalt.get_default_encoding_idx(ida_nalt.BPU_1B)

    seqstr = " ".join([f"{b:02x}" for b in seq])
    err = ida_bytes.parse_binpat_str(patterns, 0, seqstr, 16, encoding)

    if not err:
        while True:
            ea = ida_bytes.bin_search(start, end, patterns, ida_bytes.BIN_SEARCH_FORWARD)
            if ea == idaapi.BADADDR:
                break
            start = ea + 1
            yield ea


def get_functions(
    start: Optional[int] = None, end: Optional[int] = None, skip_thunks: bool = False, skip_libs: bool = False
) -> Iterator[FunctionHandle]:
    """get functions, range optional

    args:
        start: min virtual address
        end: max virtual address
    """
    for ea in idautils.Functions(start=start, end=end):
        f = idaapi.get_func(ea)
        if not (skip_thunks and (f.flags & idaapi.FUNC_THUNK) or skip_libs and (f.flags & idaapi.FUNC_LIB)):
            yield FunctionHandle(address=AbsoluteVirtualAddress(ea), inner=f)


def get_segments(skip_header_segments=False) -> Iterator[idaapi.segment_t]:
    """get list of segments (sections) in the binary image

Review comment (suggestion):
Let’s change to this to return if `err` is set so we can avoid nesting the following `while` statement.

Code after:
import idc
import idaapi
import ida_nalt
import idautils
import ida_bytes
import ida_segment

from capa.features.address import AbsoluteVirtualAddress
from capa.features.extractors.base_extractor import FunctionHandle

IDA_NALT_ENCODING = ida_nalt.get_default_encoding_idx(ida_nalt.BPU_1B)  # use one byte-per-character encoding


def find_byte_sequence(start: int, end: int, seq: bytes) -> Iterator[int]:
    """yield all ea of a given byte sequence

    args:
        start: min virtual address
        end: max virtual address
        seq: bytes to search e.g. b"\x01\x03"
    """
    patterns = ida_bytes.compiled_binpat_vec_t()

    seqstr = " ".join([f"{b:02x}" for b in seq])
    err = ida_bytes.parse_binpat_str(patterns, 0, seqstr, 16, IDA_NALT_ENCODING)

    if err:
        return

    while True:
        ea = ida_bytes.bin_search(start, end, patterns, ida_bytes.BIN_SEARCH_FORWARD)
        if ea == idaapi.BADADDR:
            break
        start = ea + 1
        yield ea


def get_functions(
    start: Optional[int] = None, end: Optional[int] = None, skip_thunks: bool = False, skip_libs: bool = False
) -> Iterator[FunctionHandle]:
    """get functions, range optional

    args:
        start: min virtual address
        end: max virtual address
    """
    for ea in idautils.Functions(start=start, end=end):
        f = idaapi.get_func(ea)
        if not (skip_thunks and (f.flags & idaapi.FUNC_THUNK) or skip_libs and (f.flags & idaapi.FUNC_LIB)):
            yield FunctionHandle(address=AbsoluteVirtualAddress(ea), inner=f)