binary ninja: optimize the order of computing LLIL to partially address #2402

capa/features/extractors/binja/extractor.py

@@ -5,15 +5,19 @@
 # Unless required by applicable law or agreed to in writing, software distributed under the License
 #  is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and limitations under the License.
+import logging
 from typing import Iterator
+from collections import defaultdict
 
 import binaryninja as binja
-from binaryninja import ILException
+from binaryninja import Function, BinaryView, SymbolType, ILException, RegisterValueType, LowLevelILOperation
 
+import capa.perf
 import capa.features.extractors.elf
 import capa.features.extractors.binja.file
 import capa.features.extractors.binja.insn
 import capa.features.extractors.binja.global_
+import capa.features.extractors.binja.helpers
 import capa.features.extractors.binja.function
 import capa.features.extractors.binja.basicblock
 from capa.features.common import Feature
@@ -26,6 +30,8 @@
     StaticFeatureExtractor,
 )
 
+logger = logging.getLogger(__name__)
+
 
 class BinjaFeatureExtractor(StaticFeatureExtractor):
     def __init__(self, bv: binja.BinaryView):
@@ -36,6 +42,9 @@ def __init__(self, bv: binja.BinaryView):
         self.global_features.extend(capa.features.extractors.binja.global_.extract_os(self.bv))
         self.global_features.extend(capa.features.extractors.binja.global_.extract_arch(self.bv))
 
+        with capa.perf.timing("binary ninja: computing call graph"):
+            self._call_graph = self._build_call_graph()
+
     def get_base_address(self):
         return AbsoluteVirtualAddress(self.bv.start)
 
@@ -45,9 +54,27 @@ def extract_global_features(self):
     def extract_file_features(self):
         yield from capa.features.extractors.binja.file.extract_features(self.bv)
 
+    def _build_call_graph(self):
+        # from function address to function addresses
+        calls_from: defaultdict[int, set[int]] = defaultdict(set)
+        calls_to: defaultdict[int, set[int]] = defaultdict(set)
+
+        f: Function
+        for f in self.bv.functions:
+            for caller in f.callers:
+                calls_from[caller.start].add(f.start)
+                calls_to[f.start].add(caller.start)
+
+        call_graph = {
+            "calls_to": calls_to,
+            "calls_from": calls_from,
+        }
+
+        return call_graph
+
     def get_functions(self) -> Iterator[FunctionHandle]:
         for f in self.bv.functions:
-            yield FunctionHandle(address=AbsoluteVirtualAddress(f.start), inner=f)
+            yield FunctionHandle(address=AbsoluteVirtualAddress(f.start), inner=f, ctx={"call_graph": self._call_graph})
 
     def extract_function_features(self, fh: FunctionHandle) -> Iterator[tuple[Feature, Address]]:
         yield from capa.features.extractors.binja.function.extract_features(fh)
@@ -76,13 +103,16 @@ def extract_basic_block_features(self, fh: FunctionHandle, bbh: BBHandle) -> Ite
         yield from capa.features.extractors.binja.basicblock.extract_features(fh, bbh)
 
     def get_instructions(self, fh: FunctionHandle, bbh: BBHandle) -> Iterator[InsnHandle]:
-        import capa.features.extractors.binja.helpers as binja_helpers
+        f: binja.Function = fh.inner
 
-        bb: tuple[binja.BasicBlock, binja.MediumLevelILBasicBlock] = bbh.inner
-        addr = bb[0].start
+        bb: binja.BasicBlock
+        mlbb: binja.MediumLevelILBasicBlock
+        bb, mlbb = bbh.inner
 
-        for text, length in bb[0]:
-            insn = binja_helpers.DisassemblyInstruction(addr, length, text)
+        addr: int = bb.start
+        for text, length in bb:
+            llil = f.get_llils_at(addr)
+            insn = capa.features.extractors.binja.helpers.DisassemblyInstruction(addr, length, text, llil)
             yield InsnHandle(address=AbsoluteVirtualAddress(addr), inner=insn)
             addr += length
 

capa/features/extractors/binja/function.py

@@ -7,7 +7,7 @@
 # See the License for the specific language governing permissions and limitations under the License.
 from typing import Iterator
 
-from binaryninja import Function, BinaryView, SymbolType, ILException, RegisterValueType, LowLevelILOperation
+from binaryninja import Function, BinaryView, SymbolType
 
 from capa.features.file import FunctionName
 from capa.features.common import Feature, Characteristic
