Python bindings for custom basic block analysis

1 files changed, 323 insertions, 38 deletions

diff --git a/python/architecture.py b/python/architecture.py
index 827d4bef..099b56f1 100644
--- a/python/architecture.py
+++ b/python/architecture.py
@@ -20,7 +20,7 @@
 
 import traceback
 import ctypes
-from typing import Generator, Union, List, Optional, Mapping, Tuple, NewType, Dict
+from typing import Generator, Union, List, Optional, Mapping, Tuple, NewType, Dict, Set
 from dataclasses import dataclass, field
 
 # Binary Ninja components
@@ -39,8 +39,8 @@ from . import callingconvention
 from . import typelibrary
 from . import function
 from . import binaryview
-from . import deprecation
-from .variable import IndirectBranchInfo
+from . import variable
+from . import basicblock
 
 RegisterIndex = NewType('RegisterIndex', int)
 RegisterStackIndex = NewType('RegisterStackIndex', int)
@@ -66,14 +66,324 @@ SemanticGroupType = Union[SemanticGroupName, 'lowlevelil.ILSemanticFlagGroup', S
 IntrinsicType = Union[IntrinsicName, 'lowlevelil.ILIntrinsic', IntrinsicIndex]
 
 
-@dataclass(frozen=True)
+@dataclass
 class BasicBlockAnalysisContext:
-	"""Used by ``analyze_basic_blocks`` and contains analysis settings and other contextual information."""
-	indirect_branches: List[IndirectBranchInfo]
-	analysis_skip_override: core.FunctionAnalysisSkipOverrideEnum
-	translate_tail_calls: bool
-	disallow_branch_to_string: bool
-	max_function_size: int
+    """Used by ``analyze_basic_blocks`` and contains analysis settings and other contextual information.
+
+    .. note:: This class is meant to be used by Architecture plugins only
+    """
+
+    _handle: core.BNBasicBlockAnalysisContext
+    _function: "function.Function"
+    _contextual_returns_dirty: bool
+
+    # In
+    _indirect_branches: List["variable.IndirectBranchInfo"]
+    _indirect_no_return_calls: Set["function.ArchAndAddr"]
+    _analysis_skip_override: core.FunctionAnalysisSkipOverride
+    _translate_tail_calls: bool
+    _disallow_branch_to_string: bool
+    _max_function_size: int
+    _halt_on_invalid_instruction: bool
+    _max_size_reached: bool
+
+    # In/Out
+    _contextual_returns: Dict["function.ArchAndAddr", bool]
+
+    # Out
+    _direct_code_references: Dict[int, "function.ArchAndAddr"]
+    _direct_no_return_calls: Set["function.ArchAndAddr"]
+    _halted_disassembly_addresses: Set["function.ArchAndAddr"]
+
+    @staticmethod
+    def from_core_struct(bn_bb_context: core.BNBasicBlockAnalysisContext) -> "BasicBlockAnalysisContext":
+        """Create a BasicBlockAnalysisContext from a core.BNBasicBlockAnalysisContext structure."""
+
+        indirect_branches = []
+        for i in range(0, bn_bb_context.indirectBranchesCount):
+            ibi = variable.IndirectBranchInfo(
+                source_arch=CoreArchitecture._from_cache(bn_bb_context.indirectBranches[i].sourceArch),
+                source_addr=bn_bb_context.indirectBranches[i].sourceAddr,
+                dest_arch=CoreArchitecture._from_cache(bn_bb_context.indirectBranches[i].destArch),
+                dest_addr=bn_bb_context.indirectBranches[i].destAddr,
+                auto_defined=bn_bb_context.indirectBranches[i].autoDefined,
+            )
+            indirect_branches.append(ibi)
+
+        indirect_no_return_calls = set()
+        for i in range(0, bn_bb_context.indirectNoReturnCallsCount):
+            loc = function.ArchAndAddr(
+                CoreArchitecture._from_cache(bn_bb_context.indirectNoReturnCalls[i].arch),
+                bn_bb_context.indirectNoReturnCalls[i].address,
+            )
+            indirect_no_return_calls.add(loc)
+
+        contextual_returns = {}
+        for i in range(0, bn_bb_context.contextualFunctionReturnCount):
+            loc = function.ArchAndAddr(
+                CoreArchitecture._from_cache(bn_bb_context.contextualFunctionReturnLocations[i].arch),
+                bn_bb_context.contextualFunctionReturnLocations[i].address,
+            )
+            contextual_returns[loc] = bn_bb_context._contextualFunctionReturnValues[i]
+
+        direct_code_references = {}
+        for i in range(0, bn_bb_context.directRefCount):
+            src = function.ArchAndAddr(
+                CoreArchitecture._from_cache(bn_bb_context.directRefSources[i].arch),
+                bn_bb_context.directRefSources[i].address,
+            )
+            direct_code_references[bn_bb_context.directRefTargets[i]] = src
+
+        direct_no_return_calls = set()
+        for i in range(0, bn_bb_context.directNoReturnCallsCount):
+            loc = function.ArchAndAddr(
+                CoreArchitecture._from_cache(bn_bb_context.directNoReturnCallLocations[i].arch),
+                bn_bb_context.directNoReturnCallLocations[i].address,
+            )
+            direct_no_return_calls.add(loc)
+
+        halted_disassembly_addresses = set()
+        for i in range(0, bn_bb_context.haltedDisassemblyAddressesCount):
+            addr = function.ArchAndAddr(
+                CoreArchitecture._from_cache(bn_bb_context.haltedDisassemblyAddresses[i].arch),
+                bn_bb_context.haltedDisassemblyAddresses[i].address,
+            )
+            halted_disassembly_addresses.add(addr)
+
+        view = binaryview.BinaryView(handle=core.BNGetFunctionData(bn_bb_context.function))
+        return BasicBlockAnalysisContext(
+            _handle=bn_bb_context,
+            _function=function.Function(view, core.BNNewFunctionReference(bn_bb_context.function)),
+            _indirect_branches=indirect_branches,
+            _indirect_no_return_calls=indirect_no_return_calls,
+            _analysis_skip_override=bn_bb_context.analysisSkipOverride,
+            _translate_tail_calls=bn_bb_context.translateTailCalls,
+            _disallow_branch_to_string=bn_bb_context.disallowBranchToString,
+            _max_function_size=bn_bb_context.maxFunctionSize,
+            _halt_on_invalid_instruction=bn_bb_context.haltOnInvalidInstructions,
+            _max_size_reached=bn_bb_context.maxSizeReached,
+            _contextual_returns=contextual_returns,
+            _contextual_returns_dirty=False,
+            _direct_code_references=direct_code_references,
+            _direct_no_return_calls=direct_no_return_calls,
+            _halted_disassembly_addresses=halted_disassembly_addresses,
+        )
+
+    @property
+    def indirect_branches(self) -> List["variable.IndirectBranchInfo"]:
+        """Get the list of indirect branches in this context."""
+
+        return self._indirect_branches
+
+    @property
+    def indirect_no_return_calls(self) -> Set["function.ArchAndAddr"]:
+        """Get the set of indirect no-return calls in this context."""
+
+        return self._indirect_no_return_calls
+
+    @property
+    def analysis_skip_override(self) -> core.FunctionAnalysisSkipOverride:
+        """Get the analysis skip override setting for this context."""
+
+        return self._analysis_skip_override
+
+    @property
+    def translate_tail_calls(self) -> bool:
+        """Get setting from context that determines if tail calls should be translated."""
+
+        return self._translate_tail_calls
+
+    @property
+    def disallow_branch_to_string(self) -> bool:
+        """Get setting from context that determines if branches to string addresses should be disallowed."""
+
+        return self._disallow_branch_to_string
+
+    @property
+    def max_function_size(self) -> int:
+        """Get the maximum function size setting for this context."""
+
+        return self._max_function_size
+
+    @property
+    def halt_on_invalid_instruction(self) -> bool:
+        """Get the setting from context that determines if analysis should halt on invalid instructions."""
+
+        return self._halt_on_invalid_instruction
+
+    @property
+    def max_size_reached(self) -> bool:
+        """Get boolean that indicates if the maximum function size has been reached."""
+
+        return self._max_size_reached
+
+    @property
+    def contextual_returns(self) -> Dict["function.ArchAndAddr", bool]:
+        """Get the mapping of contextual function return locations to their values."""
+
+        return self._contextual_returns
+
+    def add_contextual_return(self, loc: "function.ArchAndAddr", value: bool) -> None:
+        """
+        ``add_contextual_return`` adds a contextual function return location and its value to the current function.
+
+        :param function.ArchAndAddr loc: The location of the contextual function return
+        :param bool value: The value of the contextual function return
+        """
+        if not isinstance(value, bool):
+            raise TypeError("value must be a boolean")
+
+        if not isinstance(loc, function.ArchAndAddr):
+            raise TypeError("loc must be an instance of function.ArchAndAddr")
+
+        # Update existing value if it exists
+        if loc in self._contextual_returns:
+            if self._contextual_returns[loc] == value:
+                return
+
+        self._contextual_returns[loc] = value
+        self._contextual_returns_dirty = True
+
+    @property
+    def direct_code_references(self) -> Dict[int, "function.ArchAndAddr"]:
+        """Get the mapping of direct code reference targets to their source locations."""
+
+        return self._direct_code_references
+
+    def add_direct_code_reference(self, target: int, source: "function.ArchAndAddr") -> None:
+        """
+        ``add_direct_code_reference`` adds a direct code reference to the current function.
+
+        :param int target: The target address of the direct code reference
+        :param function.ArchAndAddr source: The source location of the direct code reference
+        """
+
+        if not isinstance(target, int):
+            raise TypeError("target must be an integer")
+
+        if not isinstance(source, function.ArchAndAddr):
+            raise TypeError("source must be an instance of function.ArchAndAddr")
+
+        self._direct_code_references[target] = source
+
+    @property
+    def direct_no_return_calls(self) -> Set["function.ArchAndAddr"]:
+        """Get the set of direct no-return call locations in this context."""
+
+        return self._direct_no_return_calls
+
+    def add_direct_no_return_call(self, loc: "function.ArchAndAddr") -> None:
+        """
+        ``add_direct_no_return_call`` adds a direct no-return call location to the current function.
+
+        :param function.ArchAndAddr loc: The location of the direct no-return call
+        """
+        if not isinstance(loc, function.ArchAndAddr):
+            raise TypeError("loc must be an instance of function.ArchAndAddr")
+
+        self._direct_no_return_calls.add(loc)
+
+    @property
+    def halted_disassembly_addresses(self) -> Set["function.ArchAndAddr"]:
+        """Get the set of addresses where disassembly has been halted."""
+
+        return self._halted_disassembly_addresses
+
+    def add_halted_disassembly_address(self, loc: "function.ArchAndAddr") -> None:
+        """
+        ``add_halted_disassembly_address`` adds an address to the set of halted disassembly addresses.
+
+        :param function.ArchAndAddr loc: The location of the halted disassembly address
+        """
+        if not isinstance(loc, function.ArchAndAddr):
+            raise TypeError("loc must be an instance of function.ArchAndAddr")
+
+        self._halted_disassembly_addresses.add(loc)
+
+    def create_basic_block(self, arch: "Architecture", start: int) -> Optional["basicblock.BasicBlock"]:
+        """
+        ``create_basic_block`` creates a new BasicBlock at the specified address for the given Architecture.
+
+        :param Architecture arch: Architecture of the BasicBlock to create
+        :param int start: Address of the BasicBlock to create
+        """
+
+        if not isinstance(arch, Architecture):
+            raise TypeError("arch must be an instance of architecture.Architecture")
+
+        bnblock = core.BNAnalyzeBasicBlocksContextCreateBasicBlock(self._handle, arch.handle, start)
+        if not bnblock:
+            return None
+
+        view = binaryview.BinaryView(handle=core.BNGetFunctionData(self._function.handle))
+        return basicblock.BasicBlock(bnblock, view)
+
+    def add_basic_block(self, block: "basicblock.BasicBlock") -> None:
+        """
+        ``add_basic_block`` adds a BasicBlock to the current function.
+
+        :param basicblock.BasicBlock block: The BasicBlock to add
+        """
+        if not isinstance(block, basicblock.BasicBlock):
+            raise TypeError("block must be an instance of basicblock.BasicBlock")
+
+        core.BNAnalyzeBasicBlocksContextAddBasicBlockToFunction(self._handle, block.handle)
+
+    def add_temp_outgoing_reference(self, target: "function.Function") -> None:
+        """
+        ``add_temp_outgoing_reference`` adds a temporary outgoing reference to the specified function.
+
+        :param function.Function target: The target function to add a temporary outgoing reference to
+        """
+        if not isinstance(target, function.Function):
+            raise TypeError("target must be an instance of function.Function")
+
+        core.BNAnalyzeBasicBlocksContextAddTempReference(self._handle, target.handle)
+
+    def finalize(self) -> None:
+        """
+        ``finalize`` finalizes the function's basic block analysis
+        """
+
+        if self._direct_code_references:
+            total = len(self._direct_code_references)
+            sources = (core.BNArchitectureAndAddress * total)()
+            targets = (ctypes.c_ulonglong * total)()
+            for i, (target, src) in enumerate(self._direct_code_references.items()):
+                sources[i].arch = src.arch.handle
+                sources[i].address = src.addr
+                targets[i] = target
+
+            core.BNAnalyzeBasicBlocksContextSetDirectCodeReferences(self._handle, sources, targets, total)
+
+        if self._direct_no_return_calls:
+            total = len(self._direct_no_return_calls)
+            direct_no_return_calls = (core.BNArchitectureAndAddress * total)()
+            for i, loc in enumerate(self._direct_no_return_calls):
+                direct_no_return_calls[i].arch = loc.arch.handle
+                direct_no_return_calls[i].address = loc.addr
+            core.BNAnalyzeBasicBlocksContextSetDirectNoReturnCalls(self._handle, direct_no_return_calls, total)
+
+        self._halted_disassembly_addresses.add(function.ArchAndAddr(self._function.arch, 0))
+        if self._halted_disassembly_addresses:
+            total = len(self._halted_disassembly_addresses)
+            halted_addresses = (core.BNArchitectureAndAddress * total)()
+            for i, loc in enumerate(self._halted_disassembly_addresses):
+                halted_addresses[i].arch = loc.arch.handle
+                halted_addresses[i].address = loc.addr
+            core.BNAnalyzeBasicBlocksContextSetHaltedDisassemblyAddresses(self._handle, halted_addresses, total)
+
+        if self._contextual_returns_dirty:
+            total = len(self._contextual_returns)
+            values = (ctypes.c_bool * total)()
+            returns = (core.BNArchitectureAndAddress * total)()
+            for i, (loc, value) in enumerate(self._contextual_returns.items()):
+                returns[i].arch = loc.arch.handle
+                returns[i].address = loc.addr
+                values[i] = value
+            core.BNAnalyzeBasicBlocksContextSetContextualFunctionReturns(self._handle, returns, values, total)
+
+        core.BNAnalyzeBasicBlocksContextFinalize(self._handle)
 
 
 @dataclass(frozen=True)
@@ -725,20 +1035,8 @@ class Architecture(metaclass=_ArchitectureMetaClass):
 
 	def _analyze_basic_blocks(self, ctx, func, ptr_bn_bb_context):
 		try:
-			return core.BNArchitectureDefaultAnalyzeBasicBlocks(func, ptr_bn_bb_context)
 			bn_bb_context = ptr_bn_bb_context.contents
-			indirect_branches = []
-			for i in range(0, bn_bb_context.indirectBranchesCount):
-				ibi = IndirectBranchInfo()
-				ibi.source_arch = CoreArchitecture._from_cache(bn_bb_context.indirectBranches[i].sourceArch)
-				ibi.source_addr = bn_bb_context.indirectBranches[i].sourceAddr
-				ibi.dest_arch = CoreArchitecture._from_cache(bn_bb_context.indirectBranches[i].destArch)
-				ibi.dest_addr = bn_bb_context.indirectBranches[i].destAddr
-				ibi.auto_defined = bn_bb_context.indirectBranches[i].autoDefined
-				indirect_branches.append(ibi)
-
-			context = BasicBlockAnalysisContext(indirect_branches, bn_bb_context.analysisSkipOverride,
-				bn_bb_context.translateTailCalls, bn_bb_context.disallowBranchToString, bn_bb_context.maxFunctionSize)
+			context = BasicBlockAnalysisContext.from_core_struct(bn_bb_context)
 			self.analyze_basic_blocks(function.Function(handle=core.BNNewFunctionReference(func)), context)
 		except:
 			log_error(traceback.format_exc())
@@ -1465,24 +1763,11 @@ class Architecture(metaclass=_ArchitectureMetaClass):
 		.. note:: Architecture subclasses should only implement this method if function-level analysis is required
 
 		:param Function func: the function to analyze
-		:param BNBasicBlockAnalysisContext context: the analysis context
+		:param BasicBlockAnalysisContext context: the analysis context
 		"""
 
 		try:
-			bn_bb_context = core.BNBasicBlockAnalysisContext()
-			bn_bb_context.indirectBranchesCount = len(context.indirect_branches)
-			bn_bb_context.analysisSkipOverride = context.analysis_skip_override
-			bn_bb_context.translateTailCalls = context.translate_tail_calls
-			bn_bb_context.disallowBranchToString = context.disallow_branch_to_string
-			bn_bb_context.maxFunctionSize = context.max_function_size
-			bn_bb_context.indirectBranches = (core.BNIndirectBranchInfo * len(context.indirect_branches))()
-			for i in range(0, len(context.indirect_branches)):
-				bn_bb_context.indirectBranches[i].sourceArch = context.indirect_branches[i].source_arch.handle
-				bn_bb_context.indirectBranches[i].sourceAddr = context.indirect_branches[i].source_addr
-				bn_bb_context.indirectBranches[i].destArch = context.indirect_branches[i].dest_arch.handle
-				bn_bb_context.indirectBranches[i].destAddr = context.indirect_branches[i].dest_addr
-				bn_bb_context.indirectBranches[i].autoDefined = context.indirect_branches[i].auto_defined
-			core.BNArchitectureDefaultAnalyzeBasicBlocks(func.handle, ctypes.byref(bn_bb_context))
+			core.BNArchitectureDefaultAnalyzeBasicBlocks(func.handle, context._handle)
 		except:
 			log_error(traceback.format_exc())