From d527d6f7d3cdc7c938a6f20fdab7e9cc1b59c920 Mon Sep 17 00:00:00 2001 From: KyleMiles Date: Tue, 10 Oct 2023 13:34:33 -0400 Subject: 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 --- python/mediumlevelil.py | 73 +++++++++++++++++++++++++++++++++++++++---------- 1 file changed, 59 insertions(+), 14 deletions(-) (limited to 'python/mediumlevelil.py') 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)""" -- cgit v1.3.1