From b5e865430a62fa209a55517f9d29a6f18c586f3a Mon Sep 17 00:00:00 2001 From: OBarronCS <55004530+OBarronCS@users.noreply.github.com> Date: Tue, 24 Sep 2024 00:52:00 -0700 Subject: [PATCH] Make annotation reg/memory assignment symbol configurable (#2449) * Make annotation reg/memory assignment symbol configurable * Hardcode the separator symbol * Fix tests for store instructions --- pwndbg/aglib/disasm/aarch64.py | 3 +- pwndbg/aglib/disasm/arch.py | 42 +++++++++++++--- pwndbg/aglib/disasm/mips.py | 5 +- pwndbg/aglib/disasm/riscv.py | 9 ++-- pwndbg/aglib/disasm/x86.py | 53 +++++++++++++-------- tests/qemu-tests/tests/user/test_aarch64.py | 8 ++-- tests/qemu-tests/tests/user/test_arm.py | 8 ++-- tests/qemu-tests/tests/user/test_mips.py | 6 +-- tests/qemu-tests/tests/user/test_riscv64.py | 2 +- 9 files changed, 89 insertions(+), 47 deletions(-) diff --git a/pwndbg/aglib/disasm/aarch64.py b/pwndbg/aglib/disasm/aarch64.py index 6b529397..8248a5b9 100644 --- a/pwndbg/aglib/disasm/aarch64.py +++ b/pwndbg/aglib/disasm/aarch64.py @@ -14,6 +14,7 @@ import pwndbg.aglib.memory import pwndbg.aglib.regs import pwndbg.enhance import pwndbg.lib.disasm.helpers as bit_math +from pwndbg.aglib.disasm.arch import register_assign from pwndbg.aglib.disasm.instruction import ALL_JUMP_GROUPS from pwndbg.aglib.disasm.instruction import EnhancedOperand from pwndbg.aglib.disasm.instruction import InstructionCondition @@ -288,7 +289,7 @@ class DisassemblyAssistant(pwndbg.aglib.disasm.arch.DisassemblyAssistant): telescope = self._telescope_format_list(addresses, TELESCOPE_DEPTH, emu) - instruction.annotation = f"{result_operand.str} => {telescope}" + instruction.annotation = register_assign(result_operand.str, telescope) @override def _condition( diff --git a/pwndbg/aglib/disasm/arch.py b/pwndbg/aglib/disasm/arch.py index 57590350..d88618a6 100644 --- a/pwndbg/aglib/disasm/arch.py +++ b/pwndbg/aglib/disasm/arch.py @@ -126,6 +126,21 @@ DO_NOT_EMULATE = { } +def register_assign(left: str, right: str) -> str: + return f"{left} => {right}" + + +def memory_assign(left: str, right: str) -> str: + return f"{left} <= {right}" + + +def memory_or_register_assign(left: str, right: str, mem_assign: bool) -> str: + """ + Used when we don't know until runtime whether a codepath will annotate a register or memory location. + """ + return memory_assign(left, right) if mem_assign else register_assign(left, right) + + # Enhances disassembly with memory values & symbols by adding member variables to an instruction # The only public method that should be called is "enhance" # The enhance function is passed an instance of the Unicorn emulator @@ -809,7 +824,9 @@ class DisassemblyAssistant: if not telescope_addresses: return - instruction.annotation = f"{left.str} => {self._telescope_format_list(telescope_addresses, TELESCOPE_DEPTH, emu)}" + instruction.annotation = register_assign( + left.str, self._telescope_format_list(telescope_addresses, TELESCOPE_DEPTH, emu) + ) def _common_cmp_annotator_builder( self, flags_register_name: str, char_to_separate_operands: str = "-" @@ -845,7 +862,7 @@ class DisassemblyAssistant: emu_eflags = emu.read_register(flags_register_name) eflags_formatted = C.format_flags(emu_eflags, eflags_bits) - display_result = f"{FLAG_REG_NAME_DISPLAY} => {eflags_formatted}" + display_result = register_assign(FLAG_REG_NAME_DISPLAY, eflags_formatted) if instruction.annotation is None: # First part of this function usually sets .annotation to a string. But if the instruction @@ -928,7 +945,7 @@ class DisassemblyAssistant: instruction.annotation = f"{dest_str}, {source_str}" if telescope_print is not None: - instruction.annotation += f" => {telescope_print}" + instruction.annotation = register_assign(instruction.annotation, telescope_print) def _common_store_annotator( self, @@ -971,7 +988,9 @@ class DisassemblyAssistant: emu, ) - instruction.annotation = f"{address_str} => {self._telescope_format_list(telescope_addresses, TELESCOPE_DEPTH, emu)}" + instruction.annotation = memory_assign( + address_str, self._telescope_format_list(telescope_addresses, TELESCOPE_DEPTH, emu) + ) def _common_move_annotator(self, instruction: PwndbgInstruction, emu: Emulator): """ @@ -993,7 +1012,9 @@ class DisassemblyAssistant: if not telescope_addresses: return - instruction.annotation = f"{left.str} => {self._telescope_format_list(telescope_addresses, TELESCOPE_DEPTH, emu)}" + instruction.annotation = register_assign( + left.str, self._telescope_format_list(telescope_addresses, TELESCOPE_DEPTH, emu) + ) def _common_binary_op_annotator( self, @@ -1003,6 +1024,7 @@ class DisassemblyAssistant: op_one: int | None, op_two: int | None, char_to_separate_operands: str, + memory_assignment=False, ) -> None: # Ex: "0x198723 + 0x2b8" math_string = None @@ -1014,11 +1036,17 @@ class DisassemblyAssistant: # Using emulation, we can determine the resulting value if target_operand.after_value_resolved is not None: - instruction.annotation = f"{target_operand.str} => {MemoryColor.get_address_and_symbol(target_operand.after_value_resolved)}" + instruction.annotation = memory_or_register_assign( + target_operand.str, + MemoryColor.get_address_and_symbol(target_operand.after_value_resolved), + memory_assignment, + ) if math_string: instruction.annotation += f" ({math_string})" elif math_string: - instruction.annotation = f"{target_operand.str} => {math_string}" + instruction.annotation = memory_or_register_assign( + target_operand.str, math_string, memory_assignment + ) generic_assistant = DisassemblyAssistant(None) diff --git a/pwndbg/aglib/disasm/mips.py b/pwndbg/aglib/disasm/mips.py index 89d65217..ac0ea88b 100644 --- a/pwndbg/aglib/disasm/mips.py +++ b/pwndbg/aglib/disasm/mips.py @@ -24,6 +24,7 @@ from typing_extensions import override import pwndbg.aglib.disasm.arch import pwndbg.color.memory as MemoryColor import pwndbg.lib.disasm.helpers as bit_math +from pwndbg.aglib.disasm.arch import register_assign from pwndbg.aglib.disasm.instruction import FORWARD_JUMP_GROUP from pwndbg.aglib.disasm.instruction import InstructionCondition from pwndbg.aglib.disasm.instruction import PwndbgInstruction @@ -200,8 +201,8 @@ class DisassemblyAssistant(pwndbg.aglib.disasm.arch.DisassemblyAssistant): # Resolve it manually without emulation address = right.before_value << 16 - instruction.annotation = ( - f"{result_operand.str} => {MemoryColor.get_address_and_symbol(address)}" + instruction.annotation = register_assign( + result_operand.str, MemoryColor.get_address_and_symbol(address) ) @override diff --git a/pwndbg/aglib/disasm/riscv.py b/pwndbg/aglib/disasm/riscv.py index d63ceac3..d3d2c305 100644 --- a/pwndbg/aglib/disasm/riscv.py +++ b/pwndbg/aglib/disasm/riscv.py @@ -11,6 +11,7 @@ import pwndbg.aglib.disasm.arch import pwndbg.aglib.regs import pwndbg.color.memory as MemoryColor import pwndbg.lib.disasm.helpers as bit_math +from pwndbg.aglib.disasm.arch import register_assign from pwndbg.aglib.disasm.instruction import InstructionCondition from pwndbg.aglib.disasm.instruction import PwndbgInstruction @@ -217,8 +218,8 @@ class DisassemblyAssistant(pwndbg.aglib.disasm.arch.DisassemblyAssistant): # Resolve it manually without emulation address = instruction.address + (right.before_value << 12) - instruction.annotation = ( - f"{result_operand.str} => {MemoryColor.get_address_and_symbol(address)}" + instruction.annotation = register_assign( + result_operand.str, MemoryColor.get_address_and_symbol(address) ) def _lui_annotator(self, instruction: PwndbgInstruction, emu: Emulator) -> None: @@ -228,8 +229,8 @@ class DisassemblyAssistant(pwndbg.aglib.disasm.arch.DisassemblyAssistant): # Resolve it manually without emulation address = right.before_value << 12 - instruction.annotation = ( - f"{result_operand.str} => {MemoryColor.get_address_and_symbol(address)}" + instruction.annotation = register_assign( + result_operand.str, MemoryColor.get_address_and_symbol(address) ) def _resolve_compressed_target_addr( diff --git a/pwndbg/aglib/disasm/x86.py b/pwndbg/aglib/disasm/x86.py index 3ea8f37f..38dd9ad7 100644 --- a/pwndbg/aglib/disasm/x86.py +++ b/pwndbg/aglib/disasm/x86.py @@ -18,6 +18,8 @@ import pwndbg.chain import pwndbg.color.memory as MemoryColor import pwndbg.color.message as MessageColor import pwndbg.enhance +from pwndbg.aglib.disasm.arch import memory_or_register_assign +from pwndbg.aglib.disasm.arch import register_assign from pwndbg.aglib.disasm.instruction import EnhancedOperand from pwndbg.aglib.disasm.instruction import InstructionCondition from pwndbg.aglib.disasm.instruction import PwndbgInstruction @@ -90,6 +92,7 @@ class DisassemblyAssistant(pwndbg.aglib.disasm.arch.DisassemblyAssistant): instruction.operands[0].before_value_resolved, instruction.operands[1].before_value_resolved, X86_MATH_INSTRUCTIONS[instruction.id], + instruction.operands[0].type == CS_OP_MEM, ) else: self.annotation_handlers.get(instruction.id, lambda *a: None)(instruction, emu) @@ -121,19 +124,7 @@ class DisassemblyAssistant(pwndbg.aglib.disasm.arch.DisassemblyAssistant): ) elif left.type == CS_OP_REG and right.before_value is not None: # MOV REG, REG|IMM - TELESCOPE_DEPTH = max(0, int(pwndbg.config.disasm_telescope_depth)) - - telescope_addresses = super()._telescope( - right.before_value, - TELESCOPE_DEPTH + 1, - instruction, - emu, - read_size=right.cs_op.size, - ) - if not telescope_addresses: - return - - instruction.annotation = f"{left.str} => {super()._telescope_format_list(telescope_addresses, TELESCOPE_DEPTH, emu)}" + self._common_move_annotator(instruction, emu) def handle_vmovaps(self, instruction: PwndbgInstruction, emu: Emulator) -> None: # If the source or destination is in memory, it must be aligned to: @@ -168,7 +159,9 @@ class DisassemblyAssistant(pwndbg.aglib.disasm.arch.DisassemblyAssistant): telescope_addresses = super()._telescope( right.before_value, TELESCOPE_DEPTH, instruction, emu ) - instruction.annotation = f"{left.str} => {super()._telescope_format_list(telescope_addresses, TELESCOPE_DEPTH, emu)}" + instruction.annotation = register_assign( + left.str, super()._telescope_format_list(telescope_addresses, TELESCOPE_DEPTH, emu) + ) def handle_xchg(self, instruction: PwndbgInstruction, emu: Emulator) -> None: left, right = instruction.operands @@ -176,7 +169,19 @@ class DisassemblyAssistant(pwndbg.aglib.disasm.arch.DisassemblyAssistant): if left.before_value_resolved is not None and right.before_value_resolved is not None: # Display the exchanged values. Doing it this way (instead of using .after_value) allows this to work without emulation # Don't telescope here for the sake of screen space - instruction.annotation = f"{left.str} => {MemoryColor.get_address_or_symbol(right.before_value_resolved)}, {right.str} => {MemoryColor.get_address_or_symbol(left.before_value_resolved)}" + instruction.annotation = ( + memory_or_register_assign( + left.str, + MemoryColor.get_address_or_symbol(right.before_value_resolved), + left.type == CS_OP_MEM, + ) + + ", " + + memory_or_register_assign( + right.str, + MemoryColor.get_address_or_symbol(left.before_value_resolved), + right.type == CS_OP_MEM, + ) + ) def handle_pop(self, instruction: PwndbgInstruction, emu: Emulator) -> None: pc_is_at_instruction = self.can_reason_about_process_state(instruction) @@ -190,13 +195,15 @@ class DisassemblyAssistant(pwndbg.aglib.disasm.arch.DisassemblyAssistant): if reg_operand.type == CS_OP_REG: if emu and reg_operand.after_value is not None: # After emulation, the register has taken on the popped value - instruction.annotation = f"{reg_operand.str} => {MemoryColor.get_address_and_symbol(reg_operand.after_value)}" + instruction.annotation = register_assign( + reg_operand.str, MemoryColor.get_address_and_symbol(reg_operand.after_value) + ) elif pc_is_at_instruction: - # Attempt to read from the stop of the stack + # Attempt to read from the top of the stack try: value = pwndbg.aglib.memory.pvoid(pwndbg.aglib.regs.sp) - instruction.annotation = ( - f"{reg_operand.str} => {MemoryColor.get_address_and_symbol(value)}" + instruction.annotation = register_assign( + reg_operand.str, MemoryColor.get_address_and_symbol(value) ) except Exception: pass @@ -206,7 +213,7 @@ class DisassemblyAssistant(pwndbg.aglib.disasm.arch.DisassemblyAssistant): # If zeroing the register with XOR A, A. Can reason about this no matter where the instruction is if left.type == CS_OP_REG and right.type == CS_OP_REG and left.reg == right.reg: - instruction.annotation = f"{left.str} => 0" + instruction.annotation = register_assign(left.str, "0") else: self._common_binary_op_annotator( instruction, @@ -222,7 +229,11 @@ class DisassemblyAssistant(pwndbg.aglib.disasm.arch.DisassemblyAssistant): operand = instruction.operands[0] if operand.after_value_resolved is not None: - instruction.annotation = f"{operand.str} => {MemoryColor.get_address_and_symbol(operand.after_value_resolved)}" + instruction.annotation = memory_or_register_assign( + operand.str, + MemoryColor.get_address_and_symbol(operand.after_value_resolved), + operand.type == CS_OP_MEM, + ) def handle_dec(self, instruction: PwndbgInstruction, emu: Emulator) -> None: self.handle_inc(instruction, emu) diff --git a/tests/qemu-tests/tests/user/test_aarch64.py b/tests/qemu-tests/tests/user/test_aarch64.py index b9a98366..8c9f3dd0 100644 --- a/tests/qemu-tests/tests/user/test_aarch64.py +++ b/tests/qemu-tests/tests/user/test_aarch64.py @@ -341,13 +341,13 @@ def test_aarch64_store_instructions(qemu_assembly_run): "LEGEND: STACK | HEAP | CODE | DATA | WX | RODATA\n" "─────────────────────[ DISASM / aarch64 / set emulate on ]──────────────────────\n" " ► 0x10000028 ldr x4, #stores+56 X4, 0x10000060 => 0x4010e8 (value1) ◂— 0\n" - " 0x1000002c strb w0, [x4] [value1] => 0xf0\n" + " 0x1000002c strb w0, [x4] [value1] <= 0xf0\n" " 0x10000030 ldr x5, #stores+64 X5, 0x10000068 => 0x4010e9 (value2) ◂— 0\n" - " 0x10000034 strh w0, [x5] [value2] => 0xdef0\n" + " 0x10000034 strh w0, [x5] [value2] <= 0xdef0\n" " 0x10000038 ldr x6, #stores+72 X6, 0x10000070 => 0x4010eb (value4) ◂— 0\n" - " 0x1000003c str w0, [x6] [value4] => 0x9abcdef0\n" + " 0x1000003c str w0, [x6] [value4] <= 0x9abcdef0\n" " 0x10000040 ldr x7, #stores+80 X7, 0x10000078 => 0x4010ef (value8) ◂— 0\n" - " 0x10000044 str x0, [x7] [value8] => 0x123456789abcdef0\n" + " 0x10000044 str x0, [x7] [value8] <= 0x123456789abcdef0\n" " 0x10000048 mov x8, #0x5d X8 => 0x5d\n" " 0x1000004c mov x0, #0 X0 => 0\n" " 0x10000050 svc #0 \n" diff --git a/tests/qemu-tests/tests/user/test_arm.py b/tests/qemu-tests/tests/user/test_arm.py index 78363027..21bee3b6 100644 --- a/tests/qemu-tests/tests/user/test_arm.py +++ b/tests/qemu-tests/tests/user/test_arm.py @@ -383,7 +383,7 @@ def test_arm_stack_pointer_check(qemu_assembly_run): " 0x10000004 <_start+4> mov r1, #3 R1 => 3\n" " 0x10000008 <_start+8> add r2, r0, r1 R2 => 7 (4 + 3)\n" " 0x1000000c <_start+12> sub r3, r2, #2 R3 => 5 (7 - 2)\n" - f" 0x10000010 <_start+16> str r3, [sp, #-4]! [{hex(pwndbg.gdblib.regs.sp - 4)}] => 5\n" + f" 0x10000010 <_start+16> str r3, [sp, #-4]! [{hex(pwndbg.gdblib.regs.sp - 4)}] <= 5\n" " 0x10000014 <_start+20> pop {r4}\n" " 0x10000018 <_start+24> mul r4, r2, r1 R4 => 21 (7 * 3)\n" " 0x1000001c <_start+28> add r4, r4, #1 R4 => 22 (0x15 + 0x1)\n" @@ -521,9 +521,9 @@ def test_arm_exclusive_store(qemu_assembly_run): " ► 0x10000000 <_start> ldr r0, [pc, #0x34] R0, [_start+60] => 0x11094 (value1) ◂— 0\n" " 0x10000004 <_start+4> ldr r1, [pc, #0x34] R1, [_start+64] => 0x87654321\n" " 0x10000008 <_start+8> ldr r2, [pc, #0x34] R2, [_start+68] => 0x12345678\n" - " 0x1000000c <_start+12> str r1, [r0] [value1] => 0x87654321\n" - " 0x10000010 <_start+16> strex r3, r2, [r0] [value1] => 0x12345678\n" - " 0x10000014 <_start+20> str r1, [r0], #1 [value1] => 0x87654321\n" + " 0x1000000c <_start+12> str r1, [r0] [value1] <= 0x87654321\n" + " 0x10000010 <_start+16> strex r3, r2, [r0] [value1] <= 0x12345678\n" + " 0x10000014 <_start+20> str r1, [r0], #1 [value1] <= 0x87654321\n" " 0x10000018 <_start+24> add r0, r0, r1 R0 => 0x876653b6 (0x11095 + 0x87654321)\n" " 0x1000001c <_start+28> nop \n" " 0x10000020 <_start+32> nop \n" diff --git a/tests/qemu-tests/tests/user/test_mips.py b/tests/qemu-tests/tests/user/test_mips.py index 54c75136..bfef54b5 100644 --- a/tests/qemu-tests/tests/user/test_mips.py +++ b/tests/qemu-tests/tests/user/test_mips.py @@ -282,13 +282,13 @@ def test_mips32_store_instruction(qemu_assembly_run): " 0x10000004 <_start+4> ori $t0, $t0, 0x5678 T0 => 0x12345678 (0x12340000 | 0x5678)\n" " 0x10000008 <_start+8> lui $s0, 0x40 S0 => 0x400000\n" " 0x1000000c <_start+12> addiu $s0, $s0, 0x1130 S0 => 0x401130 (value1) (0x400000 + 0x1130)\n" - " 0x10000010 <_start+16> sw $t0, ($s0) [value1] => 0x12345678\n" + " 0x10000010 <_start+16> sw $t0, ($s0) [value1] <= 0x12345678\n" " 0x10000014 <_start+20> lui $s1, 0x40 S1 => 0x400000\n" " 0x10000018 <_start+24> addiu $s1, $s1, 0x1134 S1 => 0x401134 (value2) (0x400000 + 0x1134)\n" - " 0x1000001c <_start+28> sh $t0, ($s1) [value2] => 0x5678\n" + " 0x1000001c <_start+28> sh $t0, ($s1) [value2] <= 0x5678\n" " 0x10000020 <_start+32> lui $s2, 0x40 S2 => 0x400000\n" " 0x10000024 <_start+36> addiu $s2, $s2, 0x1136 S2 => 0x401136 (value3) (0x400000 + 0x1136)\n" - " 0x10000028 <_start+40> sb $t0, ($s2) [value3] => 0x78\n" + " 0x10000028 <_start+40> sb $t0, ($s2) [value3] <= 0x78\n" "────────────────────────────────────────────────────────────────────────────────\n" ) diff --git a/tests/qemu-tests/tests/user/test_riscv64.py b/tests/qemu-tests/tests/user/test_riscv64.py index 01379917..46f9c38b 100644 --- a/tests/qemu-tests/tests/user/test_riscv64.py +++ b/tests/qemu-tests/tests/user/test_riscv64.py @@ -120,7 +120,7 @@ def test_riscv64_compressed_loads(qemu_assembly_run): expected = ( "LEGEND: STACK | HEAP | CODE | DATA | WX | RODATA\n" "───────────────────────[ DISASM / rv64 / set emulate on ]───────────────────────\n" - " ► 0x10000028 c.sd a0, 0(a2) [data] => 0x1234567890abcdef\n" + " ► 0x10000028 c.sd a0, 0(a2) [data] <= 0x1234567890abcdef\n" " 0x1000002a c.ld a1, 0(a2) A1, [data] => 0x1234567890abcdef\n" " 0x1000002c c.li a1, 0x10 A1 => 0x10\n" " 0x1000002e addi a2, zero, 0x26 A2 => 38 (0x0 + 0x26)\n"