Firmware Ninja reference tree API

4 files changed, 491 insertions, 72 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 8b41ea76..2a2e5e9b 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -18397,6 +18397,59 @@ namespace BinaryNinja {
 		size_t unique;
 	};
 
+
+	/*! FirmwareNinjaReferenceNode is a class used to build reference trees to memory regions, functions, and data
+		variables. This class is only available in the Ultimate Edition of Binary Ninja.
+
+		\ingroup firmwareninja
+	*/
+	class FirmwareNinjaReferenceNode : public CoreRefCountObject<BNFirmwareNinjaReferenceNode, BNNewFirmwareNinjaReferenceNodeReference, BNFreeFirmwareNinjaReferenceNode>
+	{
+		BNFirmwareNinjaReferenceNode* m_object;
+	public:
+		FirmwareNinjaReferenceNode(BNFirmwareNinjaReferenceNode* node);
+		~FirmwareNinjaReferenceNode();
+
+		/*! Determine if the reference tree node is for a function
+
+			\return true if the reference tree node is for a function, false otherwise
+		 */
+		bool IsFunction();
+
+		/*! Determine if the reference tree node is for a data variable
+
+			\return true if the reference tree node is for a data variable, false otherwise
+		 */
+		bool IsDataVariable();
+
+		/*! Determine if the reference tree node contains child nodes
+
+			\return true if the reference tree node contains child nodes, false otherwise
+		 */
+		bool HasChildren();
+
+		/*! Query the function contained in the reference tree node
+
+			\param function Output function object
+			\return true if the function was queried successfully, false otherwise
+		 */
+		bool GetFunction(Ref<Function>& function);
+
+		/*! Query the data variable contained in the reference tree node
+
+			\param function Output data variable object
+			\return true if the data variable was queried successfully, false otherwise
+		 */
+		bool GetDataVariable(DataVariable& variable);
+
+		/*! Query the child nodes contained in the reference tree node
+
+			\return Vector of child reference tree nodes
+		 */
+		std::vector<Ref<FirmwareNinjaReferenceNode>> GetChildren();
+	};
+
+
 	/*! FirmwareNinja is a class containing features specific to embedded firmware analysis. This class is only
 		available in the Ultimate Edition of Binary Ninja.
 
@@ -18485,6 +18538,47 @@ namespace BinaryNinja {
 		 */
 		std::vector<FirmwareNinjaDeviceAccesses> GetBoardDeviceAccesses(
 			const std::vector<FirmwareNinjaFunctionMemoryAccesses>& fma);
+
+
+		/*! Returns a tree of reference nodes that reference the memory region represented by the given device
+
+			\param device Firmware Ninja device
+			\param fma Vector of Firmware Ninja function memory accesses information
+			\param value (Optional) only include components that originate with a write of this value to the device
+			\return Root reference node of tree
+		 */
+		Ref<FirmwareNinjaReferenceNode> GetReferenceTree(
+			FirmwareNinjaDevice& device,
+			const std::vector<FirmwareNinjaFunctionMemoryAccesses>& fma,
+			uint64_t* value = nullptr
+		);
+
+		/*! Returns a tree of reference nodes that reference the memory region represented by the given section
+
+			\param device Firmware Ninja device
+			\param fma Vector of Firmware Ninja function memory accesses information
+			\param value (Optional) only include components that originate with a write of this value to the device
+			\return Root reference node of tree
+		 */
+		Ref<FirmwareNinjaReferenceNode> GetReferenceTree(
+			Section& section,
+			const std::vector<FirmwareNinjaFunctionMemoryAccesses>& fma,
+			uint64_t* value = nullptr
+		);
+
+
+		/*! Returns a tree of reference nodes that reference the given address
+
+			\param device Firmware Ninja device
+			\param fma Vector of Firmware Ninja function memory accesses information
+			\param value (Optional) only include components that originate with a write of this value to the device
+			\return Root reference node of tree
+		 */
+		Ref<FirmwareNinjaReferenceNode> GetReferenceTree(
+			uint64_t address,
+			const std::vector<FirmwareNinjaFunctionMemoryAccesses>& fma,
+			uint64_t* value = nullptr
+		);
 	};
 
 
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 2fe08614..b13a7a7a 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -37,7 +37,7 @@
 // Current ABI version for linking to the core. This is incremented any time
 // there are changes to the API that affect linking, including new functions,
 // new types, or modifications to existing functions or types.
-#define BN_CURRENT_CORE_ABI_VERSION 87
+#define BN_CURRENT_CORE_ABI_VERSION 88
 
 // Minimum ABI version that is supported for loading of plugins. Plugins that
 // are linked to an ABI version less than this will not be able to load and
@@ -301,6 +301,7 @@ extern "C"
 	typedef struct BNUndoEntry BNUndoEntry;
 	typedef struct BNDemangler BNDemangler;
 	typedef struct BNFirmwareNinja BNFirmwareNinja;
+	typedef struct BNFirmwareNinjaReferenceNode BNFirmwareNinjaReferenceNode;
 
 	//! Console log levels
 	typedef enum BNLogLevel
@@ -3537,7 +3538,6 @@ extern "C"
 		size_t unique;
 	} BNFirmwareNinjaDeviceAccesses;
 
