From 446f5d5ae25265fbc17667b202def7423fcc8754 Mon Sep 17 00:00:00 2001
From: Disconnect3d <dominik.b.czarnota@gmail.com>
Date: Mon, 19 Sep 2022 15:30:45 -0700
Subject: [PATCH] add patch command (#1150)

* add patch command

This commit adds the `patch`, `patch_list` and `patch_revert` commands
and adds the `pwntools==4.8.0` as Pwndbg dependency.

The current implementation could be further improved by:
- adding tests :)
- maybe moving `patch_list` and `patch_revert` to `patch --list` and
  `patch --revert` flags?
- better handling of incorrect args/pwnlib exceptions

* lint
---
 gdbinit.py               |  2 +
 pwndbg/__init__.py       |  1 +
 pwndbg/commands/patch.py | 81 ++++++++++++++++++++++++++++++++++++++++
 requirements.txt         |  3 +-
 4 files changed, 86 insertions(+), 1 deletion(-)
 create mode 100644 pwndbg/commands/patch.py

diff --git a/gdbinit.py b/gdbinit.py
index 984b74f1..ae87a9b5 100644
--- a/gdbinit.py
+++ b/gdbinit.py
@@ -70,4 +70,6 @@ if encoding != "UTF-8":
     print("Make sure that en_US.UTF-8 is activated in /etc/locale.gen and you called locale-gen")
     print("******")
 
+environ["PWNLIB_NOTERM"] = "1"
+
 import pwndbg  # noqa: F401
diff --git a/pwndbg/__init__.py b/pwndbg/__init__.py
index 82871586..e3f00726 100755
--- a/pwndbg/__init__.py
+++ b/pwndbg/__init__.py
@@ -31,6 +31,7 @@ import pwndbg.commands.misc
 import pwndbg.commands.mprotect
 import pwndbg.commands.next
 import pwndbg.commands.p2p
+import pwndbg.commands.patch
 import pwndbg.commands.peda
 import pwndbg.commands.pie
 import pwndbg.commands.probeleak
diff --git a/pwndbg/commands/patch.py b/pwndbg/commands/patch.py
new file mode 100644
index 00000000..11210c4a
--- /dev/null
+++ b/pwndbg/commands/patch.py
@@ -0,0 +1,81 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+import argparse
+
+from pwnlib.asm import asm
+from pwnlib.asm import disasm
+
+import pwndbg.color.message as message
+import pwndbg.commands
+import pwndbg.gdblib.memory
+import pwndbg.lib.memoize
+
+# Keep old patches made so we can revert them
+patches = {}
+
+
+parser = argparse.ArgumentParser(description="Patches given instruction with given code or bytes")
+parser.add_argument("address", type=int, help="The address to patch")
+parser.add_argument("ins", type=str, help="instruction[s]")
+
+
+@pwndbg.commands.ArgparsedCommand(parser)
+@pwndbg.commands.OnlyWhenRunning
+def patch(address, ins):
+    new_mem = asm(ins, arch=pwndbg.gdblib.arch.current)
+
+    old_mem = pwndbg.gdblib.memory.read(address, len(new_mem))
+
+    patches[address] = (old_mem, new_mem)
+
+    pwndbg.gdblib.memory.write(address, new_mem)
+
+    pwndbg.lib.memoize.reset()
+
+
+parser2 = argparse.ArgumentParser(description="Revert patch at given address")
+parser2.add_argument("address", type=int, help="Address to revert patch on")
+
+
+@pwndbg.commands.ArgparsedCommand(parser2)
+@pwndbg.commands.OnlyWhenRunning
+def patch_revert(address):
+    if not patches:
+        print(message.notice("No patches to revert"))
+        return
+
+    if address == -1:
+        for addr, (old, _new) in patches.items():
+            pwndbg.gdblib.memory.write(addr, old)
+            print(message.notice("Reverted patch at %#x" % addr))
+        patches.clear()
+    else:
+        old, _new = patches[address]
+        pwndbg.gdblib.memory.write(address, old)
+
+    pwndbg.lib.memoize.reset()
+
+
+parser3 = argparse.ArgumentParser(description="List all patches")
+
+
+@pwndbg.commands.ArgparsedCommand(parser3)
+@pwndbg.commands.OnlyWhenRunning
+def patch_list():
+    if not patches:
+        print(message.hint("No patches to list"))
+        return
+
+    print(message.hint("Patches:"))
+    for addr, (old, new) in patches.items():
+        old_insns = disasm(old, arch=pwndbg.gdblib.arch.current)
+        new_insns = disasm(new, arch=pwndbg.gdblib.arch.current)
+
+        print(
+            message.hint("Patch at"),
+            message.warning("%#x:" % addr),
+            message.hint("from"),
+            message.warning(old_insns.replace("\n", "; ")),
+            message.hint("to"),
+            message.warning(new_insns.replace("\n", "; ")),
+        )
diff --git a/requirements.txt b/requirements.txt
index 0e84fb31..f93a96c3 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -20,4 +20,5 @@ testresources==2.0.1
 tomli==2.0.1; python_version >= '3.7'
 tomli==1.2.3; python_version < '3.7'
 unicorn==2.0.0; python_version >= '3.7'
-unicorn==2.0.0rc7; python_version < '3.7'
\ No newline at end of file
+unicorn==2.0.0rc7; python_version < '3.7'
+pwntools==4.8.0