Fix xor and memfrob (#1057)

* added basic xor and memfrob tests

* refactor xor and memforb + add more tests

pwndbg/commands/xor.py

@@ -18,49 +18,25 @@ def xor_memory(address, key, count):
 
     return mem
 
+
 parser = argparse.ArgumentParser(description="XOR `count` bytes at address` with the key key`.")
-parser.add_argument("address", type=int, help="The address to start xoring at.")
+parser.add_argument("address", type=pwndbg.commands.sloppy_gdb_parse, help="The address to start xoring at.")
 parser.add_argument("key", type=str, help="The key to use.")
 parser.add_argument("count", type=int, help="The number of bytes to xor.")
 @pwndbg.commands.ArgparsedCommand(parser)
 @pwndbg.commands.OnlyWhenRunning
 def xor(address, key, count):
-    '''xor(address, key, count)
-
-    XOR ``count`` bytes at ``address`` with the key ``key``.
-    '''
-    if not isinstance(address, int):
-        try:
-            address = int(address, 0)
-        except ValueError:
-            print('Invalid address %s' % address)
-            return
-
     try:
         xorred_memory = xor_memory(address, key, count)
         pwndbg.memory.write(address, xorred_memory)
     except gdb.error as e:
         print(e)
 
-parser = argparse.ArgumentParser(description="Memfrobs a region of memory.")
+
+parser = argparse.ArgumentParser(description="Memfrobs a region of memory (xor with '*').")
 parser.add_argument("address", type=int, help="The address to start xoring at.")
 parser.add_argument("count", type=int, help="The number of bytes to xor.")
 @pwndbg.commands.ArgparsedCommand(parser)
 @pwndbg.commands.OnlyWhenRunning
 def memfrob(address, count):
-    '''memfrob(address, count)
-
-    Run the memfrob command on a region of memory
-    '''
-    if not isinstance(address, int):
-        try:
-            address = int(address, 0)
-        except ValueError:
-            print('Invalid address %s' % address)
-            return
-
-    try:
-        xorred_memory = xor_memory(address, '*', count)
-        pwndbg.memory.write(address, xorred_memory)
-    except gdb.error as e:
-        print(e)
+    return xor(address, '*', count)

tests/test_command_xor.py

@@ -0,0 +1,60 @@
+import gdb
+
+import pwndbg.memory
+import pwndbg.regs
+import tests
+from pwndbg.commands.xor import memfrob
+from pwndbg.commands.xor import xor
+
+REFERENCE_BINARY = tests.binaries.get("reference-binary.out")
+
+
+def test_command_xor_with_gdb_execute(start_binary):
+    """
+    Tests simple xoring
+    """
+    start_binary(REFERENCE_BINARY)
+
+    before = pwndbg.regs.rsp
+    pwndbg.memory.write(before, b"aaaaaaaa")
+    gdb.execute("xor $rsp ' ' 4")
+    after = pwndbg.memory.read(before, 8)
+    assert after == b"AAAAaaaa"
+
+def test_command_xor_with_int(start_binary):
+    """
+    Tests simple xoring
+    """
+    start_binary(REFERENCE_BINARY)
+
+    before = pwndbg.regs.rsp
+    assert isinstance(before, int)
+    pwndbg.memory.write(before, b"aaaaaaaa")
+    gdb.execute(f"xor {before} ' ' 4")
+    after = pwndbg.memory.read(before, 8)
+    assert after == b"AAAAaaaa"
+
+
+def test_command_xor_with_hex(start_binary):
+    """
+    Tests simple xoring
+    """
+    start_binary(REFERENCE_BINARY)
+
+    before = pwndbg.regs.rsp
+    before_hex = hex(before)
+    assert isinstance(before_hex, str)
+    pwndbg.memory.write(before, b"aaaaaaaa")
+    gdb.execute(f"xor {before_hex} ' ' 4")
+    after = pwndbg.memory.read(before, 8)
+    assert after == b"AAAAaaaa"
+
+
+def test_command_memfrob(start_binary):
+    start_binary(REFERENCE_BINARY)
+
+    before = pwndbg.regs.rsp
+    pwndbg.memory.write(before, b"aaaaaaaa")
+    memfrob(before, 4)
+    after = pwndbg.memory.read(before, 8)
+    assert after == b"KKKKaaaa"