Querying base detection reasons from Python API

1 files changed, 81 insertions, 44 deletions

diff --git a/python/basedetection.py b/python/basedetection.py
index ac12c92c..e8a1ef95 100644
--- a/python/basedetection.py
+++ b/python/basedetection.py
@@ -21,15 +21,26 @@
 
 import ctypes
 from typing import Optional
+from dataclasses import dataclass
 from . import _binaryninjacore as core
 
 
+@dataclass
+class BaseAddressDetectionReason:
+    """``class BaseAddressDetectionReason`` is a class that is used to store information about why a base address is a
+    candidate"""
+
+    pointer: int
+    offset: int
+    type: int
+
+
 class BaseAddressDetection:
     """
     ``class BaseAddressDetection`` is a class that is used to detect the base address of position-dependent raw binaries
 
     >>> from binaryninja import *
-    >>> bv = load("firmware.bin")
+    >>> bv = BinaryViewType["Raw"].open("firmware.bin")
     >>> bad = BaseAddressDetection(bv)
     >>> bad.detect_base_address()
     True
@@ -51,6 +62,53 @@ class BaseAddressDetection:
         if core is not None:
             core.BNFreeBaseAddressDetection(self._handle)
 
+    @property
+    def scores(self) -> list[tuple[int, int]]:
+        """
+        ``scores`` returns a list of base addresses and their scores
+
+        :return: list of tuples containing each base address and score
+        :rtype: OrderedDict
+        """
+
+        return self._scores
+
+    @property
+    def confidence(self) -> "BaseAddressDetectionConfidenceEnum":
+        """
+        ``confidence`` returns an enum that indicates confidence that the top base address candidate is correct
+
+        :return: confidence of the base address detection results
+        :rtype: BaseAddressDetectionConfidenceEnum
+        """
+
+        return self._confidence
+
+    @property
+    def last_tested_base_address(self) -> int:
+        """
+        ``last_tested_base_address`` returns the last base address candidate that was tested
+
+        :return: last base address tested
+        :rtype: int
+        """
+
+        return self._last_tested_base_address
+
+    @property
+    def preferred_base_address(self) -> int:
+        """
+        ``preferred_base_address`` returns the base address that is preferred by analysis
+
+        :return: preferred base address
+        :rtype: int
+        """
+
+        if not self._scores:
+            return None
+
+        return self._scores[0][0]
+
     def detect_base_address(
         self,
         arch: Optional[str] = "",
@@ -85,8 +143,14 @@ class BaseAddressDetection:
             raise ValueError("upper boundary must be greater than lower boundary")
 
         settings = core.BNBaseAddressDetectionSettings(
-            arch.encode('utf-8'), analysis.encode('utf-8'), minstrlen, alignment, lowerboundary, upperboundary,
-            poi_analysis, max_pointers
+            arch.encode(),
+            analysis.encode(),
+            minstrlen,
+            alignment,
+            lowerboundary,
+            upperboundary,
+            poi_analysis,
+            max_pointers,
         )
 
         if not core.BNDetectBaseAddress(self._handle, settings):
@@ -108,49 +172,22 @@ class BaseAddressDetection:
         self._last_tested_base_address = last_base.value
         return True
 
-    @property
-    def scores(self) -> list[tuple[int, int]]:
-        """
-        ``get_scores`` returns a list of base addresses and their scores
-
-        :return: list of tuples containing each base address and score
-        :rtype: OrderedDict
-        """
-
-        return self._scores
-
-    @property
-    def confidence(self) -> "BaseAddressDetectionConfidenceEnum":
+    def get_reasons_for_base_address(self, base_address: int) -> list[BaseAddressDetectionReason]:
         """
-        ``get_confidence`` returns an enum that indicates confidence that the top base address candidate is correct
+        ``get_reasons_for_base_address`` returns a list of reasons why the specified base address is a candidate
 
-        :return: confidence of the base address detection results
-        :rtype: BaseAddressDetectionConfidenceEnum
+        :param int base_address: base address to get reasons for
+        :return: list of reasons for the specified base address
+        :rtype: list
         """
 
-        return self._confidence
+        count = ctypes.c_size_t()
+        reasons = core.BNGetBaseAddressDetectionReasons(self._handle, base_address, ctypes.byref(count))
+        if count.value == 0:
+            return []
 
-    @property
-    def last_tested_base_address(self) -> int:
-        """
-        ``last_tested_base_address`` returns the last base address candidate that was tested
-
-        :return: last base address tested
-        :rtype: int
-        """
-
-        return self._last_tested_base_address
-
-    @property
-    def preferred_base_address(self) -> int:
-        """
-        ``preferred_base_address`` returns the base address that is preferred by analysis
-
-        :return: preferred base address
-        :rtype: int
-        """
-
-        if not self._scores:
-            return None
-
-        return self._scores[0][0]
+        result = list()
+        for i in range(count.value):
+            result.append(BaseAddressDetectionReason(reasons[i].Pointer, reasons[i].POIOffset, reasons[i].POIType))
+        core.BNFreeBaseAddressDetectionReasons(reasons)
+        return result