-
 	BINARYNINJACOREAPI char* BNAllocString(const char* contents);
 	BINARYNINJACOREAPI char* BNAllocStringWithLength(const char* contents, size_t len);
 	BINARYNINJACOREAPI void BNFreeString(char* str);
@@ -8038,12 +8038,24 @@ extern "C"
 	BINARYNINJACOREAPI int BNFirmwareNinjaQueryBoardDevices(BNFirmwareNinja* fn, BNArchitecture* arch, const char* board, BNFirmwareNinjaDevice** devices);
 	BINARYNINJACOREAPI int BNFirmwareNinjaFindSectionsWithEntropy(BNFirmwareNinja* fn, BNFirmwareNinjaSection** sections, float highCodeEntropyThreshold, float lowCodeEntropyThreshold, size_t blockSize, BNFirmwareNinjaSectionAnalysisMode mode);
 	BINARYNINJACOREAPI void BNFirmwareNinjaFreeSections(BNFirmwareNinjaSection *sections, int size);
-	BINARYNINJACOREAPI int BNFirmwareNinjaGetFunctionMemoryAccesses(BNFirmwareNinja* fn, BNFirmwareNinjaFunctionMemoryAccesses*** mmio, BNProgressFunction progress, void* progressContext);
-	BINARYNINJACOREAPI void BNFirmwareNinjaFreeFunctionMemoryAccesses(BNFirmwareNinjaFunctionMemoryAccesses **mmio, int size);
-	BINARYNINJACOREAPI void BNFirmwareNinjaStoreFunctionMemoryAccessesToMetadata(BNFirmwareNinja* fn, BNFirmwareNinjaFunctionMemoryAccesses** mmio, int size);
-	BINARYNINJACOREAPI int BNFirmwareNinjaQueryFunctionMemoryAccessesFromMetadata(BNFirmwareNinja* fn, BNFirmwareNinjaFunctionMemoryAccesses*** mmio);
-	BINARYNINJACOREAPI int BNFirmwareNinjaGetBoardDeviceAccesses(BNFirmwareNinja* fn, BNFirmwareNinjaFunctionMemoryAccesses** mmio, int size, BNFirmwareNinjaDeviceAccesses** accesses, BNArchitecture* arch);
+	BINARYNINJACOREAPI int BNFirmwareNinjaGetFunctionMemoryAccesses(BNFirmwareNinja* fn, BNFirmwareNinjaFunctionMemoryAccesses*** fma, BNProgressFunction progress, void* progressContext);
+	BINARYNINJACOREAPI void BNFirmwareNinjaFreeFunctionMemoryAccesses(BNFirmwareNinjaFunctionMemoryAccesses **fma, int size);
+	BINARYNINJACOREAPI void BNFirmwareNinjaStoreFunctionMemoryAccessesToMetadata(BNFirmwareNinja* fn, BNFirmwareNinjaFunctionMemoryAccesses** fma, int size);
+	BINARYNINJACOREAPI int BNFirmwareNinjaQueryFunctionMemoryAccessesFromMetadata(BNFirmwareNinja* fn, BNFirmwareNinjaFunctionMemoryAccesses*** fma);
+	BINARYNINJACOREAPI int BNFirmwareNinjaGetBoardDeviceAccesses(BNFirmwareNinja* fn, BNFirmwareNinjaFunctionMemoryAccesses** fma, int size, BNFirmwareNinjaDeviceAccesses** accesses, BNArchitecture* arch);
 	BINARYNINJACOREAPI void BNFirmwareNinjaFreeBoardDeviceAccesses(BNFirmwareNinjaDeviceAccesses *accesses, int size);
+	BINARYNINJACOREAPI BNFirmwareNinjaReferenceNode* BNFirmwareNinjaGetMemoryRegionReferenceTree(BNFirmwareNinja* fn, uint64_t start, uint64_t end, BNFirmwareNinjaFunctionMemoryAccesses** fma, int size, uint64_t* value);
+	BINARYNINJACOREAPI BNFirmwareNinjaReferenceNode* BNFirmwareNinjaGetAddressReferenceTree(BNFirmwareNinja* fn, uint64_t address, BNFirmwareNinjaFunctionMemoryAccesses** fma, int size, uint64_t* value);
+
+	BINARYNINJACOREAPI bool BNFirmwareNinjaReferenceNodeIsFunction(BNFirmwareNinjaReferenceNode* node);
+	BINARYNINJACOREAPI bool BNFirmwareNinjaReferenceNodeIsDataVariable(BNFirmwareNinjaReferenceNode* node);
+	BINARYNINJACOREAPI bool BNFirmwareNinjaReferenceNodeHasChildren(BNFirmwareNinjaReferenceNode* node);
+	BINARYNINJACOREAPI BNFunction* BNFirmwareNinjaReferenceNodeGetFunction(BNFirmwareNinjaReferenceNode* node);
+	BINARYNINJACOREAPI BNDataVariable* BNFirmwareNinjaReferenceNodeGetDataVariable(BNFirmwareNinjaReferenceNode* node);
+	BINARYNINJACOREAPI BNFirmwareNinjaReferenceNode** BNFirmwareNinjaReferenceNodeGetChildren(BNFirmwareNinjaReferenceNode* parent, size_t* count);
+	BINARYNINJACOREAPI void BNFreeFirmwareNinjaReferenceNode(BNFirmwareNinjaReferenceNode* node);
+	BINARYNINJACOREAPI BNFirmwareNinjaReferenceNode* BNNewFirmwareNinjaReferenceNodeReference(BNFirmwareNinjaReferenceNode* node);
+	BINARYNINJACOREAPI void BNFreeFirmwareNinjaReferenceNodes(BNFirmwareNinjaReferenceNode** nodes, size_t count);
 #ifdef __cplusplus
 }
 #endif
