Putting all of these into the base class

Author: CouleeApps

python/renderlayer.py

@@ -50,12 +50,11 @@ def __getitem__(self, value):
 		return CoreRenderLayer(handle)
 
 
-class BaseRenderLayer(metaclass=_RenderLayerMetaclass):
+class RenderLayer(metaclass=_RenderLayerMetaclass):
 	"""
-	BaseRenderLayer provides a lower-level interface for manipulating graphs
-	and lines directly, whereas :py:class:`RenderLayer` provides an interface where you
-	can choose to target a specific IL level or a couple different Linear View specific
-	line constructs.
+	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
@@ -88,7 +87,7 @@ def register(cls):
 		cls._registered_layers.append(layer)
 
 	def __eq__(self, other):
-		if not isinstance(other, BaseRenderLayer):
+		if not isinstance(other, RenderLayer):
 			return False
 		return self.name == other.name
 
@@ -138,87 +137,6 @@ def _free_lines(self, ctxt, lines, count):
 		except:
 			log_error(traceback.format_exc())
 
-	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.
-
-		:param graph: Graph to modify
-		"""
-		pass
-
-	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.
-
-		: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: Lines originally rendered by the Linear View Object
-		:return: Updated set of lines to display in Linear View
-		"""
-		return lines
-
-
-class CoreRenderLayer(BaseRenderLayer):
-
-	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
-
-
-class RenderLayer(BaseRenderLayer):
-	"""
-	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.
-
-	BaseRenderLayer provides a lower-level interface for manipulating the graphs
-	and lines directly, whereas RenderLayer provides an interface where you can
-	choose to target a specific IL level or a couple different Linear View specific
-	line constructs.
-	"""
-
 	def apply_to_disassembly_block(
 			self,
 			block: 'binaryninja.BasicBlock',
@@ -286,15 +204,15 @@ def apply_to_high_level_il_block(
 		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_hlil_body` instead as they \
+		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: Lines of text for the block, to be modified by this function
 		"""
 		pass
 
-	def apply_to_hlil_body(
+	def apply_to_high_level_il_body(
 			self,
 			function: 'binaryninja.Function',
 			lines: List['binaryninja.LinearDisassemblyLine']
@@ -338,7 +256,7 @@ def apply_to_block(
 			lines: List['binaryninja.DisassemblyTextLine'],
 	):
 		"""
-		Apply to lines generated by a Basic Block, of any type. If not overriden, this
+		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
@@ -358,12 +276,16 @@ def apply_to_block(
 
 	def apply_to_flow_graph(self, graph: 'binaryninja.FlowGraph') -> None:
 		"""
-		Overridden from BaseRenderLayer to provide more specific callbacks. Check
-		:py:func:`BaseRenderLayer.apply_to_flow_graph` in case you want to
-		override this directly.
+		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):
@@ -378,9 +300,12 @@ def apply_to_linear_view_object(
 			lines: List['binaryninja.LinearDisassemblyLine']
 	) -> List['binaryninja.LinearDisassemblyLine']:
 		"""
-		Overridden from BaseRenderLayer to provide more specific callbacks. Check
-		:py:func:`BaseRenderLayer.apply_to_linear_view_object` in case you want to
-		override this directly.
+		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
@@ -389,12 +314,12 @@ def apply_to_linear_view_object(
 		:return: Updated set of lines to display in Linear View
 		"""
 		# Hack: HLIL bodies don't have basic blocks
-		if obj.identifier.name in [
+		if len(lines) > 0 and obj.identifier.name in [
 			"HLIL Function Body",
 			"HLIL SSA Function Body",
 			"Language Representation Function Body"
 		]:
-			self.apply_to_hlil_body(lines[0].function, lines)
+			self.apply_to_high_level_il_body(lines[0].function, lines)
 			return lines
 
 		block_lines = []
@@ -407,8 +332,8 @@ def apply_to_linear_view_object(
 					if last_block is not None:
 						# Convert linear lines to disassembly lines for the apply()
 						# and then convert back for linear view
-						disasm_lines = [line.contents for line in block_lines]
-						disasm_map = {id(line.contents): line for line in block_lines}
+						disasm_lines = [block_line.contents for block_line in block_lines]
+						disasm_map = {id(block_line.contents): block_line for block_line in block_lines}
 						self.apply_to_block(last_block, disasm_lines)
 						block_lines = [
 							LinearDisassemblyLine(
@@ -445,3 +370,43 @@ def apply_to_linear_view_object(
 				self.apply_to_misc_linear_lines(obj, prev, next, block_lines)
 		final_lines += block_lines
 		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