diff --git a/capa/features/extractors/binexport2/__init__.py b/capa/features/extractors/binexport2/__init__.py index ba488cd86..a09ca2a52 100644 --- a/capa/features/extractors/binexport2/__init__.py +++ b/capa/features/extractors/binexport2/__init__.py @@ -261,7 +261,9 @@ def _find_base_address(self): def _compute_thunks(self): for addr, idx in self.idx.vertex_index_by_address.items(): vertex: BinExport2.CallGraph.Vertex = self.be2.call_graph.vertex[idx] - if not capa.features.extractors.binexport2.helpers.is_vertex_type(vertex, BinExport2.CallGraph.Vertex.Type.THUNK): + if not capa.features.extractors.binexport2.helpers.is_vertex_type( + vertex, BinExport2.CallGraph.Vertex.Type.THUNK + ): continue curr_idx: int = idx diff --git a/capa/features/extractors/binexport2/basicblock.py b/capa/features/extractors/binexport2/basicblock.py index 5d7398aa1..bcb7977b4 100644 --- a/capa/features/extractors/binexport2/basicblock.py +++ b/capa/features/extractors/binexport2/basicblock.py @@ -13,6 +13,7 @@ from capa.features.basicblock import BasicBlock from capa.features.extractors.binexport2 import FunctionContext, BasicBlockContext from capa.features.extractors.base_extractor import BBHandle, FunctionHandle +from capa.features.extractors.binexport2.binexport2_pb2 import BinExport2 def extract_bb_tight_loop(fh: FunctionHandle, bbh: BBHandle) -> Iterator[Tuple[Feature, Address]]: diff --git a/capa/features/extractors/binexport2/extractor.py b/capa/features/extractors/binexport2/extractor.py index 6e1e4c633..1c3c4d393 100644 --- a/capa/features/extractors/binexport2/extractor.py +++ b/capa/features/extractors/binexport2/extractor.py @@ -75,7 +75,9 @@ def get_functions(self) -> Iterator[FunctionHandle]: be2_vertex: BinExport2.CallGraph.Vertex = self.be2.call_graph.vertex[vertex_idx] # skip thunks - if capa.features.extractors.binexport2.helpers.is_vertex_type(be2_vertex, BinExport2.CallGraph.Vertex.Type.THUNK): + if capa.features.extractors.binexport2.helpers.is_vertex_type( + be2_vertex, BinExport2.CallGraph.Vertex.Type.THUNK + ): continue yield FunctionHandle( diff --git a/capa/features/extractors/binexport2/insn.py b/capa/features/extractors/binexport2/insn.py index 1a907eb7b..92ef311ee 100644 --- a/capa/features/extractors/binexport2/insn.py +++ b/capa/features/extractors/binexport2/insn.py @@ -54,7 +54,9 @@ def extract_insn_api_features(fh: FunctionHandle, _bbh: BBHandle, ih: InsnHandle vertex_idx: int = be2_index.vertex_index_by_address[addr] vertex: BinExport2.CallGraph.Vertex = be2.call_graph.vertex[vertex_idx] - if not capa.features.extractors.binexport2.helpers.is_vertex_type(vertex, BinExport2.CallGraph.Vertex.Type.IMPORTED): + if not capa.features.extractors.binexport2.helpers.is_vertex_type( + vertex, BinExport2.CallGraph.Vertex.Type.IMPORTED + ): continue if not vertex.HasField("mangled_name"):