diff --git a/firmwareninja.cpp b/firmwareninja.cpp
index 5bf623bd..8e13fccd 100644
--- a/firmwareninja.cpp
+++ b/firmwareninja.cpp
@@ -58,10 +58,83 @@ static void FreeMemoryInfoArray(BNFirmwareNinjaFunctionMemoryAccesses** fma, siz
 }
 
 
+FirmwareNinjaReferenceNode::FirmwareNinjaReferenceNode(BNFirmwareNinjaReferenceNode* node)
+{
+	m_object = node;
+}
+
+
+FirmwareNinjaReferenceNode::~FirmwareNinjaReferenceNode()
+{
+	BNFreeFirmwareNinjaReferenceNode(m_object);
+}
+
+
+bool FirmwareNinjaReferenceNode::IsFunction()
+{
+	return BNFirmwareNinjaReferenceNodeIsFunction(m_object);
+}
+
+
+bool FirmwareNinjaReferenceNode::IsDataVariable()
+{
+	return BNFirmwareNinjaReferenceNodeIsDataVariable(m_object);
+}
+
+
+bool FirmwareNinjaReferenceNode::HasChildren()
+{
+	return BNFirmwareNinjaReferenceNodeHasChildren(m_object);
+}
+
+
+bool FirmwareNinjaReferenceNode::GetFunction(Ref<Function>& function)
+{
+	auto bnFunction = BNFirmwareNinjaReferenceNodeGetFunction(m_object);
+	if (!bnFunction)
+		return false;
+
+	function = new Function(BNNewFunctionReference(bnFunction));
+	return true;
+}
+
+
+bool FirmwareNinjaReferenceNode::GetDataVariable(DataVariable& variable)
+{
+	auto bnVariable = BNFirmwareNinjaReferenceNodeGetDataVariable(m_object);
+	if (!bnVariable)
+		return false;
+
+	variable.address = bnVariable->address;
+	variable.type = Confidence(new Type(BNNewTypeReference(bnVariable->type)), bnVariable->typeConfidence);
+	variable.autoDiscovered = bnVariable->autoDiscovered;
+	BNFreeDataVariable(bnVariable);
+	return true;
+}
+
+
+std::vector<Ref<FirmwareNinjaReferenceNode>> FirmwareNinjaReferenceNode::GetChildren()
+{
+	std::vector<Ref<FirmwareNinjaReferenceNode>> result;
+	size_t count = 0;
+	auto bnChildren = BNFirmwareNinjaReferenceNodeGetChildren(m_object, &count);
+	result.reserve(count);
+	for (size_t i = 0; i < count; ++i)
+	{
+		result.push_back(new FirmwareNinjaReferenceNode(
+			BNNewFirmwareNinjaReferenceNodeReference(bnChildren[i])));
+	}
+
+	if (count)
+		BNFreeFirmwareNinjaReferenceNodes(bnChildren, count);
+	return result;
+}
+
+
 FirmwareNinja::FirmwareNinja(Ref<BinaryView> view)
 {
 	m_view = view;
-    m_object = BNCreateFirmwareNinja(view->GetObject());
+	m_object = BNCreateFirmwareNinja(view->GetObject());
 }
 
 
