From 5028f69ce5c716dcd0ec20666ea1111a2f7d1fcc Mon Sep 17 00:00:00 2001
From: Levente Polyak <levente@leventepolyak.net>
Date: Wed, 15 Jun 2016 22:50:58 +0200
Subject: [PATCH] highlighting the current IP or line of the code/source
 section (#49)

---
 pwndbg/color.py            | 38 ++++++++++++++++++++++++++------------
 pwndbg/commands/context.py | 12 +++++++++++-
 pwndbg/commands/nearpc.py  |  6 ++++++
 pwndbg/disasm/color.py     |  1 -
 4 files changed, 43 insertions(+), 14 deletions(-)

diff --git a/pwndbg/color.py b/pwndbg/color.py
index 391b0e188..df1af0d1c 100644
--- a/pwndbg/color.py
+++ b/pwndbg/color.py
@@ -2,8 +2,8 @@ from __future__ import print_function
 from __future__ import unicode_literals
 
 import functools
-
 import six
+import re
 
 import gdb
 import pwndbg.config
@@ -19,6 +19,7 @@ YELLOW         = "\x1b[33m"
 BLUE           = "\x1b[34m"
 PURPLE         = "\x1b[35m"
 CYAN           = "\x1b[36m"
+WHITE          = "\x1b[37m"
 GREY = GRAY    = "\x1b[90m"
 BOLD           = "\x1b[1m"
 UNDERLINE      = "\x1b[4m"
@@ -29,17 +30,21 @@ pwndbg.config.Parameter('color-code', 'red', 'color for executable memory')
 pwndbg.config.Parameter('color-data', 'purple', 'color for all other writable memory')
 pwndbg.config.Parameter('color-rodata', 'normal', 'color for all read only memory')
 pwndbg.config.Parameter('color-rwx', 'underline', 'color added to all RWX memory')
-
-def normal(x): return NORMAL + str(x)
-def bold(x): return BOLD + str(x) + NORMAL
-def red(x): return RED + str(x) + NORMAL
-def blue(x): return BLUE + str(x) + NORMAL
-def gray(x): return GRAY + str(x) + NORMAL
-def green(x): return GREEN + str(x) + NORMAL
-def cyan(x): return CYAN + str(x) + NORMAL
-def yellow(x): return YELLOW + str(x) + NORMAL
-def purple(x): return PURPLE + str(x) + NORMAL
-def underline(x): return UNDERLINE + str(x) + NORMAL
+pwndbg.config.Parameter('color-highlight', 'green,bold', 'color added to highlights like source/pc')
+
+def normal(x): return colorize(x, NORMAL)
+def black(x): return colorize(x, BLACK)
+def red(x): return colorize(x, RED)
+def green(x): return colorize(x, GREEN)
+def yellow(x): return colorize(x, YELLOW)
+def blue(x): return colorize(x, BLUE)
+def purple(x): return colorize(x, PURPLE)
+def cyan(x): return colorize(x, CYAN)
+def white(x): return colorize(x, WHITE)
+def gray(x): return colorize(x, GRAY)
+def bold(x): return colorize(x, BOLD)
+def underline(x): return colorize(x, UNDERLINE)
+def colorize(x, color): return color + terminateWith(str(x), color) + NORMAL
 
 @pwndbg.memoize.reset_on_stop
 def generateColorFunctionInner(old, new):
@@ -71,6 +76,9 @@ def rodata(x):
 def rwx(x):
     return generateColorFunction(pwndbg.config.color_rwx)(x)
 
+def highlight(x):
+    return generateColorFunction(pwndbg.config.color_highlight)(x)
+
 def get(address, text = None):
     """
     Returns a colorized string representing the provided address.
@@ -111,3 +119,9 @@ def legend():
         rwx('RWX'),
         rodata('RODATA')
     ))
+
+def strip(x):
+    return re.sub('\x1b\\[\d+m', '', x)
+
+def terminateWith(x, color):
+    return re.sub('\x1b\\[0m', NORMAL + color, x)
diff --git a/pwndbg/commands/context.py b/pwndbg/commands/context.py
index ab5395921..11d20cd65 100644
--- a/pwndbg/commands/context.py
+++ b/pwndbg/commands/context.py
@@ -141,6 +141,8 @@ def context_code():
 
     return banner + result
 
+pwndbg.config.Parameter('highlight-source', True, 'whether to highlight the closest source line')
+
 def context_source():
     try:
         symtab = gdb.selected_frame().find_sal().symtab
@@ -163,8 +165,16 @@ def context_source():
         if not source or closest_line <= 1:
             return []
 
+        # highlight the current code line
+        source_lines = source.splitlines()
+        if pwndbg.config.highlight_source:
+            for i in range(len(source_lines)):
+                if source_lines[i].startswith('%s\t' % closest_line):
+                    source_lines[i] = pwndbg.color.highlight(source_lines[i])
+                    break
+
         banner = [pwndbg.color.blue(pwndbg.ui.banner("code"))]
-        banner.extend(source.splitlines())
+        banner.extend(source_lines)
         return banner
     except:
         pass
diff --git a/pwndbg/commands/nearpc.py b/pwndbg/commands/nearpc.py
index 05e993cb5..60743e4fb 100644
--- a/pwndbg/commands/nearpc.py
+++ b/pwndbg/commands/nearpc.py
@@ -21,6 +21,8 @@ import pwndbg.ui
 import pwndbg.vmmap
 
 
+pwndbg.config.Parameter('highlight-pc', True, 'whether to highlight the current instruction')
+
 @pwndbg.commands.ParsedCommand
 @pwndbg.commands.OnlyWhenRunning
 def nearpc(pc=None, lines=None, to_string=False, emulate=False):
@@ -99,6 +101,10 @@ def nearpc(pc=None, lines=None, to_string=False, emulate=False):
 
         line   = ' '.join((prefix, "%#x" % i.address, s or '', asm))
 
+        # Highlight the current line if the config is enabled
+        if pwndbg.config.highlight_pc and i.address == pc:
+            line = pwndbg.color.highlight(line)
+
         # If there was a branch before this instruction which was not
         # contiguous, put in some ellipses.
         if prev and prev.address + prev.size != i.address:
diff --git a/pwndbg/disasm/color.py b/pwndbg/disasm/color.py
index f669fb16f..ab38e78c2 100644
--- a/pwndbg/disasm/color.py
+++ b/pwndbg/disasm/color.py
@@ -51,7 +51,6 @@ def instruction(ins):
             # XXX: not sure when this ever happens
             asm += '<-- file a pwndbg bug for this'
         else:
-            colored_addr = pwndbg.color.get(ins.symbol_addr)
             asm = asm.replace(hex(ins.symbol_addr), ins.symbol)
             asm = '%-36s <%s>' % (asm, pwndbg.color.get(ins.symbol_addr))