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

Code before:

import idaapi

import capa.features.extractors.ida.helpers
from capa.features.common import Feature, Characteristic
from capa.features.address import Address
from capa.features.basicblock import BasicBlock
from capa.features.extractors.ida import helpers
from capa.features.extractors.helpers import MIN_STACKSTRING_LEN
from capa.features.extractors.base_extractor import BBHandle, FunctionHandle


def get_printable_len(op: idaapi.op_t) -> int:
    """Return string length if all operand bytes are ascii or utf16-le printable"""
    op_val = capa.features.extractors.ida.helpers.mask_op_val(op)

    if op.dtype == idaapi.dt_byte:
        chars = struct.pack("<B", op_val)
    elif op.dtype == idaapi.dt_word:
        chars = struct.pack("<H", op_val)
    elif op.dtype == idaapi.dt_dword:
        chars = struct.pack("<I", op_val)
    elif op.dtype == idaapi.dt_qword:
        chars = struct.pack("<Q", op_val)
    else:
        raise ValueError(f"Unhandled operand data type {hex(op.dtype)}.")

    def is_printable_ascii(chars_: bytes):
        return all(c < 127 and chr(c) in string.printable for c in chars_)

    def is_printable_utf16le(chars_: bytes):
        if all(c == 0x00 for c in chars_[1::2]):
            return is_printable_ascii(chars_[::2])

    if is_printable_ascii(chars):
        return idaapi.get_dtype_size(op.dtype)

    if is_printable_utf16le(chars):
        return idaapi.get_dtype_size(op.dtype) // 2

    return 0


def is_mov_imm_to_stack(insn: idaapi.insn_t) -> bool:
    """verify instruction moves immediate onto stack"""
    if insn.Op2.type != idaapi.o_imm:
        return False

    if not helpers.is_op_stack_var(insn.ea, 0):
        return False


Review comment (suggestion):
```suggestion
        raise ValueError(f"Unhandled operand data type 0x{op.dtype:x}.")
```

Code after:

import idaapi

import capa.features.extractors.ida.helpers
from capa.features.common import Feature, Characteristic
from capa.features.address import Address
from capa.features.basicblock import BasicBlock
from capa.features.extractors.ida import helpers
from capa.features.extractors.helpers import MIN_STACKSTRING_LEN
from capa.features.extractors.base_extractor import BBHandle, FunctionHandle


def get_printable_len(op: idaapi.op_t) -> int:
    """Return string length if all operand bytes are ascii or utf16-le printable"""
    op_val = capa.features.extractors.ida.helpers.mask_op_val(op)

    if op.dtype == idaapi.dt_byte:
        chars = struct.pack("<B", op_val)
    elif op.dtype == idaapi.dt_word:
        chars = struct.pack("<H", op_val)
    elif op.dtype == idaapi.dt_dword:
        chars = struct.pack("<I", op_val)
    elif op.dtype == idaapi.dt_qword:
        chars = struct.pack("<Q", op_val)
    else:
        raise ValueError(f"Unhandled operand data type 0x{op.dtype:x}.")

    def is_printable_ascii(chars_: bytes):
        return all(c < 127 and chr(c) in string.printable for c in chars_)

    def is_printable_utf16le(chars_: bytes):
        if all(c == 0x00 for c in chars_[1::2]):
            return is_printable_ascii(chars_[::2])

    if is_printable_ascii(chars):
        return idaapi.get_dtype_size(op.dtype)

    if is_printable_utf16le(chars):
        return idaapi.get_dtype_size(op.dtype) // 2

    return 0


def is_mov_imm_to_stack(insn: idaapi.insn_t) -> bool:
    """verify instruction moves immediate onto stack"""
    if insn.Op2.type != idaapi.o_imm:
        return False

    if not helpers.is_op_stack_var(insn.ea, 0):
        return False