@@ -218,6 +291,9 @@ std::vector<FirmwareNinjaFunctionMemoryAccesses> FirmwareNinja::GetFunctionMemor
 
 void FirmwareNinja::StoreFunctionMemoryAccesses(const std::vector<FirmwareNinjaFunctionMemoryAccesses>& fma)
 {
+	if (fma.empty())
+		return;
+
 	BNFirmwareNinjaFunctionMemoryAccesses** fmaArray = MemoryInfoVectorToArray(fma);
 	BNFirmwareNinjaStoreFunctionMemoryAccessesToMetadata(m_object, fmaArray, fma.size());
 	FreeMemoryInfoArray(fmaArray, fma.size());
@@ -262,6 +338,9 @@ std::vector<FirmwareNinjaDeviceAccesses> FirmwareNinja::GetBoardDeviceAccesses(
 	const std::vector<FirmwareNinjaFunctionMemoryAccesses>& fma)
 {
 	std::vector<FirmwareNinjaDeviceAccesses> result;
+	if (fma.empty())
+		return result;
+
 	auto platform = m_view->GetDefaultPlatform();
 	if (!platform)
 		return result;
@@ -273,11 +352,9 @@ std::vector<FirmwareNinjaDeviceAccesses> FirmwareNinja::GetBoardDeviceAccesses(
 	BNFirmwareNinjaFunctionMemoryAccesses** fmaArray = MemoryInfoVectorToArray(fma);
 	BNFirmwareNinjaDeviceAccesses* accesses;
 	int count = BNFirmwareNinjaGetBoardDeviceAccesses(m_object, fmaArray, fma.size(), &accesses, arch->GetObject());
+	FreeMemoryInfoArray(fmaArray, fma.size());
 	if (count <= 0)
-	{
-		FreeMemoryInfoArray(fmaArray, fma.size());
 		return result;
-	}
 
 	result.reserve(count);
 	for (size_t i = 0; i < count; i++)
@@ -290,3 +367,56 @@ std::vector<FirmwareNinjaDeviceAccesses> FirmwareNinja::GetBoardDeviceAccesses(
 
 	return result;
 }
+
+
+Ref<FirmwareNinjaReferenceNode> FirmwareNinja::GetReferenceTree(
+	FirmwareNinjaDevice& device, const std::vector<FirmwareNinjaFunctionMemoryAccesses>& fma, uint64_t* value)
+{
+	BNFirmwareNinjaFunctionMemoryAccesses** fmaArray = nullptr;
+	if (!fma.empty())
+		fmaArray = MemoryInfoVectorToArray(fma);
+
+	auto bnReferenceTree = BNFirmwareNinjaGetMemoryRegionReferenceTree(
+		m_object, device.start, device.end, fmaArray, fma.size(), value);
+
+	FreeMemoryInfoArray(fmaArray, fma.size());
+	if (!bnReferenceTree)
+		return nullptr;
+
+	return new FirmwareNinjaReferenceNode(bnReferenceTree);
+}
+
+
+Ref<FirmwareNinjaReferenceNode> FirmwareNinja::GetReferenceTree(
+	Section& section, const std::vector<FirmwareNinjaFunctionMemoryAccesses>& fma, uint64_t* value)
+{
+	BNFirmwareNinjaFunctionMemoryAccesses** fmaArray = nullptr;
+	if (!fma.empty())
+		fmaArray = MemoryInfoVectorToArray(fma);
+
+	auto bnReferenceTree = BNFirmwareNinjaGetMemoryRegionReferenceTree(
+		m_object, section.GetStart(), section.GetStart() + section.GetLength(), fmaArray, fma.size(), value);
+
+	FreeMemoryInfoArray(fmaArray, fma.size());
+	if (!bnReferenceTree)
+		return nullptr;
+
+	return new FirmwareNinjaReferenceNode(bnReferenceTree);
+}
+
+
+Ref<FirmwareNinjaReferenceNode> FirmwareNinja::GetReferenceTree(
+	uint64_t address, const std::vector<FirmwareNinjaFunctionMemoryAccesses>& fma, uint64_t* value)
+{
+	BNFirmwareNinjaFunctionMemoryAccesses** fmaArray = nullptr;
+	if (!fma.empty())
+		fmaArray = MemoryInfoVectorToArray(fma);
+
+	auto bnReferenceTree = BNFirmwareNinjaGetAddressReferenceTree(m_object, address, fmaArray, fma.size(), value);
+
+	FreeMemoryInfoArray(fmaArray, fma.size());
+	if (!bnReferenceTree)
+		return nullptr;
+
+	return new FirmwareNinjaReferenceNode(bnReferenceTree);
+}
diff --git a/python/firmwareninja.py b/python/firmwareninja.py
index b000ecc9..834cdd80 100644
--- a/python/firmwareninja.py
+++ b/python/firmwareninja.py
@@ -21,8 +21,8 @@
 
 import ctypes
 from dataclasses import dataclass
-from typing import Callable
-from .binaryview import BinaryView
+from typing import Callable, Union, Optional
+from .binaryview import BinaryView, Section, DataVariable
 from .variable import RegisterValue
 from .enums import (
     FirmwareNinjaMemoryHeuristic,
@@ -34,6 +34,111 @@ from .function import Function
 from . import _binaryninjacore as core
 
 
+class FirmwareNinjaReferenceNode:
+    """
+    ``class FirmwareNinjaReferenceNode`` is a class for building reference trees for functions, data variables, and
+    memory regions.
+    """
+
+    def __init__(self, handle=None, view=None):
+        assert handle is not None, "Cannot create reference node directly, run `FirmwareNinja.get_reference_tree`"
+        self._handle = handle
+        self._view = view
+
+    def __del__(self):
+        if core is not None:
+            core.BNFreeFirmwareNinjaReferenceNode(self._handle)
+
+    def is_function(self) -> bool:
+        """
+        ``is_function`` determines if the reference tree node is for a function
+
+        :return: True if the reference tree node is for a function, False otherwise
+        :rtype: bool
+        """
+
+        return core.BNFirmwareNinjaReferenceNodeIsFunction(self._handle)
+
+    def is_data_variable(self) -> bool:
+        """
+        ``is_data_variable`` determines if the reference tree node is for a data variable
+
+        :return: True if the reference tree node is for a data variable, False otherwise
+        :rtype: bool
+        """
+
+        return core.BNFirmwareNinjaReferenceNodeIsDataVariable(self._handle)
+
+    def has_children(self) -> bool:
+        """
+        ``has_children`` determines if the reference tree node contains child reference tree nodes
+
+        :return: True if the reference tree node contains children, False otherwise
+        :rtype: bool
+        """
+
+        return core.BNFirmwareNinjaReferenceNodeHasChildren(self._handle)
+
+    @property
+    def function(self) -> Function:
+        """
+        ``function`` query the function from the reference tree node
+
+        :return: Function contained in the reference tree node
+        :rtype: Function
+        """
+
+        bn_function = core.BNFirmwareNinjaReferenceNodeGetFunction(self._handle)
+        if not bn_function:
+            return None
+
+        return Function(handle=bn_function)
+
+    @property
+    def data_variable(self) -> DataVariable:
+        """
+        ``data_variable`` query the data variable from the reference tree node
+
+        :return: Data variable contained in the reference tree node
+        :rtype: DataVariable
+        """
+
+        try:
+            bn_data_var = core.BNFirmwareNinjaReferenceNodeGetDataVariable(
+                self._handle)
+            if not bn_data_var:
+                return None
+
+            data_var = DataVariable.from_core_struct(bn_data_var.contents, self._view)
+        finally:
+            core.BNFreeDataVariable(bn_data_var)
+        return data_var
+
+    @property
+    def children(self) -> list['FirmwareNinjaReferenceNode']:
+        """
+        ``children`` query the child reference tree nodes
+
+        :return: Child nodes contained in the reference tree node
+        :rtype: list[FirmwareNinjaReferenceNode]
+        """
+
+        count = ctypes.c_ulonglong(0)
+        nodes = []
+        try:
+            bn_nodes = core.BNFirmwareNinjaReferenceNodeGetChildren(
+                self._handle, count)
+            for i in range(count.value):
+                nodes.append(
+                    FirmwareNinjaReferenceNode(
+                        core.BNNewFirmwareNinjaReferenceNodeReference(
+                            bn_nodes[i]), self._view))
+        finally:
+            core.BNFreeFirmwareNinjaReferenceNodes(bn_nodes, count.value)
+
+        return nodes
+
+
 @dataclass
 class FirmwareNinjaDevice:
     """
@@ -74,7 +179,9 @@ class FirmwareNinjaMemoryAccess:
     value: RegisterValue
 
     @classmethod
-    def from_BNFirmwareNinjaMemoryAccess(cls, access: core.BNFirmwareNinjaMemoryAccess) -> "FirmwareNinjaMemoryAccess":
+    def from_BNFirmwareNinjaMemoryAccess(
+        cls, access: core.BNFirmwareNinjaMemoryAccess
+    ) -> "FirmwareNinjaMemoryAccess":
         return cls(
             instr_address=access.instrAddress,
             mem_address=RegisterValue.from_BNRegisterValue(access.memAddress),
@@ -84,7 +191,9 @@ class FirmwareNinjaMemoryAccess:
         )
 
     @classmethod
-    def to_BNFirmwareNinjaMemoryAccess(cls, access: "FirmwareNinjaMemoryAccess") -> core.BNFirmwareNinjaMemoryAccess:
+    def to_BNFirmwareNinjaMemoryAccess(
+        cls, access: "FirmwareNinjaMemoryAccess"
+    ) -> core.BNFirmwareNinjaMemoryAccess:
         return core.BNFirmwareNinjaMemoryAccess(
             instrAddress=access.instr_address,
             memAddress=RegisterValue.to_BNRegisterValue(access.mem_address),
@@ -113,7 +222,9 @@ class FirmwareNinjaFunctionMemoryAccesses:
         accesses = []
         for i in range(info.count):
             access = info.accesses[i]
-            accesses.append(FirmwareNinjaMemoryAccess.from_BNFirmwareNinjaMemoryAccess(access.contents))
+            accesses.append(
+                FirmwareNinjaMemoryAccess.from_BNFirmwareNinjaMemoryAccess(
+                    access.contents))
 
         return cls(
             function=view.get_function_at(info.start),
@@ -156,9 +267,10 @@ class FirmwareNinja:
         if core is not None:
             core.BNFreeFirmwareNinja(self._handle)
 
-    def store_custom_device(self, name: str, start: int, size: int, info: str) -> bool:
+    def store_custom_device(self, name: str, start: int, size: int,
+                            info: str) -> bool:
         """
-        ``store_custom_device`` store a user-defined Firmware Ninja device in the binary view metadata
+        ``store_custom_device`` stores a user-defined Firmware Ninja device in the binary view metadata
 
         :param str name: Name of the device
         :param int start: Start address of the device
@@ -168,7 +280,8 @@ class FirmwareNinja:
         :rtype: bool
         """
 
-        return core.BNFirmwareNinjaStoreCustomDevice(self._handle, name, start, start + size, info)
+        return core.BNFirmwareNinjaStoreCustomDevice(self._handle, name, start,
+                                                     start + size, info)
 
     def remove_custom_device(self, name: str) -> bool:
         """
@@ -186,12 +299,13 @@ class FirmwareNinja:
         """
         ``query_custom_devices`` queries user-defined Firmware Ninja devices from the binary view metadata
 
-        :return: List of Firmware Ninja device objects
+        :return: List of Firmware Ninja devices
         :rtype: list[FirmwareNinjaDevice]
         """
 
         devices = ctypes.POINTER(core.BNFirmwareNinjaDevice)()
-        count = core.BNFirmwareNinjaQueryCustomDevices(self._handle, ctypes.byref(devices))
+        count = core.BNFirmwareNinjaQueryCustomDevices(self._handle,
+                                                       ctypes.byref(devices))
         if count == -1:
             raise RuntimeError("BNFirmwareNinjaQueryCustomDevices")
 
@@ -204,8 +318,7 @@ class FirmwareNinja:
                         start=devices[i].start,
                         size=devices[i].end - devices[i].start,
                         info=devices[i].info,
-                    )
-                )
+                    ))
 
             return device_list
         finally:
