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/highlevelil.py b/python/highlevelil.py
index 0041c780..e2922b46 100644
--- a/python/highlevelil.py
+++ b/python/highlevelil.py
@@ -20,7 +20,7 @@
 
 import ctypes
 import struct
-from typing import Optional, Generator, List, Union, NewType, Tuple, ClassVar, Mapping, Set, Callable
+from typing import Optional, Generator, List, Union, NewType, Tuple, ClassVar, Mapping, Set, Callable, Any
 from dataclasses import dataclass
 from enum import Enum
 
@@ -791,12 +791,10 @@ class HighLevelILInstruction(BaseILInstruction):
 		if cb(name, self, "HighLevelILInstruction", parent) == False:
 			return False
 		for name, op, opType in self.detailed_operands:
-			if opType == "HighLevelILInstruction":
-				assert isinstance(op, HighLevelILInstruction)
+			if isinstance(op, HighLevelILInstruction):
 				if not op.visit_all(cb, name, self):
 					return False
-			elif opType == "List[HighLevelILInstruction]":
-				assert isinstance(op, list) and all(isinstance(i, HighLevelILInstruction) for i in op)
+			elif isinstance(op, list) and all(isinstance(i, HighLevelILInstruction) for i in op):
 				for i in op:
 					if not i.visit_all(cb, name, self): # type: ignore
 						return False
@@ -813,12 +811,10 @@ class HighLevelILInstruction(BaseILInstruction):
 		:return: True if all instructions were visited, False if the callback returned False
 		"""
 		for name, op, opType in self.detailed_operands:
-			if opType == "HighLevelILInstruction":
-				assert isinstance(op, HighLevelILInstruction)
+			if isinstance(op, HighLevelILInstruction):
 				if not op.visit_operands(cb, name, self):
 					return False
-			elif opType == "List[HighLevelILInstruction]":
-				assert isinstance(op, list) and all(isinstance(i, HighLevelILInstruction) for i in op)
+			elif isinstance(op, list) and all(isinstance(i, HighLevelILInstruction) for i in op):
 				for i in op:
 					if not i.visit_operands(cb, name, self): # type: ignore
 						return False
@@ -847,18 +843,67 @@ class HighLevelILInstruction(BaseILInstruction):
 		"""
 		if cb(name, self, "HighLevelILInstruction", parent) == False:
 			return False
-		for name, op, opType in self.detailed_operands:
-			if opType == "HighLevelILInstruction":
-				assert isinstance(op, HighLevelILInstruction)
+		for name, op, _ in self.detailed_operands:
+			if isinstance(op, HighLevelILInstruction):
 				if not op.visit(cb, name, self):
 					return False
-			elif opType == "List[HighLevelILInstruction]":
-				assert isinstance(op, list) and all(isinstance(i, HighLevelILInstruction) for i in op)
+			elif isinstance(op, list) and all(isinstance(i, HighLevelILInstruction) 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[['HighLevelILInstruction', Any], Any], *args: Any, **kwargs: Any) -> Any:
+		"""
+		Traverses all HighLevelILInstructions 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.hlil_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.hlil_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, HighLevelILInstruction):
+				if (result := op.traverse(cb, *args, **kwargs)) is not None:
+					return result
+			elif isinstance(op, list) and all(isinstance(i, HighLevelILInstruction) for i in op):
+				for i in op:
+					if (result := i.traverse(cb, *args, **kwargs)) is not None:
+						return result
+		return None
+
 
 @dataclass(frozen=True, repr=False, eq=False)
 class HighLevelILUnaryBase(HighLevelILInstruction, UnaryOperation):