Python API : Add a more flexible/ergonomic visitor pattern called traverse to LLIL, MLIL, and HLIL

Also cleans up some old code patterns in the old visitors

1 files changed, 59 insertions, 14 deletions

diff --git a/python/mediumlevelil.py b/python/mediumlevelil.py
index 9e8fb3b4..31c4d363 100644
--- a/python/mediumlevelil.py
+++ b/python/mediumlevelil.py
@@ -21,7 +21,7 @@
 import ctypes
 import struct
 from typing import (Optional, List, Union, Mapping,
-	Generator, NewType, Tuple, ClassVar, Dict, Set, Callable)
+	Generator, NewType, Tuple, ClassVar, Dict, Set, Callable, Any)
 from dataclasses import dataclass
 
 # Binary Ninja components
@@ -437,12 +437,10 @@ class MediumLevelILInstruction(BaseILInstruction):
 		if cb(name, self, "MediumLevelILInstruction", parent) == False:
 			return False
 		for name, op, opType in self.detailed_operands:
-			if opType == "MediumLevelILInstruction":
-				assert isinstance(op, MediumLevelILInstruction)
+			if isinstance(op, MediumLevelILInstruction):
 				if not op.visit_all(cb, name, self):
 					return False
-			elif opType == "List[MediumLevelILInstruction]":
-				assert isinstance(op, list) and all(isinstance(i, MediumLevelILInstruction) for i in op)
+			elif isinstance(op, list) and all(isinstance(i, MediumLevelILInstruction) for i in op):
 				for i in op:
 					if not i.visit_all(cb, name, self): # type: ignore
 						return False
@@ -459,12 +457,10 @@ class MediumLevelILInstruction(BaseILInstruction):
 		:return: True if all instructions were visited, False if the callback returned False
 		"""
 		for name, op, opType in self.detailed_operands:
-			if opType == "MediumLevelILInstruction":
-				assert isinstance(op, MediumLevelILInstruction)
+			if isinstance(op, MediumLevelILInstruction):
 				if not op.visit_operands(cb, name, self):
 					return False
-			elif opType == "List[MediumLevelILInstruction]":
-				assert isinstance(op, list) and all(isinstance(i, MediumLevelILInstruction) for i in op)
+			elif isinstance(op, list) and all(isinstance(i, MediumLevelILInstruction) for i in op):
 				for i in op:
 					if not i.visit_operands(cb, name, self): # type: ignore
 						return False
@@ -493,18 +489,67 @@ class MediumLevelILInstruction(BaseILInstruction):
 		"""
 		if cb(name, self, "MediumLevelILInstruction", parent) == False:
 			return False
-		for name, op, opType in self.detailed_operands:
-			if opType == "MediumLevelILInstruction":
-				assert isinstance(op, MediumLevelILInstruction)
+		for name, op, _ in self.detailed_operands:
+			if isinstance(op, MediumLevelILInstruction):
 				if not op.visit(cb, name, self):
 					return False
-			elif opType == "List[MediumLevelILInstruction]":
-				assert isinstance(op, list) and all(isinstance(i, MediumLevelILInstruction) for i in op)
+			elif isinstance(op, list) and all(isinstance(i, MediumLevelILInstruction) for i in op):
 				for i in op:
 					if not i.visit(cb, name, self): # type: ignore
 						return False
 		return True
 
+	def traverse(self, cb: Callable[['MediumLevelILInstruction', Any], Any], *args: Any, **kwargs: Any) -> Any:
+		"""
+		Traverses all MediumLevelILInstructions in the operands of this instruction and any sub-instructions.
+		The callback you provide only needs to accept a single instruction, but accepts anything, and can return whatever you want.
+
+		None is treated as a reserved value to indicate that the traverser should continue descending into subexpressions.
+
+		:param cb: Callback function that takes only the instruction
+		:param args: Custom user-defined arguments
+		:param kwargs: Custom user-defined keyword arguments
+		:return: None if your callback doesn't return anything and all instructions were traversed, otherwise it returns the value from your callback.
+		:Example:
+		>>> # This traverser allows for simplified function signatures in your callback
+		>>> def traverser(inst) -> int:
+		>>>  if isinstance(inst, Constant):
+		>>>   return inst.constant # Stop recursion and return the constant
+		>>>  return None # Continue descending into subexpressions
+
+		>>> # Finds all constants used in the program
+		>>> for inst in bv.mlil_instructions:
+		>>>  if const := inst.traverse(traverser):
+		>>>   print(f"Found constant {const}")
+
+
+		>>> # But it also allows for complex function signatures in your callback
+		>>> def traverser(inst, search_constant, skip_list: List[int] = []) -> int:
+		>>>  if inst.address in skip_list:
+		>>>   return None # Skip this instruction
+		>>>  if isinstance(inst, Constant):
+		>>>   if inst.constant == search_constant:
+		>>>    return inst.address # Stop recursion and return the address of this use
+		>>>  return None # Continue descending into subexpressions
+
+		>>> # Finds all instances of 0xdeadbeaf used in the program
+		>>> for inst in bv.mlil_instructions:
+		>>>  if use_addr := inst.traverse(traverser, 0xdeadbeaf, skip_list=[0x12345678]):
+		>>>   print(f"Found 0xdeadbeef use at {use_addr}")
+		"""
+
+		if (result := cb(self, *args, **kwargs)) is not None:
+			return result
+		for _, op, _ in self.detailed_operands:
+			if isinstance(op, MediumLevelILInstruction):
+				if (result := op.traverse(cb, *args, **kwargs)) is not None:
+					return result
+			elif isinstance(op, list) and all(isinstance(i, MediumLevelILInstruction) for i in op):
+				for i in op:
+					if (result := i.traverse(cb, *args, **kwargs)) is not None:
+						return result
+		return None
+
 	@property
 	def tokens(self) -> TokenList:
 		"""MLIL tokens (read-only)"""