fix: mprotect command and add tests for it

Turns out the mprotect command didn't ever work, as it was amd64 only, but used x86 syscall numbers to call mprotect. I have refactored the command to use shellcraft to generate the shellcode that calls mprotect. I have also unit-tested this command.

pwndbg/commands/mprotect.py

@@ -1,7 +1,5 @@
 import argparse
-
 import gdb
-
 import pwndbg.chain
 import pwndbg.commands
 import pwndbg.enhance
@@ -9,6 +7,8 @@ import pwndbg.file
 import pwndbg.lib.which
 import pwndbg.wrappers.checksec
 import pwndbg.wrappers.readelf
+import pwnlib
+from pwnlib import asm
 
 parser = argparse.ArgumentParser(description="Calls mprotect. x86_64 only.")
 parser.add_argument("addr", help="Page-aligned address to all mprotect on.", type=int)
@@ -51,24 +51,22 @@ def mprotect(addr, length, prot):
     saved_rdx = pwndbg.gdblib.regs.rdx
     saved_rip = pwndbg.gdblib.regs.rip
 
+
     prot_int = prot_str_to_val(prot)
-    gdb.execute("set $rax={}".format(SYS_MPROTECT))
-    gdb.execute("set $rbx={}".format(addr))
-    gdb.execute("set $rcx={}".format(length))
-    gdb.execute("set $rdx={}".format(prot_int))
 
-    saved_instruction_2bytes = pwndbg.gdblib.memory.read(pwndbg.gdblib.regs.rip, 2)
+    shellcode_asm = pwnlib.shellcraft.syscall("SYS_mprotect", int(addr), int(length), int(prot_int))
+    shellcode = asm.asm(shellcode_asm)
 
-    # int 0x80
-    pwndbg.gdblib.memory.write(pwndbg.gdblib.regs.rip, b"\xcd\x80")
+    saved_instruction_bytes = pwndbg.gdblib.memory.read(pwndbg.gdblib.regs.rip, len(shellcode))
 
+    pwndbg.gdblib.memory.write(pwndbg.gdblib.regs.rip, shellcode)
+   
     # execute syscall
+    gdb.execute("nextsyscall")
     gdb.execute("stepi")
 
-    print("mprotect returned {}".format(pwndbg.gdblib.regs.rax))
-
     # restore registers and memory
-    pwndbg.gdblib.memory.write(saved_rip, saved_instruction_2bytes)
+    pwndbg.gdblib.memory.write(saved_rip, saved_instruction_bytes)
 
     gdb.execute("set $rax={}".format(saved_rax))
     gdb.execute("set $rbx={}".format(saved_rbx))

pwndbg/gdblib/arch.py

@@ -3,11 +3,22 @@ import gdb
 import pwndbg.proc
 from pwndbg.gdblib import typeinfo
 from pwndbg.lib.arch import Arch
-
+import pwnlib
 # TODO: x86-64 needs to come before i386 in the current implementation, make
 # this order-independent
 ARCHS = ("x86-64", "i386", "aarch64", "mips", "powerpc", "sparc", "arm")
 
+# mapping between gdb and pwntools arch names
+pwnlib_archs_mapping = {
+    "x86-64": "amd64",
+    "i386": "i386",
+    "aarch64": "aarch64",
+    "mips": "mips",
+    "powerpc": "powerpc",
+    "sparc": "sparc",
+    "arm": "arm",
+}
+
 arch = Arch("i386", typeinfo.ptrsize, "little")
 
 
@@ -45,3 +56,5 @@ def update():
     # object. Instead, we call `__init__` again with the new args
     arch_name, ptrsize, endian = _get_arch(typeinfo.ptrsize)
     arch.__init__(arch_name, ptrsize, endian)
+    pwnlib.context.context.arch = pwnlib_archs_mapping[arch_name]
+    pwnlib.context.context.bits = ptrsize * 8

tests/binaries/mprotect.c

@@ -0,0 +1,17 @@
+#include <stdlib.h>
+#include <stdint.h>
+
+// 
+// pwnlib.asm.asm(pwnlib.shellcraft.amd64.linux.echo("mprotect_ok"), arch='amd64').hex()
+uint8_t func[] = {
+    0x36, 0x38, 0x35, 0x65, 0x36, 0x65, 0x36, 0x61, 0x30, 0x31, 0x38, 0x31, 0x33, 0x34, 0x32, 0x34, 0x30, 0x31, 0x30, 0x31, 0x30, 0x31, 0x30, 0x31, 0x34, 0x38, 0x62, 0x38, 0x36, 0x64, 0x37, 0x30, 0x37, 0x32, 0x36, 0x66, 0x37, 0x34, 0x36, 0x35, 0x36, 0x33, 0x37, 0x34, 0x35, 0x30, 0x36, 0x61, 0x30, 0x31, 0x35, 0x38, 0x36, 0x61, 0x30, 0x31, 0x35, 0x66, 0x36, 0x61, 0x30, 0x62, 0x35, 0x61, 0x34, 0x38, 0x38, 0x39, 0x65, 0x36, 0x30, 0x66, 0x30, 0x35, 0x33, 0x31, 0x66, 0x66, 0x36, 0x61, 0x33, 0x63, 0x35, 0x38, 0x30, 0x66, 0x30, 0x35
+};
+
+
+int main() {
+
+    void (*f)() = (void (*)())func;
+    printf("ptr = %p", f);
+    // f();
+    return 1;
+}

tests/test_mprotect.py

@@ -0,0 +1,48 @@
+import gdb
+
+import pwndbg
+import tests
+
+MPROTECT_BINARY = tests.binaries.get("mprotect.out")
+
+def test_mprotect(start_binary):
+    """
+    Tests mprotect command
+    It will mark some memory as executable, then this binary will print "mprotect_ok"
+    """
+    start_binary(MPROTECT_BINARY)
+
+    gdb.execute("starti")
+    # get addr of func
+    addr = int(gdb.parse_and_eval("&func"))
+    addr_aligned = pwndbg.lib.memory.page_align(addr)
+  
+    # sizeof
+    size = int(gdb.parse_and_eval("sizeof(func)"))
+    size_aligned = pwndbg.lib.memory.page_align(size)
+
+    vmmaps_before = gdb.execute("vmmap -x", to_string=True).splitlines()
+   
+    # mark memory as executable
+    gdb.execute("mprotect {} {} PROT_EXEC|PROT_READ|PROT_WRITE".format(hex(addr_aligned), pwndbg.lib.memory.PAGE_SIZE))
+   
+    vmmaps_after = gdb.execute("vmmap -x", to_string=True).splitlines()
+
+    # expect vmmaps_after to be one element longer than vmmaps_before
+    assert len(vmmaps_after) == len(vmmaps_before) + 1
+
+    # get the changed vmmap entry
+    vmmap_entry = [x for x in vmmaps_after if x not in vmmaps_before][0]
+
+    assert vmmap_entry.split()[2] == "rwxp"
+   
+    # continue execution
+    gdb.execute("continue")
+    
+
+
+def test_cannot_run_mprotect_when_not_running(start_binary):
+
+
+    # expect error message
+    assert "mprotect: The program is not being run.\n" == gdb.execute("mprotect 0x0 0x1000 PROT_EXEC|PROT_READ|PROT_WRITE", to_string=True)