Python: Add LLILFunction.translate and sample workflow using it

2 files changed, 181 insertions, 0 deletions

diff --git a/python/examples/wf_rewrite_push_pop.py b/python/examples/wf_rewrite_push_pop.py
new file mode 100644
index 00000000..8ae65212
--- /dev/null
+++ b/python/examples/wf_rewrite_push_pop.py
@@ -0,0 +1,150 @@
+import json
+from binaryninja import Workflow, Activity, AnalysisContext
+from binaryninja.lowlevelil import *
+
+
+def rewrite_action(context: AnalysisContext):
+    def translate_instr(
+            new_func: LowLevelILFunction,
+            old_block: LowLevelILBasicBlock,
+            old_instr: LowLevelILInstruction
+    ) -> ExpressionIndex:
+        """
+        Copy and translate ``old_instr`` in ``old_block`` into ``new_func``
+
+        :param new_func: new function to receive translated instructions
+        :param old_block: original block containing old_instr
+        :param old_instr: original instruction
+        :return: expression index of newly created instruction in ``new_func``
+        """
+
+        if old_instr.operation == LowLevelILOperation.LLIL_PUSH:
+            # push(x)
+            # -----------------------
+            # sp = sp - sizeof(void*)
+            # *(sp) = x
+
+            old_instr: LowLevelILPush
+            # sp = sp - sizeof(void*)
+            new_func.append(
+                new_func.set_reg(
+                    old_instr.size,
+                    old_block.arch.stack_pointer,
+                    new_func.sub(
+                        old_instr.size,
+                        new_func.reg(
+                            old_instr.size,
+                            old_block.arch.stack_pointer,
+                            loc=ILSourceLocation.from_instruction(old_instr)
+                        ),
+                        new_func.const(
+                            old_instr.size,
+                            old_block.arch.address_size,
+                            loc=ILSourceLocation.from_instruction(old_instr)
+                        ),
+                        loc=ILSourceLocation.from_instruction(old_instr)
+                    ),
+                    loc=ILSourceLocation.from_instruction(old_instr)
+                )
+            )
+            # *(sp) = x
+            return new_func.store(
+                old_instr.size,
+                new_func.reg(
+                    old_instr.size,
+                    old_block.arch.stack_pointer,
+                    loc=ILSourceLocation.from_instruction(old_instr)
+                ),
+                old_instr.src.copy_to(new_func),
+                loc=ILSourceLocation.from_instruction(old_instr)
+            )
+        elif old_instr.operation == LowLevelILOperation.LLIL_POP:
+            # pop
+            # -----------------------
+            # sp = sp + sizeof(void*)
+            # *(sp - sizeof(void*))
+
+            # We need to append any helper instructions first and then return an expression
+            # that replaces the ``pop`` in the original IL (since ``pop`` has a value).
+            # So anything that is ``rax = pop`` becomes ``sp = sp + 8 ; rax = *(sp - 8)``
+
+            old_instr: LowLevelILPop
+            # sp = sp + sizeof(void*)
+            new_func.append(
+                new_func.set_reg(
+                    old_instr.size,
+                    old_block.arch.stack_pointer,
+                    new_func.add(
+                        old_instr.size,
+                        new_func.reg(
+                            old_instr.size,
+                            old_block.arch.stack_pointer,
+                            loc=ILSourceLocation.from_instruction(old_instr)
+                        ),
+                        new_func.const(
+                            old_instr.size,
+                            old_block.arch.address_size,
+                            loc=ILSourceLocation.from_instruction(old_instr)
+                        ),
+                        loc=ILSourceLocation.from_instruction(old_instr)
+                    ),
+                    loc=ILSourceLocation.from_instruction(old_instr)
+                )
+            )
+            # *(sp - sizeof(void*))
+            return new_func.load(
+                old_instr.size,
+                new_func.sub(
+                    old_instr.size,
+                    new_func.reg(
+                        old_instr.size,
+                        old_block.arch.stack_pointer,
+                        loc=ILSourceLocation.from_instruction(old_instr)
+                    ),
+                    new_func.const(
+                        old_instr.size,
+                        old_block.arch.address_size,
+                        loc=ILSourceLocation.from_instruction(old_instr)
+                    ),
+                    loc=ILSourceLocation.from_instruction(old_instr)
+                ),
+                loc=ILSourceLocation.from_instruction(old_instr)
+            )
+        else:
+            # All other instructions: copy as-is
+            return old_instr.copy_to(
+                new_func,
+                lambda sub_instr: translate_instr(new_func, old_block, sub_instr)
+            )
+
+    # Modify the existing Lifted IL function by our translator above
+    translated_func = context.lifted_il.translate(translate_instr)
+    # Clean up blocks and prepare this function for the rest of analysis
+    translated_func.finalize()
+    # Tell the analysis to use the new form of this function
+    context.lifted_il = translated_func
+
+
+# Create and register the workflow for translating these instructions
+wf = Workflow("core.function.metaAnalysis").clone("RewritePushPop")
+
+# Define the custom activity configuration
+wf.register_activity(Activity(
+    configuration=json.dumps({
+        "name": "extension.rewrite_push_pop.rewrite_action",
+        "title": "Rewrite LLIL_PUSH/LLIL_POP",
+        "description": "Rewrites LLIL_PUSH/LLIL_POP instructions into their component store/load/register parts, demonstrating modifying and inserting Lifted IL instructions.",
+        "eligibility": {
+            "auto": {
+                "default": True
+            }
+        }
+    }),
+    action=rewrite_action
+))
+
+# This action is run right after generateLiftedIL so we can poke the IL before LLIL flag and stack
+# adjustment resolution happens.
+# core.function.generateLiftedIL -> (this) -> core.function.resetIndirectBranchesOnFullUpdate
+wf.insert("core.function.resetIndirectBranchesOnFullUpdate", ["extension.rewrite_push_pop.rewrite_action"])
+wf.register()
diff --git a/python/lowlevelil.py b/python/lowlevelil.py
index b8fe0b94..f37e53a2 100644
--- a/python/lowlevelil.py
+++ b/python/lowlevelil.py
@@ -4126,6 +4126,37 @@ class LowLevelILFunction:
 
 		raise NotImplementedError(f"unknown expr operation {expr.operation} in copy_expr_to")
 