@@ -221,8 +334,7 @@ class FirmwareNinja:
 
         boards = ctypes.POINTER(ctypes.c_char_p)()
         count = core.BNFirmwareNinjaQueryBoardNamesForArchitecture(
-            self._handle, self._view.arch.handle, ctypes.byref(boards)
-        )
+            self._handle, self._view.arch.handle, ctypes.byref(boards))
         if count == -1:
             raise RuntimeError("BNFirmwareNinjaQueryBoardNamesForArchitecture")
 
@@ -235,7 +347,8 @@ class FirmwareNinja:
         finally:
             core.BNFirmwareNinjaFreeBoardNames(boards, count)
 
-    def query_devices_by_board_name(self, name: str) -> list[FirmwareNinjaDevice]:
+    def query_devices_by_board_name(self,
+                                    name: str) -> list[FirmwareNinjaDevice]:
         """
         ``query_devices_by_board_name`` queries the hardware device information for a specific board
 
@@ -246,12 +359,15 @@ class FirmwareNinja:
             FirmwareNinjaDevice(name='nand@12f', start=303, size=1024, info='marvell,orion-nand')
 
         :param str name: Name of the board
-        :return: List of Firmware Ninja device objects
+        :return: List of Firmware Ninja devices
         :rtype: list[FirmwareNinjaDevice]
         """
 
         devices = ctypes.POINTER(core.BNFirmwareNinjaDevice)()
