Add `stepuntilasm` command (#1798)

* Add `stepuntilasm` command

This commit adds a `stepuntilasm` command that, given a mnemonic and,
optionally, a set of operands, will step until a instruction that
matches both is found. Matching is string-based, as the user will likely
want to spell out the asm directive they want as text, and interpreting
assembly language conventions for all of the platforms pwndbg supports
is probably outside the scope of this change.

* next.py: small code cleanup

* next.py: fix bug introduced in previous commit

op.str -> op_str

* Update next.py

* Update next.py

* Update next.py

---------

Co-authored-by: Disconnect3d <dominik.b.czarnota@gmail.com>

pwndbg/commands/next.py

@@ -132,3 +132,22 @@ def stepsyscall() -> None:
 
     if pwndbg.gdblib.proc.alive:
         pwndbg.commands.context.context()
+
+
+parser = argparse.ArgumentParser(description="Breaks on the next matching instruction.")
+parser.add_argument("mnemonic", type=str, help="The mnemonic of the instruction")
+parser.add_argument(
+    "op_str",
+    type=str,
+    nargs="*",
+    help="The operands of the instruction",
+)
+
+
+@pwndbg.commands.ArgparsedCommand(parser, command_name="stepuntilasm")
+@pwndbg.commands.OnlyWhenRunning
+def stepuntilasm(mnemonic, op_str):
+    if len(op_str) == 0:
+        op_str = None
+
+    pwndbg.gdblib.next.break_on_next_matching_instruction(mnemonic, op_str)

pwndbg/gdblib/next.py

@@ -4,6 +4,7 @@ instruction of some type (call, branch, etc.)
 """
 
 import re
+from itertools import chain
 
 import capstone
 import gdb
@@ -14,9 +15,9 @@ import pwndbg.gdblib.proc
 import pwndbg.gdblib.regs
 from pwndbg.color import message
 
-jumps = set((capstone.CS_GRP_CALL, capstone.CS_GRP_JUMP, capstone.CS_GRP_RET, capstone.CS_GRP_IRET))
+jumps = {capstone.CS_GRP_CALL, capstone.CS_GRP_JUMP, capstone.CS_GRP_RET, capstone.CS_GRP_IRET}
 
-interrupts = set((capstone.CS_GRP_INT,))
+interrupts = {capstone.CS_GRP_INT}
 
 
 def clear_temp_breaks() -> None:
@@ -42,9 +43,10 @@ def next_int(address=None):
 
     ins = pwndbg.disasm.one(address)
     while ins:
-        if set(ins.groups) & jumps:
+        ins_groups = set(ins.groups)
+        if ins_groups & jumps:
             return None
-        if set(ins.groups) & interrupts:
+        elif ins_groups & interrupts:
             return ins
         ins = pwndbg.disasm.one(ins.next)
 
@@ -67,6 +69,48 @@ def next_branch(address=None):
     return None
 
 
+def next_matching_until_branch(address=None, mnemonic=None, op_str=None):
+    """
+    Finds the next instruction that matches the arguments between the given
+    address and the branch closest to it.
+    """
+    if address is None:
+        address = pwndbg.gdblib.regs.pc
+
+    ins = pwndbg.disasm.one(address)
+    while ins:
+        # Check whether or not the mnemonic matches if it was specified
+        mnemonic_match = ins.mnemonic.casefold() == mnemonic.casefold() if mnemonic else True
+
+        # Check whether or not the operands match if they were specified
+        op_str_match = True
+        if op_str is not None:
+            op_str_match = False
+
+            # Remove whitespace and fold the case of both targets.
+            ops = "".join(ins.op_str.split()).casefold()
+            if isinstance(op_str, str):
+                op_str = "".join(op_str.split()).casefold()
+            elif isinstance(op_str, list):
+                op_str = "".join(chain.from_iterable(op.split() for op in op_str)).casefold()
+            else:
+                raise ValueError("op_str value is of an unsupported type")
+            op_str_match = ops == op_str
+
+        # If all of the parameters that were specified match, this is the
+        # instruction we want to stop at.
+        if mnemonic_match and op_str_match:
+            return ins
+
+        if set(ins.groups) & jumps:
+            # No matching instruction until the next branch, and we're
+            # not trying to match the branch instruction itself.
+            return None
+
+        ins = pwndbg.disasm.one(ins.next)
+    return None
+
+
 def break_next_branch(address=None):
     ins = next_branch(address)
 
@@ -128,6 +172,52 @@ def break_next_ret(address=None):
             return ins
 
 
+def break_on_next_matching_instruction(mnemonic=None, op_str=None) -> bool:
+    """
+    Breaks on next instuction that matches the arguments.
+    """
+    # Make sure we have something to break on.
+    if mnemonic is None and op_str is None:
+        return False
+
+    while pwndbg.gdblib.proc.alive:
+        # Break on signal as it may be a segfault
+        if pwndbg.gdblib.proc.stopped_with_signal:
+            return False
+
+        ins = next_matching_until_branch(mnemonic=mnemonic, op_str=op_str)
+        if ins is not None:
+            if ins.address != pwndbg.gdblib.regs.pc:
+                print("Found instruction")
+                # Only set breakpoints at a different PC location, otherwise we
+                # will continue until we hit a breakpoint that's not related to
+                # this opeeration, or the program halts.
+                gdb.Breakpoint("*%#x" % ins.address, internal=True, temporary=True)
+                gdb.execute("continue", from_tty=False, to_string=True)
+                return ins
+            else:
+                # We don't want to be spinning in place, nudge execution forward
+                # and try again.
+                pass
+        else:
+            # Move to the next branch instruction.
+            print("Moving to next branch")
+            nb = next_branch(pwndbg.gdblib.regs.pc)
+            if nb is not None:
+                if nb.address != pwndbg.gdblib.regs.pc:
+                    # Stop right at the next branch instruction.
+                    gdb.Breakpoint("*%#x" % nb.address, internal=True, temporary=True)
+                    gdb.execute("continue", from_tty=False, to_string=True)
+                else:
+                    # Nudge execution so we take the branch we're on top of.
+                    pass
+
+        if pwndbg.gdblib.proc.alive:
+            gdb.execute("si")
+
+    return False
+
+
 def break_on_program_code() -> bool:
     """
     Breaks on next instruction that belongs to process' objfile code

tests/gdb-tests/tests/binaries/stepuntilasm_x64.asm

@@ -0,0 +1,45 @@
+section .text
+    global _start
+    global break_here
+    global stop1
+    global stop2
+    global stop3
+    global stop4
+_start: 
+break_here:
+    xor rax, rax
+stop1: 
+    nop ; Stop point #1: No operands
+stop2: 
+    xor rax, rax ; Stop point #2: Some simple operands
+    
+    lea rax, [some_data]
+stop3: 
+    ; Stop point #3: More complex operands.
+    mov qword [rax], 0x20 
+    
+    call loop
+    lea rax, [some_data]
+stop4: 
+    ; Stop point #4: Even more complex operands, after loop.
+    mov dword [rax+4], 0x20 
+
+exit:
+    ; Terminate the process by calling sys_exit(0) in Linux.
+    mov rax, 60
+    mov rdi, 0
+    syscall
+
+
+; Loop subroutine. Loops for a while so we can test whether stepuntilasm can get 
+; to a directive that's sitting after a few iterations of a loop.
+loop:
+    mov rax, 100
+loop_iter:
+    sub rax, 1    
+    jnz loop_iter
+
+    ret 
+
+section .bss
+some_data: resq 1

tests/gdb-tests/tests/test_command_stepuntilasm.py

@@ -0,0 +1,23 @@
+import gdb
+
+import pwndbg.gdblib
+import tests
+
+STEPUNTILASM_X64_BINARY = tests.binaries.get("stepuntilasm_x64.out")
+
+
+def test_command_untilasm_x64(start_binary):
+    start_binary(STEPUNTILASM_X64_BINARY)
+    gdb.execute("break break_here")
+    gdb.execute("run")
+
+    run_and_verify("stop1", "nop")
+    run_and_verify("stop2", "xor rax, rax")
+    run_and_verify("stop3", "mov qword ptr [rax], 0x20")
+    run_and_verify("stop4", "mov dword ptr [rax+4], 0x20")
+
+
+def run_and_verify(stop_label, asm):
+    gdb.execute(f"stepuntilasm {asm}")
+    address = int(gdb.parse_and_eval(f"&{stop_label}"))
+    assert pwndbg.gdblib.regs.pc == address