+	def translate(
+		self, expr_handler: Callable[['LowLevelILFunction', 'LowLevelILBasicBlock', 'LowLevelILInstruction'], ExpressionIndex]
+	) -> 'LowLevelILFunction':
+		"""
+		``translate`` clones an IL function and modifies its expressions as specified by
+		a given ``expr_handler``, returning the updated IL function.
+
+		:param expr_handler: Function to modify an expression and copy it to the new function.
+		                     The function should have the following signature:
+
+		                     .. function:: expr_handler(new_func: LowLevelILFunction, old_block: LowLevelILBasicBlock, old_instr: LowLevelILInstruction) -> ExpressionIndex
+
+		                     Where:
+		                         - **new_func** (*LowLevelILFunction*): New function to receive translated instructions
+		                         - **old_block** (*LowLevelILBasicBlock*): Original block containing old_instr
+		                         - **old_instr** (*LowLevelILInstruction*): Original instruction
+		                         - **returns** (*ExpressionIndex*): Expression index of newly created instruction in ``new_func``
+		:return: Cloned IL function with modifications
+		"""
+
+		propagated_func = LowLevelILFunction(self.arch, source_func=self.source_function)
+		propagated_func.prepare_to_copy_function(self)
+		for block in self.basic_blocks:
+			propagated_func.prepare_to_copy_block(block)
+			for instr_index in range(block.start, block.end):
+				instr: LowLevelILInstruction = self[InstructionIndex(instr_index)]
+				propagated_func.set_current_address(instr.address, block.arch)
+				propagated_func.append(expr_handler(propagated_func, block, instr))
+
+		return propagated_func
+
 	def set_expr_attributes(self, expr: InstructionOrExpression, value: ILInstructionAttributeSet):
 		"""
 		``set_expr_attributes`` allows modification of instruction attributes but ONLY during lifting.