From cc92959fcc4163777a4db4f85c3dcc950cc32340 Mon Sep 17 00:00:00 2001
From: GGyul-E <39231485+GGyul-E@users.noreply.github.com>
Date: Mon, 14 Dec 2020 16:12:51 +0900
Subject: [PATCH] Added comment command  (#857)

* comment function is added

* Added file exception handling

* Added checking filename

* Update pwndbg/commands/comments.py

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

* Update pwndbg/commands/comments.py

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

* Change reading time for better performance

* Fixed exception handling

* small style fixes

Co-authored-by: Disconnect3d <dominik.b.czarnota@gmail.com>
---
 pwndbg/__init__.py          |  4 +++
 pwndbg/color/context.py     |  4 +++
 pwndbg/commands/comments.py | 50 +++++++++++++++++++++++++++++++++++++
 pwndbg/commands/nearpc.py   |  7 ++++++
 4 files changed, 65 insertions(+)
 create mode 100644 pwndbg/commands/comments.py

diff --git a/pwndbg/__init__.py b/pwndbg/__init__.py
index 06378a0e3..a649c71a9 100755
--- a/pwndbg/__init__.py
+++ b/pwndbg/__init__.py
@@ -50,6 +50,7 @@ import pwndbg.commands.vmmap
 import pwndbg.commands.windbg
 import pwndbg.commands.xinfo
 import pwndbg.commands.xor
+import pwndbg.commands.comments
 import pwndbg.constants
 import pwndbg.disasm
 import pwndbg.disasm.arm
@@ -158,3 +159,6 @@ signal.signal(signal.SIGWINCH, lambda signum, frame: gdb.execute("set width %i"
 # After GDB gets the fix, we should disable this only for bugged GDB versions.
 if 1:
     gdb.execute('set remote search-memory-packet off')
+
+# Reading Comment file 
+pwndbg.commands.comments.init()
diff --git a/pwndbg/color/context.py b/pwndbg/color/context.py
index 8ec8fe303..e4038d0cc 100644
--- a/pwndbg/color/context.py
+++ b/pwndbg/color/context.py
@@ -17,6 +17,7 @@ config_banner_color             = theme.ColoredParameter('banner-color', 'blue',
 config_banner_title             = theme.ColoredParameter('banner-title-color', 'none', 'color for banner title')
 config_register_changed_color   = theme.ColoredParameter('context-register-changed-color', 'normal', 'color for registers label (change marker)')
 config_register_changed_marker  = theme.Parameter('context-register-changed-marker', '*', 'change marker for registers label')
+config_comment                  = theme.ColoredParameter('comment-color', 'gray', 'color for comment')
 
 def prefix(x):
     return generateColorFunction(config.code_prefix_color)(x)
@@ -51,6 +52,9 @@ def banner(x):
 def banner_title(x):
     return generateColorFunction(config.banner_title_color)(x)
 
+def comment(x):
+    return generateColorFunction(config.comment_color)(x)
+
 def format_flags(value, flags, last=None):
     desc = flag_value('%#x' % value)
     if not flags:
diff --git a/pwndbg/commands/comments.py b/pwndbg/commands/comments.py
new file mode 100644
index 000000000..7f8a44b55
--- /dev/null
+++ b/pwndbg/commands/comments.py
@@ -0,0 +1,50 @@
+import argparse
+import pwndbg.commands
+from pwndbg.color import message
+
+parser = argparse.ArgumentParser(description="Put comments in assembly code")
+parser.add_argument("--addr", metavar='address', default=None, type=str, help="Address to write comments")
+parser.add_argument("comment", type=str, default=None,  help="The text you want to comment")
+
+file_lists = {} # This saves all comments.
+
+@pwndbg.commands.ArgparsedCommand(parser)
+@pwndbg.commands.OnlyWhenRunning
+def comm(addr=None, comment=None):
+    if addr is None:
+        addr = hex(pwndbg.regs.pc)
+    try: 
+        with open(".gdb_comments", "a+") as f:
+            target = int(addr,0)
+
+            if not pwndbg.memory.peek(target):
+                print(message.error("Invalid Address %#x" % target))
+            
+            else:
+                f.write("file:%s=" % pwndbg.proc.exe)
+                f.write("%#x:%s\n" % (target, comment))
+                if not pwndbg.proc.exe in file_lists.keys():
+                    file_lists[pwndbg.proc.exe] = {}
+                file_lists[pwndbg.proc.exe][hex(target)] = comment
+    except:
+        print(message.error("Permission denied to create file"))
+
+def init():
+    try:
+        with open(".gdb_comments","r") as f:
+            text = f.read()
+            text = text.split("\n")
+            for i in range(len(text)-1):
+                text1, text2 = text[i].split("=")
+
+                # split Filename, comments
+                filename = text1.split(":")[1]
+                addr_comm = text2.split(":")
+
+                if not filename in file_lists:
+                    file_lists[filename] = {}
+
+                file_lists[filename][addr_comm[0]] = addr_comm[1]
+
+    except:
+        pass 
diff --git a/pwndbg/commands/nearpc.py b/pwndbg/commands/nearpc.py
index 63840dbea..1a9099150 100644
--- a/pwndbg/commands/nearpc.py
+++ b/pwndbg/commands/nearpc.py
@@ -22,6 +22,7 @@ import pwndbg.strings
 import pwndbg.symbol
 import pwndbg.ui
 import pwndbg.vmmap
+import pwndbg.commands.comments
 from pwndbg.color import message
 
 
@@ -162,6 +163,12 @@ def nearpc(pc=None, lines=None, to_string=False, emulate=False):
             if syscall_name:
                 line += ' <%s>' % N.syscall_name(syscall_name)
 
+        # For Comment Function
+        try:
+            line += " "*10 + C.comment(pwndbg.commands.comments.file_lists[pwndbg.proc.exe][hex(instr.address)])
+        except:
+            pass
+
         result.append(line)
 
         # For call instructions, attempt to resolve the target and