Render Layers

1 files changed, 450 insertions, 0 deletions

diff --git a/python/renderlayer.py b/python/renderlayer.py
new file mode 100644
index 00000000..6f5984e0
--- /dev/null
+++ b/python/renderlayer.py
@@ -0,0 +1,450 @@
+# Copyright (c) 2015-2024 Vector 35 Inc
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to
+# deal in the Software without restriction, including without limitation the
+# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+# sell copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+# IN THE SOFTWARE.
+
+import ctypes
+import traceback
+
+# Binary Ninja components
+import binaryninja
+from . import _binaryninjacore as core, LinearDisassemblyLine
+from .enums import LinearDisassemblyLineType, RenderLayerDefaultEnableState
+from . import binaryview
+from . import types
+from .log import log_error
+from typing import Iterable, List, Optional, Union, Tuple
+
+
+class _RenderLayerMetaclass(type):
+	def __iter__(self):
+		binaryninja._init_plugins()
+		count = ctypes.c_ulonglong()
+		instances = core.BNGetRenderLayerList(count)
+		try:
+			for i in range(0, count.value):
+				yield self._handle_to_instance(instances[i])
+		finally:
+			core.BNFreeRenderLayerList(instances)
+
+	def __getitem__(self, value):
+		binaryninja._init_plugins()
+		handle = core.BNGetRenderLayerByName(str(value))
+		if handle is None:
+			raise KeyError(f"'{value}' is not a valid RenderLayer")
+		return self._handle_to_instance(handle)
+
+	def _handle_to_instance(self, handle):
+		handle_ptr = ctypes.cast(handle, ctypes.c_void_p)
+		if handle_ptr.value in RenderLayer._registered_instances:
+			return RenderLayer._registered_instances[handle_ptr.value]
+		return CoreRenderLayer(handle)
+
+
+class RenderLayer(metaclass=_RenderLayerMetaclass):
+	"""
+	RenderLayer is a plugin class that allows you to customize the presentation of
+	Linear and Graph view output, adding, changing, or removing lines before they are
+	presented in the UI.
+	"""
+
+	name = None
+	default_enable_state = RenderLayerDefaultEnableState.DisabledByDefaultRenderLayerDefaultEnableState
+	_registered_instances = {}
+	_pending_lines = {}
+
+	def __init__(self, handle=None):
+		if handle is not None:
+			self.handle = core.handle_of_type(handle, core.BNRenderLayer)
+			self.__dict__["name"] = core.BNGetRenderLayerName(handle)
+			self.__dict__["default_enable_state"] = core.BNGetRenderLayerDefaultEnableState(handle)
+		else:
+			self.handle = None
+
+	@classmethod
+	def register(cls):
+		"""
+		Register a custom Render Layer.
+		"""
+		layer = cls()
+
+		assert layer.__class__.name is not None
+		assert layer.handle is None
+
+		layer._cb = core.BNRenderLayerCallbacks()
+		layer._cb.context = 0
+		layer._cb.applyToFlowGraph = layer._cb.applyToFlowGraph.__class__(layer._apply_to_flow_graph)
+		layer._cb.applyToLinearViewObject = layer._cb.applyToLinearViewObject.__class__(layer._apply_to_linear_view_object)
+		layer._cb.freeLines = layer._cb.freeLines.__class__(layer._free_lines)
+		layer.handle = core.BNRegisterRenderLayer(layer.__class__.name, layer._cb, layer.default_enable_state)
+		handle_ptr = ctypes.cast(layer.handle, ctypes.c_void_p)
+		cls._registered_instances[handle_ptr.value] = layer
+
+	def __eq__(self, other):
+		if not isinstance(other, RenderLayer):
+			return False
+		return self.name == other.name
+
+	def __str__(self):
+		return f'<RenderLayer: {self.name}>'
+
+	def __repr__(self):
+		return f'<RenderLayer: {self.name}>'
+
+	def _apply_to_flow_graph(self, ctxt, graph):
+		try:
+			self.apply_to_flow_graph(binaryninja.FlowGraph(handle=core.BNNewFlowGraphReference(graph)))
+		except:
+			log_error(traceback.format_exc())
+
+	def _apply_to_linear_view_object(self, ctxt, obj, prev, next, in_lines, in_line_count, out_lines, out_line_count):
+		try:
+			obj_obj = binaryninja.LinearViewObject(core.BNNewLinearViewObjectReference(obj))
+			prev_obj = binaryninja.LinearViewObject(core.BNNewLinearViewObjectReference(prev)) if prev else None
+			next_obj = binaryninja.LinearViewObject(core.BNNewLinearViewObjectReference(next)) if next else None
+
+			lines = []
+			for i in range(in_line_count):
+				lines.append(LinearDisassemblyLine._from_core_struct(in_lines[i], obj=obj_obj))
+
+			lines = self.apply_to_linear_view_object(obj_obj, prev_obj, next_obj, lines)
+
+			out_line_count[0] = len(lines)
+			out_lines_buf = (core.BNLinearDisassemblyLine * len(lines))()
+			for i, r in enumerate(lines):
+				out_lines_buf[i] = r._to_core_struct()
+			out_lines_ptr = ctypes.cast(out_lines_buf, ctypes.c_void_p)
+			out_lines[0] = out_lines_buf
+			self._pending_lines[out_lines_ptr.value] = (out_lines_ptr.value, out_lines_buf)
+		except:
+			log_error(traceback.format_exc())
+			out_lines[0] = None
+			out_line_count[0] = 0
+
+	def _free_lines(self, ctxt, lines, count):
+		try:
+			buf = ctypes.cast(lines, ctypes.c_void_p)
+			if buf.value is not None:
+				if buf.value not in self._pending_lines:
+					raise ValueError("freeing lines list that wasn't allocated")
+				del self._pending_lines[buf.value]
+		except:
+			log_error(traceback.format_exc())
+
+	def apply_to_disassembly_block(
+			self,
+			block: 'binaryninja.BasicBlock',
+			lines: List['binaryninja.DisassemblyTextLine']
+	):
+		"""
+		Apply this Render Layer to a single Basic Block of Disassembly lines.
+		Subclasses should modify the input `lines` list to make modifications to
+		the presentation of the block.
+
+		.. note:: This function will only handle Disassembly lines, and not any ILs.
+
+		:param block: Basic Block containing those lines
+		:param lines: Original lines of text for the block
+		:return: Modified list of lines
+		"""
+		return lines
+
+	def apply_to_low_level_il_block(
+			self,
+			block: 'binaryninja.LowLevelILBasicBlock',
+			lines: List['binaryninja.DisassemblyTextLine']
+	):
+		"""
+		Apply this Render Layer to a single Basic Block of Low Level IL lines.
+		Subclasses should modify the input `lines` list to make modifications to
+		the presentation of the block.
+
+		.. note:: This function will only handle Lifted IL/LLIL/LLIL(SSA) lines. \
+		You can use the block's `function_graph_type` property to determine which is being handled.
+
+		:param block: Basic Block containing those lines
+		:param lines: Original lines of text for the block
+		:return: Modified list of lines
+		"""
+		return lines
+
+	def apply_to_medium_level_il_block(
+			self,
+			block: 'binaryninja.MediumLevelILBasicBlock',
+			lines: List['binaryninja.DisassemblyTextLine']
+	):
+		"""
+		Apply this Render Layer to a single Basic Block of Medium Level IL lines.
+		Subclasses should modify the input `lines` list to make modifications to
+		the presentation of the block.
+
+		.. note:: This function will only handle MLIL/MLIL(SSA)/Mapped MLIL/Mapped MLIL(SSA) lines. \
+		You can use the block's `function_graph_type` property to determine which is being handled.
+
+		:param block: Basic Block containing those lines
+		:param lines: Original lines of text for the block
+		:return: Modified list of lines
+		"""
+		return lines
+
+	def apply_to_high_level_il_block(
+			self,
+			block: 'binaryninja.HighLevelILBasicBlock',
+			lines: List['binaryninja.DisassemblyTextLine']
+	):
+		"""
+		Apply this Render Layer to a single Basic Block of High Level IL lines.
+		Subclasses should modify the input `lines` list to make modifications to
+		the presentation of the block.
+
+		.. note:: This function will only handle HLIL/HLIL(SSA)/Language Representation lines. \
+		You can use the block's `function_graph_type` property to determine which is being handled.
+
+		.. warning:: This function will NOT apply to High Level IL bodies as displayed \
+		in Linear View! Those are handled by `apply_to_high_level_il_body` instead as they \
+		do not have a Basic Block associated with them.
+
+		:param block: Basic Block containing those lines
+		:param lines: Original lines of text for the block
+		:return: Modified list of lines
+		"""
+		return lines
+
+	def apply_to_high_level_il_body(
+			self,
+			function: 'binaryninja.Function',
+			lines: List['binaryninja.LinearDisassemblyLine']
+	):
+		"""
+		Apply this Render Layer to the entire body of a High Level IL function.
+		Subclasses should modify the input `lines` list to make modifications to
+		the presentation of the function.
+
+		.. warning:: This function only applies to Linear View, and not to Graph View! \
+		If you want to handle Graph View too, you will need to use `apply_to_high_level_il_block` \
+		and handle the lines one block at a time.
+
+		:param function: Function containing those lines
+		:param lines: Original lines of text for the function
+		:return: Modified list of lines
+		"""
+		return lines
+
+	def apply_to_misc_linear_lines(
+			self,
+			obj: 'binaryninja.LinearViewObject',
+			prev: Optional['binaryninja.LinearViewObject'],
+			next: Optional['binaryninja.LinearViewObject'],
+			lines: List['binaryninja.LinearDisassemblyLine']
+	):
+		"""
+		Apply to lines generated by Linear View that are not part of a function.
+		It is up to your implementation to figure out which type of Linear View Object
+		lines these are, and what to do with them.
+
+		:param obj: Linear View Object being rendered
+		:param prev: Linear View Object located directly above this one
+		:param next: Linear View Object located directly below this one
+		:param lines: Original lines rendered by `obj`
+		:return: Modified list of lines
+		"""
+		return lines
+
+	def apply_to_block(
+			self,
+			block: 'binaryninja.BasicBlock',
+			lines: List['binaryninja.DisassemblyTextLine'],
+	):
+		"""
+		Apply to lines generated by a Basic Block, of any type. If not overridden, this
+		function will call the appropriate apply_to_X_level_il_block function.
+
+		:param block: Basic Block containing those lines
+		:param lines: Original lines of text for the block
+		:return: Modified list of lines
+		"""
+		if not block.is_il:
+			return self.apply_to_disassembly_block(block, lines)
+		elif block.is_low_level_il:
+			return self.apply_to_low_level_il_block(block, lines)
+		elif block.is_medium_level_il:
+			return self.apply_to_medium_level_il_block(block, lines)
+		elif block.is_high_level_il:
+			return self.apply_to_high_level_il_block(block, lines)
+		else:
+			# ???
+			return lines
+
+	def apply_to_flow_graph(self, graph: 'binaryninja.FlowGraph') -> None:
+		"""
+		Apply this Render Layer to a Flow Graph, potentially modifying its nodes,
+		their edges, their lines, and their lines' content.
+
+		.. note:: If you override this function, you will need to call the ``super()`` \
+		implementation if you want to use the higher level ``apply_to_X_level_il_block`` \
+		functionality.
+
+		:param graph: Graph to modify
+		"""
+		pass
+		for i, node in enumerate(graph.nodes):
+			lines = node.lines
+			if node.basic_block is not None and isinstance(node.basic_block, binaryninja.BasicBlock):
+				lines = self.apply_to_block(node.basic_block, lines)
+			node.lines = lines
+
+	def apply_to_linear_view_object(
+			self,
+			obj: 'binaryninja.LinearViewObject',
+			prev: Optional['binaryninja.LinearViewObject'],
+			next: Optional['binaryninja.LinearViewObject'],
+			lines: List['binaryninja.LinearDisassemblyLine']
+	) -> List['binaryninja.LinearDisassemblyLine']:
+		"""
+		Apply this Render Layer to the lines produced by a LinearViewObject for rendering
+		in Linear View, potentially modifying the lines and their contents.
+
+		.. note:: If you override this function, you will need to call the ``super()`` \
+		implementation if you want to use the higher level ``apply_to_X_level_il_block`` \
+		functionality.
+
+		:param obj: Linear View Object being rendered
+		:param prev: Linear View Object located directly above this one
+		:param next: Linear View Object located directly below this one
+		:param lines: Original lines rendered by the Linear View Object
+		:return: Modified list of lines to display in Linear View
+		"""
+		# Hack: HLIL bodies don't have basic blocks
+		if len(lines) > 0 and obj.identifier.name in [
+			"HLIL Function Body",
+			"HLIL SSA Function Body",
+			"Language Representation Function Body"
+		]:
+			return self.apply_to_high_level_il_body(lines[0].function, lines)
+
+		block_lines = []
+		final_lines = []
+		last_block = None
+
+		def finish_block():
+			nonlocal block_lines
+			nonlocal final_lines
+			if len(block_lines) > 0:
+				if last_block is not None:
+					# Convert linear lines to disassembly lines for the apply()
+					# and then convert back for linear view
+					new_block_lines = []
+					disasm_lines = []
+					misc_lines = []
+
+					def process_disasm():
+						nonlocal disasm_lines
+
+						if len(disasm_lines) > 0:
+							disasm_lines = self.apply_to_block(last_block, disasm_lines)
+							func = block_lines[0].function
+							block = block_lines[0].block
+							for block_line in disasm_lines:
+								new_block_lines.append(
+									LinearDisassemblyLine(
+										LinearDisassemblyLineType.CodeDisassemblyLineType,
+										func,
+										block,
+										block_line
+									)
+								)
+							disasm_lines = []
+
+					def process_misc():
+						nonlocal misc_lines
+						nonlocal new_block_lines
+
+						if len(misc_lines) > 0:
+							misc_lines = self.apply_to_misc_linear_lines(obj, prev, next, misc_lines)
+							new_block_lines += misc_lines
+							misc_lines = []
+
+					for block_line in block_lines:
+						# Lines in the block get sent to process_disasm, anything else goes
+						# to process_misc so we preserve line information
+						if block_line.type == LinearDisassemblyLineType.CodeDisassemblyLineType:
+							process_misc()
+							disasm_lines.append(block_line.contents)
+						else:
+							process_disasm()
+							misc_lines.append(block_line)
+
+					# At the end, zero or one of these has lines in it
+					process_misc()
+					process_disasm()
+					block_lines = new_block_lines
+				else:
+					block_lines = self.apply_to_misc_linear_lines(obj, prev, next, block_lines)
+				final_lines += block_lines
+				block_lines = []
+
+		for line in lines:
+			# Assume we've finished a block when the line's block changes
+			if line.block != last_block:
+				finish_block()
+			block_lines.append(line)
+			last_block = line.block
+
+		# And we've finished a block when we're done with every line
+		finish_block()
+		return final_lines
+
+
+class CoreRenderLayer(RenderLayer):
+
+	def apply_to_flow_graph(self, graph: 'binaryninja.FlowGraph') -> None:
+		core.BNApplyRenderLayerToFlowGraph(self.handle, graph.handle)
+
+	def apply_to_linear_view_object(
+			self,
+			obj: 'binaryninja.LinearViewObject',
+			prev: Optional['binaryninja.LinearViewObject'],
+			next: Optional['binaryninja.LinearViewObject'],
+			lines: List['binaryninja.LinearDisassemblyLine']
+	) -> List['binaryninja.LinearDisassemblyLine']:
+
+		in_lines_buf = (core.BNLinearDisassemblyLine * len(lines))()
+		for i, r in enumerate(lines):
+			in_lines_buf[i] = r._to_core_struct()
+
+		out_lines = ctypes.POINTER(core.BNLinearDisassemblyLine)()
+		out_line_count = ctypes.c_size_t(0)
+
+		core.BNApplyRenderLayerToLinearViewObject(
+			self.handle,
+			obj.handle,
+			prev.handle if prev is not None else None,
+			next.handle if next is not None else None,
+			in_lines_buf,
+			len(lines),
+			out_lines,
+			out_line_count
+		)
+
+		result = []
+		for i in range(out_line_count.value):
+			result.append(binaryninja.LinearDisassemblyLine._from_core_struct(out_lines[i], obj=obj))
+
+		core.BNFreeLinearDisassemblyLines(out_lines, out_line_count.value)
+
+		return result