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

Code before:
    def __init__(self):
        super().__init__()
        self.global_features: List[Tuple[Feature, Address]] = []
        self.global_features.extend(capa.features.extractors.ghidra.file.extract_file_format())
        self.global_features.extend(capa.features.extractors.ghidra.global_.extract_os())
        self.global_features.extend(capa.features.extractors.ghidra.global_.extract_arch())

    def get_base_address(self):
        return AbsoluteVirtualAddress(currentProgram().getImageBase().getOffset())  # type: ignore [name-defined] # noqa: F821

    def extract_global_features(self):
        yield from self.global_features

    def extract_file_features(self):
        yield from capa.features.extractors.ghidra.file.extract_features()

    def get_functions(self) -> Iterator[FunctionHandle]:
        import capa.features.extractors.ghidra.helpers as ghidra_helpers

        yield from ghidra_helpers.get_function_symbols()

    @staticmethod
    def get_function(addr: int) -> FunctionHandle:
        get_addr = currentAddress().getAddress(hex(addr))  # type: ignore [name-defined] # noqa: F821
        func = getFunctionContaining(get_addr)  # type: ignore [name-defined] # noqa: F821
        return FunctionHandle(address=AbsoluteVirtualAddress(func.getAddress().getOffset()), inner=func)

    def extract_function_features(self, fh: FunctionHandle) -> Iterator[Tuple[Feature, Address]]:
        yield from capa.features.extractors.ghidra.function.extract_features(fh)

    def get_basic_blocks(self, fh: FunctionHandle) -> Iterator[BBHandle]:
        import capa.features.extractors.ghidra.helpers as ghidra_helpers

        yield from ghidra_helpers.get_function_blocks(fh)

    def extract_basic_block_features(self, fh: FunctionHandle, bbh: BBHandle) -> Iterator[Tuple[Feature, Address]]:
        yield from capa.features.extractors.ghidra.basicblock.extract_features(fh, bbh)

    def get_instructions(self, fh: FunctionHandle, bbh: BBHandle) -> Iterator[InsnHandle]:
        import capa.features.extractors.ghidra.helpers as ghidra_helpers

        yield from ghidra_helpers.get_insn_in_range(bbh)

    def extract_insn_features(self, fh: FunctionHandle, bbh: BBHandle, ih: InsnHandle):
        yield from capa.features.extractors.ghidra.insn.extract_features(fh, bbh, ih)


Review comment (suggestion):
```suggestion
        func = getFunctionContaining(toAddr(addr))  # type: ignore [name-defined] # noqa: F821
        return FunctionHandle(address=AbsoluteVirtualAddress(func.getAddress().getOffset()), inner=func)
```

Code after:
    def __init__(self):
        super().__init__()
        self.global_features: List[Tuple[Feature, Address]] = []
        self.global_features.extend(capa.features.extractors.ghidra.file.extract_file_format())
        self.global_features.extend(capa.features.extractors.ghidra.global_.extract_os())
        self.global_features.extend(capa.features.extractors.ghidra.global_.extract_arch())

    def get_base_address(self):
        return AbsoluteVirtualAddress(currentProgram().getImageBase().getOffset())  # type: ignore [name-defined] # noqa: F821

    def extract_global_features(self):
        yield from self.global_features

    def extract_file_features(self):
        yield from capa.features.extractors.ghidra.file.extract_features()

    def get_functions(self) -> Iterator[FunctionHandle]:
        import capa.features.extractors.ghidra.helpers as ghidra_helpers

        yield from ghidra_helpers.get_function_symbols()

    @staticmethod
    def get_function(addr: int) -> FunctionHandle:
        func = getFunctionContaining(toAddr(addr))  # type: ignore [name-defined] # noqa: F821
        return FunctionHandle(address=AbsoluteVirtualAddress(func.getAddress().getOffset()), inner=func)

    def extract_function_features(self, fh: FunctionHandle) -> Iterator[Tuple[Feature, Address]]:
        yield from capa.features.extractors.ghidra.function.extract_features(fh)

    def get_basic_blocks(self, fh: FunctionHandle) -> Iterator[BBHandle]:
        import capa.features.extractors.ghidra.helpers as ghidra_helpers

        yield from ghidra_helpers.get_function_blocks(fh)

    def extract_basic_block_features(self, fh: FunctionHandle, bbh: BBHandle) -> Iterator[Tuple[Feature, Address]]:
        yield from capa.features.extractors.ghidra.basicblock.extract_features(fh, bbh)

    def get_instructions(self, fh: FunctionHandle, bbh: BBHandle) -> Iterator[InsnHandle]:
        import capa.features.extractors.ghidra.helpers as ghidra_helpers

        yield from ghidra_helpers.get_insn_in_range(bbh)

    def extract_insn_features(self, fh: FunctionHandle, bbh: BBHandle, ih: InsnHandle):
        yield from capa.features.extractors.ghidra.insn.extract_features(fh, bbh, ih)