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

Code before:
    # MOV dword ptr [EBP + local_*], 0x65
    if insn.getMnemonicString().startswith("MOV"):
        found = all(insn.getOperandType(i) == mov_its_ops[i] for i in range(2))

    return found


def bb_contains_stackstring(bb: ghidra.program.model.block.CodeBlock) -> bool:
    """check basic block for stackstring indicators

    true if basic block contains enough moves of constant bytes to the stack
    """
    count = 0
    for insn in listing.getInstructions(bb, True):
        if is_mov_imm_to_stack(insn):
            count += get_printable_len(insn.getScalar(1))
        if count > MIN_STACKSTRING_LEN:
            return True
    return False


def _bb_has_tight_loop(bb: ghidra.program.model.block.CodeBlock):
    """
    parse tight loops, true if last instruction in basic block branches to bb start
    """
    last_insn = listing.getCodeUnits(bb, False).next()  # Reverse Ordered, first InstructionDB

    if last_insn.getFlowType().isJump():
        return last_insn.getAddress(0) == bb.getMinAddress()

    return False


def extract_bb_stackstring(bb: ghidra.program.model.block.CodeBlock) -> Iterator[Tuple[Feature, Address]]:
    """extract stackstring indicators from basic block"""
    if bb_contains_stackstring(bb):
        yield Characteristic("stack string"), AbsoluteVirtualAddress(bb.getMinAddress().getOffset())


def extract_bb_tight_loop(bb: ghidra.program.model.block.CodeBlock) -> Iterator[Tuple[Feature, Address]]:
    """check basic block for tight loop indicators"""
    if _bb_has_tight_loop(bb):
        yield Characteristic("tight loop"), AbsoluteVirtualAddress(bb.getMinAddress().getOffset())


BASIC_BLOCK_HANDLERS = (
    extract_bb_tight_loop,
    extract_bb_stackstring,
)



Review comment (suggestion):
I found that `getCodeUnits` worked in a similar manner. I just got the Iterator in reverse order, then pulled the first instruction. This also eliminated the possibility of having NoneTypes returned as well.

Code after:
    # MOV dword ptr [EBP + local_*], 0x65
    if insn.getMnemonicString().startswith("MOV"):
        found = all(insn.getOperandType(i) == mov_its_ops[i] for i in range(2))

    return found


def bb_contains_stackstring(bb: ghidra.program.model.block.CodeBlock) -> bool:
    """check basic block for stackstring indicators

    true if basic block contains enough moves of constant bytes to the stack
    """
    count = 0
    for insn in listing.getInstructions(bb, True):
        if is_mov_imm_to_stack(insn):
            count += get_printable_len(insn.getScalar(1))
        if count > MIN_STACKSTRING_LEN:
            return True
    return False


def _bb_has_tight_loop(bb: ghidra.program.model.block.CodeBlock):
    """
    parse tight loops, true if last instruction in basic block branches to bb start
    """
    last_insn = listing.getInstructions(bb, False).next()  # Reverse Ordered, first InstructionDB

    if last_insn.getFlowType().isJump():
        return last_insn.getAddress(0) == bb.getMinAddress()

    return False


def extract_bb_stackstring(bb: ghidra.program.model.block.CodeBlock) -> Iterator[Tuple[Feature, Address]]:
    """extract stackstring indicators from basic block"""
    if bb_contains_stackstring(bb):
        yield Characteristic("stack string"), AbsoluteVirtualAddress(bb.getMinAddress().getOffset())


def extract_bb_tight_loop(bb: ghidra.program.model.block.CodeBlock) -> Iterator[Tuple[Feature, Address]]:
    """check basic block for tight loop indicators"""
    if _bb_has_tight_loop(bb):
        yield Characteristic("tight loop"), AbsoluteVirtualAddress(bb.getMinAddress().getOffset())


BASIC_BLOCK_HANDLERS = (
    extract_bb_tight_loop,
    extract_bb_stackstring,
)