-        count = core.BNFirmwareNinjaQueryBoardDevices(self._handle, self._view.arch.handle, name, ctypes.byref(devices))
+        count = core.BNFirmwareNinjaQueryBoardDevices(self._handle,
+                                                      self._view.arch.handle,
+                                                      name,
+                                                      ctypes.byref(devices))
         if count == -1:
             raise RuntimeError("BNFirmwareNinjaQueryBoardDevices")
 
@@ -264,8 +380,7 @@ class FirmwareNinja:
                         start=devices[i].start,
                         size=devices[i].end - devices[i].start,
                         info=devices[i].info,
-                    )
-                )
+                    ))
 
             return device_list
         finally:
@@ -276,7 +391,8 @@ class FirmwareNinja:
         high_code_entropy_threshold: float = 0.910,
         low_code_entropy_threshold: float = 0.500,
         block_size: int = 4096,
-        mode: FirmwareNinjaSectionAnalysisMode = FirmwareNinjaSectionAnalysisMode.DetectStringsSectionAnalysisMode,
+        mode: FirmwareNinjaSectionAnalysisMode = FirmwareNinjaSectionAnalysisMode
+        .DetectStringsSectionAnalysisMode,
     ) -> list[FirmwareNinjaSection]:
         """
         ``find_sections`` finds sections with Firmware Ninja entropy analysis and heuristics
@@ -318,36 +434,39 @@ class FirmwareNinja:
                         start=sections[i].start,
                         size=sections[i].end - sections[i].start,
                         entropy=sections[i].entropy,
-                    )
-                )
+                    ))
 
             return section_list
         finally:
             core.BNFirmwareNinjaFreeSections(sections, count)
 
-    def get_function_memory_accesses(self, progress_func: Callable = None) -> list[FirmwareNinjaFunctionMemoryAccesses]:
+    def get_function_memory_accesses(
+        self,
+        progress_func: Callable = None
+    ) -> list[FirmwareNinjaFunctionMemoryAccesses]:
         """
         ``get_function_memory_accesses`` runs analysis to find accesses to memory regions that are not file-backed, such
-        as memory-mapped I/O and RAM.
+        as memory-mapped I/O and RAM
 
         :param callback progress_func: optional function to be called with the current progress and total count.