@@ -20,38 +20,24 @@ def extract_function_calls_to(fh: FunctionHandle):
     """extract callers to a function"""
     func: Function = fh.inner
 
-    for caller in func.caller_sites:
-        # Everything that is a code reference to the current function is considered a caller, which actually includes
-        # many other references that are NOT a caller. For example, an instruction `push function_start` will also be
-        # considered a caller to the function
-        llil = None
-        try:
-            # Temporary fix for https://github.com/Vector35/binaryninja-api/issues/6020. Since `.llil` can throw an
-            # exception rather than returning None
-            llil = caller.llil
-        except ILException:
+    caller: int
+    for caller in fh.ctx["call_graph"]["calls_to"].get(func.start, []):
+        if caller == func.start:
             continue
 
-        if (llil is None) or llil.operation not in [
-            LowLevelILOperation.LLIL_CALL,
-            LowLevelILOperation.LLIL_CALL_STACK_ADJUST,
-            LowLevelILOperation.LLIL_JUMP,
-            LowLevelILOperation.LLIL_TAILCALL,
-        ]:
-            continue
+        yield Characteristic("calls to"), AbsoluteVirtualAddress(caller)
 
-        if llil.dest.value.type not in [
-            RegisterValueType.ImportedAddressValue,
-            RegisterValueType.ConstantValue,
-            RegisterValueType.ConstantPointerValue,
-        ]:
-            continue
 
-        address = llil.dest.value.value
-        if address != func.start:
+def extract_function_calls_from(fh: FunctionHandle):
+    """extract callers from a function"""
+    func: Function = fh.inner
+
+    callee: int
+    for callee in fh.ctx["call_graph"]["calls_from"].get(func.start, []):
+        if callee == func.start:
             continue
 
-        yield Characteristic("calls to"), AbsoluteVirtualAddress(caller.address)
+        yield Characteristic("calls from"), AbsoluteVirtualAddress(callee)
 
 
 def extract_function_loop(fh: FunctionHandle):
@@ -72,13 +58,12 @@ def extract_function_loop(fh: FunctionHandle):
 def extract_recursive_call(fh: FunctionHandle):
     """extract recursive function call"""
     func: Function = fh.inner
-    bv: BinaryView = func.view
-    if bv is None:
-        return
 
-    for ref in bv.get_code_refs(func.start):
-        if ref.function == func:
+    caller: int
+    for caller in fh.ctx["call_graph"]["calls_to"].get(func.start, []):
+        if caller == func.start:
             yield Characteristic("recursive call"), fh.address
+            return
 
 
 def extract_function_name(fh: FunctionHandle):
@@ -108,4 +93,10 @@ def extract_features(fh: FunctionHandle) -> Iterator[tuple[Feature, Address]]:
             yield feature, addr
 
 
-FUNCTION_HANDLERS = (extract_function_calls_to, extract_function_loop, extract_recursive_call, extract_function_name)
+FUNCTION_HANDLERS = (
+    extract_function_calls_to,
+    extract_function_calls_from,
+    extract_function_loop,
+    extract_recursive_call,
+    extract_function_name,
+)

capa/features/extractors/binja/helpers.py

@@ -9,7 +9,7 @@
 from typing import Callable
 from dataclasses import dataclass
 
-from binaryninja import BinaryView, LowLevelILInstruction
+from binaryninja import BinaryView, LowLevelILOperation, LowLevelILInstruction
 from binaryninja.architecture import InstructionTextToken
 
 
@@ -18,6 +18,24 @@ class DisassemblyInstruction:
     address: int
     length: int
     text: list[InstructionTextToken]
+    llil: list[LowLevelILInstruction]
+
+    @property
+    def is_call(self):
+        if not self.llil:
+            return False
+
+        # TODO(williballenthin): when to use one vs many llil instructions
+        # https://github.com/Vector35/binaryninja-api/issues/6205
+        llil = self.llil[0]
+        if not llil:
+            return False
+
+        return llil.operation in [
+            LowLevelILOperation.LLIL_CALL,
+            LowLevelILOperation.LLIL_CALL_STACK_ADJUST,
+            LowLevelILOperation.LLIL_TAILCALL,
+        ]
 
 
 LLIL_VISITOR = Callable[[LowLevelILInstruction, LowLevelILInstruction, int], bool]