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

Code before:
#  is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and limitations under the License.
import functools
from typing import Any, Dict, Tuple, Iterator, Optional

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 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):

Review comment (suggestion):
`find_byte_sequence` may be called many times. Can we move the initialization of `patterns` and `encoding` to the beginning of the file at the global scope so they are only called once? If so, while making this change let's update their names to be more descriptive, e.g. `IDA_BYTES_PATTERNS` and `IDA_NALT_ENCODING`.

Code after:
#  is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and limitations under the License.
import functools
from typing import Any, Dict, Tuple, Iterator, Optional

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
    """