-        :return: List of function memory accesses objects
+        :return: List of function memory accesses
         :rtype: list[FirmwareNinjaFunctionMemoryAccesses]
         """
 
-        fma_info = ctypes.POINTER((ctypes.POINTER(core.BNFirmwareNinjaFunctionMemoryAccesses)))()
+        fma_info = ctypes.POINTER(
+            (ctypes.POINTER(core.BNFirmwareNinjaFunctionMemoryAccesses)))()
         if progress_func is None:
-            progress_cfunc = ctypes.CFUNCTYPE(ctypes.c_bool, ctypes.c_void_p, ctypes.c_ulonglong, ctypes.c_ulonglong)(
-                lambda ctxt, cur, total: True
-            )
+            progress_cfunc = ctypes.CFUNCTYPE(
+                ctypes.c_bool, ctypes.c_void_p, ctypes.c_ulonglong,
+                ctypes.c_ulonglong)(lambda ctxt, cur, total: True)
         else:
-            progress_cfunc = ctypes.CFUNCTYPE(ctypes.c_bool, ctypes.c_void_p, ctypes.c_ulonglong, ctypes.c_ulonglong)(
-                lambda ctxt, cur, total: progress_func(cur, total)
-            )
+            progress_cfunc = ctypes.CFUNCTYPE(
+                ctypes.c_bool, ctypes.c_void_p, ctypes.c_ulonglong,
+                ctypes.c_ulonglong)(
+                    lambda ctxt, cur, total: progress_func(cur, total))
 
         count = core.BNFirmwareNinjaGetFunctionMemoryAccesses(
-            self._handle, ctypes.byref(fma_info), progress_cfunc, None
-        )
+            self._handle, ctypes.byref(fma_info), progress_cfunc, None)
         if count == -1:
             raise RuntimeError("BNFirmwareNinjaGetFunctionMemoryAccesses")
 
@@ -355,33 +474,41 @@ class FirmwareNinja:
             fma_info_list = []
             for i in range(count):
                 fma_info_list.append(
-                    FirmwareNinjaFunctionMemoryAccesses.from_BNFirmwareNinjaFunctionMemoryAccesses(
-                        fma_info[i].contents, self._view
-                    )
-                )
+                    FirmwareNinjaFunctionMemoryAccesses.
+                    from_BNFirmwareNinjaFunctionMemoryAccesses(
+                        fma_info[i].contents, self._view))
 
             return fma_info_list
         finally:
             core.BNFirmwareNinjaFreeFunctionMemoryAccesses(fma_info, count)
 
-    def _fma_info_list_to_array(self, fma: list[FirmwareNinjaFunctionMemoryAccesses]) -> ctypes.POINTER:
-        fma_info_ptr_array = (ctypes.POINTER(core.BNFirmwareNinjaFunctionMemoryAccesses) * len(fma))()
+    def _fma_info_list_to_array(
+            self,
+            fma: list[FirmwareNinjaFunctionMemoryAccesses]) -> ctypes.POINTER:
+        fma_info_ptr_array = (
+            ctypes.POINTER(core.BNFirmwareNinjaFunctionMemoryAccesses) *
+            len(fma))()
         for i, info in enumerate(fma):
-            accesses_ptr_array = (ctypes.POINTER(core.BNFirmwareNinjaMemoryAccess) * len(info.accesses))()
+            accesses_ptr_array = (
+                ctypes.POINTER(core.BNFirmwareNinjaMemoryAccess) *
+                len(info.accesses))()
             for j, access in enumerate(info.accesses):
-                accesses_ptr_array[j] = ctypes.pointer(FirmwareNinjaMemoryAccess.to_BNFirmwareNinjaMemoryAccess(access))
+                accesses_ptr_array[j] = ctypes.pointer(
+                    FirmwareNinjaMemoryAccess.to_BNFirmwareNinjaMemoryAccess(
+                        access))
 
             fma_info_struct = core.BNFirmwareNinjaFunctionMemoryAccesses(
-                function=info.function.handle,
-                accesses=accesses_ptr_array,
+                start=info.function.start,
                 count=len(info.accesses),
+                accesses=accesses_ptr_array,
             )
 
             fma_info_ptr_array[i] = ctypes.pointer(fma_info_struct)
 
         return fma_info_ptr_array
 
-    def store_function_memory_accesses(self, fma: list[FirmwareNinjaFunctionMemoryAccesses]) -> None:
+    def store_function_memory_accesses(
+            self, fma: list[FirmwareNinjaFunctionMemoryAccesses]) -> None:
         """
         ``store_function_memory_accesses`` saves information on function memory accesses to binary view metadata
 
@@ -391,24 +518,28 @@ class FirmwareNinja:
             >>> fma = fwn.get_function_memory_accesses()
             >>> fwn.store_function_memory_accesses(fma)
 
-        :param list[FirmwareNinjaFunctionMemoryAccesses] fma: List of function memory accesses objects
+        :param list[FirmwareNinjaFunctionMemoryAccesses] fma: List of function memory accesses
         :return: None
         :rtype: None
         """
 
         fma_info_ptr_array = self._fma_info_list_to_array(fma)
-        core.BNFirmwareNinjaStoreFunctionMemoryAccessesToMetadata(self._handle, fma_info_ptr_array, len(fma))
+        core.BNFirmwareNinjaStoreFunctionMemoryAccessesToMetadata(
+            self._handle, fma_info_ptr_array, len(fma))
 
-    def query_function_memory_accesses(self) -> list[FirmwareNinjaFunctionMemoryAccesses]:
+    def query_function_memory_accesses(
+            self) -> list[FirmwareNinjaFunctionMemoryAccesses]:
         """
         ``query_function_memory_accesses`` queries information on function memory accesses from binary view metadata
 
