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

Code before:
from capa.features.extractors.base_extractor import BBHandle, InsnHandle, FunctionHandle, FeatureExtractor


class GhidraFeatureExtractor(FeatureExtractor):
    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

        imports = ghidra_helpers.get_file_imports()
        externs = ghidra_helpers.get_file_externs()
        fakes = ghidra_helpers.map_fake_import_addrs()

        for fhandle in ghidra_helpers.get_function_symbols():
            fh: FunctionHandle = FunctionHandle(
                address=fhandle.getEntryPoint().getOffset(),
                inner=fhandle,
                ctx={"imports_cache": imports, "externs_cache": externs, "fakes_cache": fakes},
            )
            yield fh

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


Review comment (refactor):
Let's move these to `__init__` to ensure the overhead is incurred only once per `GhidraFeatureExtractor` creation.

Code after:
from capa.features.extractors.base_extractor import BBHandle, InsnHandle, FunctionHandle, FeatureExtractor


class GhidraFeatureExtractor(FeatureExtractor):
    def __init__(self):
        super().__init__()
        import capa.features.extractors.ghidra.helpers as ghidra_helpers

        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())
        self.imports = ghidra_helpers.get_file_imports()
        self.externs = ghidra_helpers.get_file_externs()
        self.fakes = ghidra_helpers.map_fake_import_addrs()

    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

        for fhandle in ghidra_helpers.get_function_symbols():
            fh: FunctionHandle = FunctionHandle(
                address=fhandle.getEntryPoint().getOffset(),
                inner=fhandle,
                ctx={"imports_cache": self.imports, "externs_cache": self.externs, "fakes_cache": self.fakes},
            )
            yield fh

    @staticmethod
    def get_function(addr: int) -> FunctionHandle:
        func = getFunctionContaining(toAddr(addr))  # type: ignore [name-defined] # noqa: F821
        return FunctionHandle(address=AbsoluteVirtualAddress(func.getEntryPoint().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)