Render Layers

2 files changed, 422 insertions, 0 deletions

diff --git a/python/examples/args_render_layer.py b/python/examples/args_render_layer.py
new file mode 100644
index 00000000..bac0889e
--- /dev/null
+++ b/python/examples/args_render_layer.py
@@ -0,0 +1,213 @@
+import functools
+from typing import List, Mapping, Tuple, Iterator
+
+from binaryninja import DisassemblyTextLine, LowLevelILInstruction, LowLevelILOperation, \
+    TypeClass, DisassemblyTextRenderer, MediumLevelILFunction, \
+    MediumLevelILCallSsa, MediumLevelILVarSsa, MediumLevelILConstBase, \
+    MediumLevelILInstruction, MediumLevelILTailcallSsa, MediumLevelILOperation, \
+    MediumLevelILVarPhi, log_debug, RenderLayer, BasicBlock, InstructionTextTokenType, \
+    RenderLayerDefaultEnableState
+
+"""
+Render Layer that shows you where the arguments to calls are set, for Disasm/LLIL.
+- Adds "Argument '<arg>' for call at <callsite>" comments to lines that set up call args
+- Adds "Call at <callsite>" comments to the call sites
+
+=========================================================================================
+
+But how do you determine the argument to a call?
+What seems like it has worked:
+- You can't determine that an instruction is a parameter, you have to go from the call to its parameters
+- Since trying to look up the call for an instruction is impossible, instead go through every call at once for a function (and memoize it)
+- LLIL is useless for looking this up, since it has no types and the call parameters often include a list of every register
+- Using MLIL, we can find all of the parameters as MLIL instructions, but we need to map them to LLIL so we can use them in the LLIL/Disasm display
+- How do we map them? Turns out that's rather inconvenient:
+  - Register arguments are generally pretty easy because they are just MLIL vars
+  - Stack arguments somehow also work out generally, the .llil on the MLIL points to the push()
+  - Constants are a mess and I just use the MLIL's address (this is often incorrect)
+  - Flags are completely unhandled for now
+  - Phis are handled by just looking up every var they use... probably not proper but sort of works
+- This fails in a couple of scenarios though, notably __builtin_xxxxxx() functions
+  - Which instruction specifies the length of a group of `mov qword [rbx+8], rax {0}` calls? I think it just picks one?
+  - The `rep` instructions could actually have these params resolved (they use real registers) but in practice this doesn't work
+  - Thunks are unhandled
+"""
+
+
+@functools.lru_cache(maxsize=64)
+def get_param_sites(mlil: MediumLevelILFunction) -> Mapping[LowLevelILInstruction, List[Tuple[MediumLevelILInstruction, int]]]:
+    """
+    For a given function, find all LLIL instructions that are parameters to a call,
+    and return a mapping for each instruction with all the calls that it maps to,
+    their corresponding MLIL call instruction, and which numbered parameter they are
+    in the call.
+
+    :param mlil: MLIL function to search
+    :return: Map of param sites as described above
+    """
+    call_sites = {}
+    mlil = mlil.ssa_form
+
+    # As a function to handle call and tailcall identically
+    def collect_call_params(call_site, dest, params):
+        def_sites = []
+        for i, param in enumerate(params):
+            llil = param.llil
+            if llil is not None:
+                def_sites.append((param, llil))
+                continue
+
+            match param:
+                case MediumLevelILVarSsa(src=var_src):
+                    def_site = mlil.get_ssa_var_definition(var_src)
+                    if def_site is not None and def_site.llil is not None:
+                        def_sites.append((i, def_site.llil))
+                        continue
+                    # Handle phis by just looking up the def sites of all their sources
+                    match def_site:
+                        case MediumLevelILVarPhi(src=phis):
+                            for phi in phis:
+                                phi_def = mlil.get_ssa_var_definition(phi)
+                                if phi_def is not None and phi_def.llil is not None:
+                                    def_sites.append((i, phi_def.llil))
+                case MediumLevelILConstBase():
+                    # This is wrong, but it works (sometimes)
+                    # Oh god, have I just quoted php.net
+                    def_site_idx = mlil.llil.get_instruction_start(param.address)
+                    if def_site_idx is not None:
+                        def_sites.append((i, mlil.llil[def_site_idx].ssa_form))
+                        continue
+
+            if len(def_sites) == 0:
+                log_debug(f"Could not find def site for param {i} in call at {call_site.address:#x}")
+
+        call_sites[call_site] = def_sites
+
+    for instr in mlil.instructions:
+        match instr:
+            case MediumLevelILCallSsa(dest=dest, params=params) as call_site:
+                collect_call_params(call_site, dest, params)
+            case MediumLevelILTailcallSsa(dest=dest, params=params) as call_site:
+                collect_call_params(call_site, dest, params)
+
+    # Inverse args
+    all_def_sites = {}
+    for call_site, params in call_sites.items():
+        for (param_idx, llil) in params:
+            if llil not in all_def_sites:
+                all_def_sites[llil] = []
+            else:
+                print(f"got two at {llil.instr_index} @ {llil.address:#x} -> {call_site.address:#x}")
+            all_def_sites[llil].append((call_site, param_idx))
+
+    return all_def_sites
+
+
+def get_llil_arg(llil: LowLevelILInstruction) -> Iterator[Tuple[str, MediumLevelILInstruction]]:
+    args = get_param_sites(llil.function.mlil)
+
+    if llil.ssa_form in args:
+        for call_site, param_idx in args[llil.ssa_form]:
+            target_type = call_site.function.get_expr_type(call_site.dest.expr_index)
+
+            # Try getting the param name from the call's type
+            if target_type is not None:
+                if target_type.type_class == TypeClass.PointerTypeClass:
+                    target_type = target_type.target
+                if target_type.type_class == TypeClass.FunctionTypeClass:
+                    target_params = target_type.parameters
+                    if param_idx < len(target_params):
+                        param_name = target_params[param_idx].name
+                        if param_name == '':
+                            param_name = f"arg{param_idx+1}"
+                        yield param_name, call_site
+                        continue
+
+            # Some calls have extra params that aren't reflected in their type
+            yield f"arg{param_idx+1}", call_site
+    return
+
+
+def apply_to_lines(lines, get_instr, renderer):
+    # So we don't process lines twice since we're iterating over a list as we modify it
+    skip_lines = []
+
+    # Tailcalls that don't return incorrectly mark the { Does not return } line as a call
+    ignore_calls = set()
+
+    for i, line in enumerate(lines):
+        if len(line.tokens) == 0:
+            continue
+        if i in skip_lines:
+            continue
+
+        llil_instr = get_instr(line)
+        if llil_instr is not None:
+            new_lines = []
+            for (arg, call) in get_llil_arg(llil_instr):
+                if call.operation == MediumLevelILOperation.MLIL_TAILCALL_SSA:
+                    comment = f"Argument '{arg}' for tailcall at {call.address:#x}"
+                else:
+                    comment = f"Argument '{arg}' for call at {call.address:#x}"
+                renderer.wrap_comment(new_lines, line, comment, False, "  ",  "")
+                for j, token in enumerate(line.tokens):
+                    if token.type == InstructionTextTokenType.AddressSeparatorToken:
+                        line.tokens = line.tokens[:j]
+                        break
+
+            # Annotate calls too so we can see them easily next to their args
+            if llil_instr.address == line.address and llil_instr.address not in ignore_calls:
+                if llil_instr.operation in [
+                    LowLevelILOperation.LLIL_CALL,
+                    LowLevelILOperation.LLIL_CALL_SSA,
+                    LowLevelILOperation.LLIL_TAILCALL,
+                    LowLevelILOperation.LLIL_TAILCALL_SSA
+                ]:
+                    ignore_calls.add(llil_instr.address)
+                    if llil_instr.operation in [
+                        LowLevelILOperation.LLIL_TAILCALL,
+                        LowLevelILOperation.LLIL_TAILCALL_SSA
+                    ]:
+                        comment = f"Tailcall at {llil_instr.address:#x}"
+                    else:
+                        comment = f"Call at {llil_instr.address:#x}"
+                    # Creating comments is a bit unwieldy at the moment
+                    renderer.wrap_comment(new_lines, line, comment, False, "  ", "")
+                    for j, token in enumerate(line.tokens):
+                        if token.type == InstructionTextTokenType.AddressSeparatorToken:
+                            line.tokens = line.tokens[:j]
+                            break
+
+            # If any of our lines changed, swap out the existing lines with the new ones
+            if len(new_lines) > 0:
+                lines.pop(i)
+                for j, new_line in enumerate(new_lines):
+                    lines.insert(i + j, new_line)
+                    skip_lines.append(i + j)
+    return lines
+
+
+class ArgumentsRenderLayer(RenderLayer):
+    name = "Annotate Call Parameters"
+    default_enable_state = RenderLayerDefaultEnableState.EnabledByDefaultRenderLayerDefaultEnableState
+
+    def apply_to_disassembly_block(
+            self,
+            block: BasicBlock,
+            lines: List['DisassemblyTextLine']
+    ):
+        # Break this out into a helper so we don't have to write it twice
+        renderer = DisassemblyTextRenderer(block.function)
+        return apply_to_lines(lines, lambda line: block.function.get_llil_at(line.address), renderer)
+
+    def apply_to_low_level_il_block(
+            self,
+            block: BasicBlock,
+            lines: List['DisassemblyTextLine']
+    ):
+        # Break this out into a helper so we don't have to write it twice
+        renderer = DisassemblyTextRenderer(block.function)
+        return apply_to_lines(lines, lambda line: line.il_instruction, renderer)
+
+
+ArgumentsRenderLayer.register()
diff --git a/python/examples/follow_reg_render_layer.py b/python/examples/follow_reg_render_layer.py
new file mode 100644
index 00000000..36e66576
--- /dev/null
+++ b/python/examples/follow_reg_render_layer.py
@@ -0,0 +1,209 @@
+from typing import List
+
+from PySide6.QtCore import QSettings
+from binaryninja import DisassemblyTextLine, LowLevelILOperation, DisassemblyTextRenderer, \
+    MediumLevelILOperation, \
+    RenderLayer, BasicBlock, InstructionTextTokenType, RenderLayerDefaultEnableState, \
+    PluginCommand, BinaryView, interaction, InstructionTextToken, RegisterValueType
+
+
+"""
+Render layer that adds annotations to follow the value of a register, for Disasm/LLIL.
+- Adds "<reg> after: <value>" comments to every line that modifies your register
+- Adds "<reg> at block entry: <value>" comments to every block
+- Adds "<reg> at block exit: <value>" comments to every block
+You can switch registers with the plugin command "Set Followed Register"
+"""
+
+
+class FollowRegRenderLayer(RenderLayer):
+    name = "Follow Register"
+    default_enable_state = RenderLayerDefaultEnableState.EnabledByDefaultRenderLayerDefaultEnableState
+
+    def __init__(self, handle=None):
+        super().__init__(handle)
+        self.followed_reg = QSettings().value("layer/followed_reg", None)
+
+    def apply_to_lines(
+            self,
+            block: BasicBlock,
+            lines: List['DisassemblyTextLine'],
+    ):
+        if self.followed_reg is None:
+            return lines
+        if self.followed_reg not in block.arch.regs:
+            return lines
+
+        func = block.function
+        arch = block.arch
+
+        # Give our tokens a unique value so they don't highlight all the other comments
+        token_value = arch.get_reg_index(self.followed_reg) ^ 0x12345678
+
+        # ===============================================================================
+        # Render the "<reg> after" comments
+
+        for i, line in enumerate(lines):
+            # Ignore non-disassembly lines
+            if not any(token.type == InstructionTextTokenType.AddressSeparatorToken for token in line.tokens):
+                continue
+
+            # If the register has changed, we should render the update
+            written = False
+            for w in func.get_regs_written_by(line.address, arch):
+                # In case this instruction only modifies a sub register of our target
+                # Or in the case that we wanted to follow a sub register
+                if arch.regs[w].full_width_reg == arch.regs[self.followed_reg].full_width_reg:
+                    written = True
+                    break
+
+            if written:
+                # Add comment tokens to existing line
+                line.tokens.extend([
+                    InstructionTextToken(InstructionTextTokenType.CommentToken, '  // ', token_value),
+                    InstructionTextToken(InstructionTextTokenType.CommentToken, self.followed_reg, token_value),
+                    InstructionTextToken(InstructionTextTokenType.CommentToken, ' after: ', token_value),
+                ])
+
+                after = func.get_reg_value_after(line.address, self.followed_reg, arch)
+                if after.type == RegisterValueType.UndeterminedValue:
+                    if line.il_instruction is not None:
+                        instr = line.il_instruction
+                    else:
+                        instr_start = func.low_level_il.get_instruction_start(line.address, arch)
+                        instr = func.low_level_il[instr_start]
+                    if instr is not None:
+                        after_possible = instr.get_possible_reg_values_after(self.followed_reg)
+                        line.tokens.append(
+                            InstructionTextToken(InstructionTextTokenType.CommentToken, str(after_possible), token_value)
+                        )
+                    else:
+                        line.tokens.append(
+                            InstructionTextToken(InstructionTextTokenType.CommentToken, str(after), token_value)
+                        )
+                else:
+                    line.tokens.append(
+                        InstructionTextToken(InstructionTextTokenType.CommentToken, str(after), token_value)
+                    )
+
+        block_start = block.start
+        block_end = block.end
+        if block.is_low_level_il:
+            block_start = func.llil[block.start].address
+            block_end = func.llil[block.end - 1].address
+
+        # ===============================================================================
+        # Render the before-block "<reg> at block entry" comment
+
+        start_before = func.get_reg_value_at(block_start, self.followed_reg, arch)
+        line = [
+            InstructionTextToken(InstructionTextTokenType.CommentToken, '// ', token_value),
+            InstructionTextToken(InstructionTextTokenType.CommentToken, self.followed_reg, token_value),
+            InstructionTextToken(InstructionTextTokenType.CommentToken, ' at block entry: ', token_value),
+        ]
+
+        # Sometimes the first line is blank and we want to insert after it
+        first_line = 0
+        if len(lines[0].tokens) == 0:
+            first_line = 1
+        if start_before.type == RegisterValueType.UndeterminedValue:
+            if lines[first_line].il_instruction is not None:
+                instr = lines[first_line].il_instruction
+            else:
+                instr_start = func.low_level_il.get_instruction_start(block_start, arch)
+                instr = func.low_level_il[instr_start]
+            if instr is not None:
+                start_before_possible = instr.get_possible_reg_values(self.followed_reg)
+                line.append(
+                    InstructionTextToken(InstructionTextTokenType.CommentToken, str(start_before_possible), token_value)
+                )
+            else:
+                line.append(
+                    InstructionTextToken(InstructionTextTokenType.CommentToken, str(start_before), token_value)
+                )
+        else:
+            line.append(
+                InstructionTextToken(InstructionTextTokenType.CommentToken, str(start_before), token_value)
+            )
+
+        lines.insert(first_line, DisassemblyTextLine(line, block_start))
+
+        # ===============================================================================
+        # Render the after-block "<reg> at block exit" comment
+
+        end_after = func.get_reg_value_after(block_end, self.followed_reg, arch)
+        line = [
+            InstructionTextToken(InstructionTextTokenType.CommentToken, '// ', token_value),
+            InstructionTextToken(InstructionTextTokenType.CommentToken, self.followed_reg, token_value),
+            InstructionTextToken(InstructionTextTokenType.CommentToken, ' at block exit: ', token_value),
+        ]
+
+        # Sometimes the last line is blank and we want to insert before it
+        if len(lines) > 1 and len(lines[-1].tokens) == 0:
+            last_line = len(lines) - 1
+        else:
+            last_line = len(lines)
+        if end_after.type == RegisterValueType.UndeterminedValue:
+            if lines[last_line - 1].il_instruction is not None:
+                instr = lines[last_line - 1].il_instruction
+            else:
+                instr_start = func.low_level_il.get_instruction_start(block_end, arch)
+                instr = func.low_level_il[instr_start]
+            if instr is not None:
+                end_after_possible = instr.get_possible_reg_values_after(self.followed_reg)
+                line.append(
+                    InstructionTextToken(InstructionTextTokenType.CommentToken, str(end_after_possible), token_value)
+                )
+            else:
+                line.append(
+                    InstructionTextToken(InstructionTextTokenType.CommentToken, str(end_after), token_value)
+                )
+        else:
+            line.append(
+                InstructionTextToken(InstructionTextTokenType.CommentToken, str(end_after), token_value)
+            )
+
+        lines.insert(last_line, DisassemblyTextLine(line, block_end))
+
+        return lines
+
+    def apply_to_disassembly_block(
+            self,
+            block: BasicBlock,
+            lines: List['DisassemblyTextLine']
+    ):
+        # Break this out into a helper so we don't have to write it twice
+        return self.apply_to_lines(block, lines)
+
+    def apply_to_low_level_il_block(
+            self,
+            block: BasicBlock,
+            lines: List['DisassemblyTextLine']
+    ):
+        # Break this out into a helper so we don't have to write it twice
+        return self.apply_to_lines(block, lines)
+
+
+def set_follow_reg(bv: BinaryView):
+    if bv.platform is not None:
+        regs = list(bv.platform.arch.regs.keys())
+        idx = interaction.get_large_choice_input("Choose", "Choose Followed Register", regs)
+
+        # Save choice both in QSettings and on the RenderLayer object instance
+        layer = RenderLayer[FollowRegRenderLayer.name]
+        if idx is None:
+            layer.followed_reg = None
+            QSettings().remove("layer/followed_reg")
+        else:
+            layer.followed_reg = regs[idx]
+            QSettings().setValue("layer/followed_reg", regs[idx])
+
+    # Trigger view refresh (gross)
+    for func in bv.get_functions_containing(bv.offset):
+        func.reanalyze()
+
+
+FollowRegRenderLayer.register()
+
+# Using a command for this is kinda janky but currently the only option
+PluginCommand.register("Set Followed Register", "Choose which register to follow for the Follow Register Render Layer", set_follow_reg)