-        :return: List of function memory accesses objects
+        :return: List of function memory accesses
         :rtype: list[FirmwareNinjaFunctionMemoryAccesses]
         """
 
-        fma = ctypes.POINTER((ctypes.POINTER(core.BNFirmwareNinjaFunctionMemoryAccesses)))()
-        count = core.BNFirmwareNinjaQueryFunctionMemoryAccessesFromMetadata(self._handle, ctypes.byref(fma))
+        fma = ctypes.POINTER(
+            (ctypes.POINTER(core.BNFirmwareNinjaFunctionMemoryAccesses)))()
+        count = core.BNFirmwareNinjaQueryFunctionMemoryAccessesFromMetadata(
+            self._handle, ctypes.byref(fma))
         if count == -1:
             return None
 
@@ -416,10 +547,9 @@ class FirmwareNinja:
             fma_info_list = []
             for i in range(count):
                 fma_info_list.append(
-                    FirmwareNinjaFunctionMemoryAccesses.from_BNFirmwareNinjaFunctionMemoryAccesses(
-                        fma[i].contents, self._view
-                    )
-                )
+                    FirmwareNinjaFunctionMemoryAccesses.
+                    from_BNFirmwareNinjaFunctionMemoryAccesses(
+                        fma[i].contents, self._view))
 
             return fma_info_list
         finally:
@@ -439,16 +569,16 @@ class FirmwareNinja:
             >>> fwn.get_board_device_accesses(fma)[0]
             FirmwareNinjaDeviceAccesses(board_name='stm32mp157c-dhcom-picoitx', total=414, unique=2)
 
-        :param list[FirmwareNinjaFunctionMemoryAccesses] fma: List of function memory accesses objects
-        :return: List of device accesses objects
+        :param list[FirmwareNinjaFunctionMemoryAccesses] fma: List of function memory accesses
+        :return: List of device accesses
         :rtype: list[FirmwareNinjaDeviceAccesses]
         """
 
         fma_info_ptr_array = self._fma_info_list_to_array(fma)
         device_accesses = ctypes.POINTER(core.BNFirmwareNinjaDeviceAccesses)()
         count = core.BNFirmwareNinjaGetBoardDeviceAccesses(
-            self._handle, fma_info_ptr_array, len(fma), ctypes.byref(device_accesses), self._view.arch.handle
-        )
+            self._handle, fma_info_ptr_array, len(fma),
+            ctypes.byref(device_accesses), self._view.arch.handle)
         if count == -1:
             raise RuntimeError("BNFirmwareNinjaGetBoardDeviceAccesses")
 
@@ -460,9 +590,62 @@ class FirmwareNinja:
                         board_name=device_accesses[i].name,
                         total=device_accesses[i].total,
                         unique=device_accesses[i].unique,
-                    )
-                )
+                    ))
 
             return device_accesses_list
         finally:
             core.BNFirmwareNinjaFreeBoardDeviceAccesses(device_accesses, count)
+
+    def get_reference_tree(
+            self,
+            location: Union[Section, FirmwareNinjaDevice, Function, DataVariable, int],
+            fma: list[FirmwareNinjaFunctionMemoryAccesses],
+            value: Optional[int] = None) -> FirmwareNinjaReferenceNode:
+        """
+        ``get_reference_tree`` returns a tree of references for a memory region, function, or data location
+
+        :param Union[Section, FirmwareNinjaDevice, DataVariable, Function, int] location: Memory location to build the
+        reference tree for
+        :param list[FirmwareNinjaFunctionMemoryAccesses] fma: List of function memory accesses or None to use cross
+        references. None should only be supplied if location is a Function, DataVariable, or address.
+        :param Optional[int] value: Only include the node in the tree if this value is written to the location
+        :return: Root reference node containing the reference tree
+        :rtype: FirmwareNinjaReferenceNode
+        """
+
+        if fma is None and (isinstance(location, Section) or isinstance(location, FirmwareNinjaDevice)):
+            raise ValueError("Function memory accesses cannot be None for location type Section or FirmwareNinjaDevice")
+
+        value = ctypes.pointer(
+            ctypes.c_uint64(value)) if value is not None else None
+
+        fma_info_ptr_array = None
+        if fma is not None and len(fma) > 0:
+            fma_info_ptr_array = self._fma_info_list_to_array(fma)
+
+        if isinstance(location, FirmwareNinjaDevice):
+            bn_node = core.BNFirmwareNinjaGetMemoryRegionReferenceTree(
+                self._handle, location.start, location.start + location.size,
+                fma_info_ptr_array, len(fma), value)
+        elif isinstance(location, Function):
+            bn_node = core.BNFirmwareNinjaGetAddressReferenceTree(
+                self._handle, location.start, fma_info_ptr_array, len(fma),
+                value)
+        elif isinstance(location, Section):
+            bn_node = core.BNFirmwareNinjaGetMemoryRegionReferenceTree(
+                self._handle, location.start, location.start + location.length,
+                fma_info_ptr_array, len(fma), value)
+        elif isinstance(location, DataVariable):
+            bn_node = core.BNFirmwareNinjaGetAddressReferenceTree(
+                self._handle, location.address, fma_info_ptr_array, len(fma),
+                value)
+        elif isinstance(location, int):
+            bn_node = core.BNFirmwareNinjaGetAddressReferenceTree(
+                self._handle, location, fma_info_ptr_array, len(fma), value)
+        else:
+            raise ValueError("Invalid location type")
+
+        if not bn_node:
+            return None
+
+        return FirmwareNinjaReferenceNode(handle=bn_node, view=self._view)