Implement asm&source syntax highlight (#390)

* Syntax-highlight: Add asm lexer in color/lexer.py

* Syntax-highlight: Add pygments to requirements.txt

* Syntax-highlight: Update lexer for supporting ARM

Support symbol, constant, comments

* Syntax-highlight: Enable asm syntax highlight

* Syntax-highlight: Add source highlight utils in commnads/context.py

* Syntax-highlight: Add disasm highlight utils in color/disasm.py

* Syntax-highlight: Implement Source code highlighting in commands/context.py

* Syntax-highlight: Add syntax_highlight() in color/syntax_highlight.py

* Fix texts

* Add color theme and prefix config for context code

* Add missed utf8 magic comment

* Fix isort

pwndbg/color/context.py

@@ -9,6 +9,7 @@ import pwndbg.color.theme as theme
 import pwndbg.config as config
 from pwndbg.color import generateColorFunction
 
+config_prefix_color             = theme.ColoredParameter('code-prefix-color', 'none', "color for 'context code' command (prefix marker)")
 config_highlight_color          = theme.ColoredParameter('highlight-color', 'green,bold', 'color added to highlights like source/pc')
 config_register_color           = theme.ColoredParameter('context-register-color', 'bold', 'color for registers label')
 config_flag_value_color         = theme.ColoredParameter('context-flag-value-color', 'none', 'color for flags register (register value)')
@@ -21,6 +22,9 @@ config_banner_title             = theme.ColoredParameter('banner-title-color', '
 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')
 
+def prefix(x):
+    return generateColorFunction(config.code_prefix_color)(x)
+
 def highlight(x):
     return generateColorFunction(config.highlight_color)(x)
 

pwndbg/color/disasm.py

@@ -10,6 +10,7 @@ import capstone
 import pwndbg.chain
 import pwndbg.color.context as C
 import pwndbg.color.memory as M
+import pwndbg.color.syntax_highlight as H
 import pwndbg.color.theme as theme
 import pwndbg.config as config
 import pwndbg.disasm.jump
@@ -27,8 +28,15 @@ config_branch = theme.ColoredParameter('disasm-branch-color', 'bold', 'color for
 def branch(x):
     return generateColorFunction(config.disasm_branch_color)(x)
 
+
+def syntax_highlight(ins):
+    return H.syntax_highlight(ins, filename='.asm')
+
+
 def instruction(ins):
     asm = '%-06s %s' % (ins.mnemonic, ins.op_str)
+    if pwndbg.config.syntax_highlight:
+        asm = syntax_highlight(asm)
     is_branch = set(ins.groups) & capstone_branch_groups
 
     # Highlight the current line if enabled

pwndbg/color/lexer.py

@@ -0,0 +1,137 @@
+# -*- coding: utf-8 -*-
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import re
+
+import six
+from pygments.lexer import RegexLexer
+from pygments.lexer import include
+from pygments.token import Comment
+from pygments.token import Name
+from pygments.token import Number
+from pygments.token import Operator
+from pygments.token import Other
+from pygments.token import Punctuation
+from pygments.token import String
+from pygments.token import Text
+
+import pwndbg.compat
+
+__all__ = ['PwntoolsLexer']
+
+class PwntoolsLexer(RegexLexer):
+    """
+    Fork from pwntools
+    https://github.com/Gallopsled/pwntools/blob/7860eecf025135380b137dd9df85dd02a2fd1667/pwnlib/lexer.py
+
+    Edit:
+        * Remove Objdump rules
+        * Merge pygments-arm (https://github.com/heia-fr/pygments-arm)
+    """
+    name = 'PwntoolsLexer'
+    filenames = ['*.s', '*.S', '*.asm']
+
+    #: optional Comment or Whitespace
+    string = r'"(\\"|[^"])*"'
+    char = r'[\w$.@-]'
+    identifier = r'(?:[a-zA-Z$_]' + char + '*|\.' + char + '+|or)'
+    number = r'(?:0[xX][a-zA-Z0-9]+|\d+)'
+    memory = r'(?:[\]\[])'
+
+    eol = r'[\r\n]+'
+
+    tokens = {
+        'root': [
+            include('whitespace'),
+
+            # Label
+            (identifier + ':', Name.Label),
+            (number + ':', Name.Label),
+
+            # AT&T directive
+            (r'\.' + identifier, Name.Attribute, 'directive-args'),
+            (r'lock|rep(n?z)?|data\d+', Name.Attribute),
+
+            # Instructions
+            (identifier, Name.Function, 'instruction-args'),
+
+            (r'[\r\n]+', Text),
+        ],
+        'directive-args': [
+            (identifier, Name.Constant),
+            (string, String),
+            ('@' + identifier, Name.Attribute),
+            (number, Number.Integer),
+
+            (eol, Text, '#pop'),
+            (r'#.*?$', Comment, '#pop'),
+
+            include('punctuation'),
+            include('whitespace')
+        ],
+        'instruction-args': [
+            # Fun things
+            (r'([\]\[]|BYTE|DWORD|PTR|\+|\-|}|{|\^|>>|<<|&)', Text),
+
+            # Address constants
+            (identifier, Name.Constant),
+            ('=' + identifier, Name.Constant), # ARM symbol
+            (number, Number.Integer),
+
+            # Registers
+            ('%' + identifier, Name.Variable),
+            ('$' + identifier, Name.Variable),
+
+            # Numeric constants
+            ('$' + number, Number.Integer),
+            ('#' + number, Number.Integer),
+
+            # ARM predefined constants
+            ('#' + identifier, Name.Constant),
+
+            (r"$'(.|\\')'", String.Char),
+
+            (eol, Text, '#pop'),
+
+            include('punctuation'),
+            include('whitespace')
+        ],
+        'whitespace': [
+            (r'\n', Text),
+            (r'\s+', Text),
+
+            # Block comments
+            # /* */ (AT&T)
+            (r'/\*.*?\*/', Comment),
+
+            # Line comments
+            # //    (AArch64)
+            # #     (AT&T)
+            # ;     (NASM/intel, LLVM)
+            # @     (ARM)
+            (r'(//|[#;@]).*$', Comment.Single)
+        ],
+        'punctuation': [
+            (r'[-*,.():]+', Punctuation)
+        ]
+    }
+
+# Note: convert all unicode() to str() if in Python2.7 since unicode_literals is enabled
+# The pygments<=2.2.0 (latest stable when commit) in Python2.7 use 'str' type in rules matching
+# We must convert all unicode back to str()
+if pwndbg.compat.python2:
+    def _to_str(obj):
+        type_ = type(obj)
+        if type_ in (tuple, list):
+            return type_(map(_to_str, obj))
+        elif type_ is unicode:
+            return str(obj)
+        return obj
+
+    PwntoolsLexer.tokens = {
+        _to_str(k): _to_str(v)
+        for k, v in PwntoolsLexer.tokens.iteritems()
+    }

pwndbg/color/syntax_highlight.py

@@ -0,0 +1,68 @@
+# -*- coding: utf-8 -*-
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import os.path
+import re
+
+import pwndbg.config
+from pwndbg.color import message
+from pwndbg.color import theme
+
+try:
+    import pygments
+    import pygments.lexers
+    import pygments.formatters
+    from pwndbg.color.lexer import PwntoolsLexer
+except ImportError:
+    pygments = None
+
+pwndbg.config.Parameter('syntax-highlight', True, 'Source code / assembly syntax highlight')
+style = theme.Parameter('syntax-highlight-style', 'monokai', 'Source code / assembly syntax highlight stylename of pygments module')
+@pwndbg.config.Trigger([style])
+def check_style():
+    try:
+        formatter = pygments.formatters.Terminal256Formatter(
+            style=str(style)
+        )
+    except pygments.util.ClassNotFound:
+        msg = message.warn("The pygment formatter style '%s' is not found, restore to default" % style)
+        print(msg)
+        style.value = style.default
+
+
+def syntax_highlight(code, filename):
+    # No syntax highlight if pygment is not installed
+    if not pygments:
+        return code
+
+    filename = os.path.basename(filename)
+
+    formatter = pygments.formatters.Terminal256Formatter(
+        style=str(style)
+    )
+
+    lexer = None
+
+    # If source code is asm, use our customized lexer.
+    # Note: We can not register our Lexer to pygments and use their APIs,
+    # since the pygment only search the lexers installed via setuptools.
+    for glob_pat in PwntoolsLexer.filenames:
+        pat = '^' + glob_pat.replace('.', r'\.').replace('*', r'.*') + '$'
+        if re.match(pat, filename):
+            lexer = PwntoolsLexer()
+            break
+
+    if not lexer:
+        try:
+            lexer = pygments.lexers.guess_lexer_for_filename(filename, code)
+        except pygments.util.ClassNotFound:
+            # no lexer for this file or invalid style
+            pass
+
+    if lexer:
+        code = pygments.highlight(code, lexer, formatter).rstrip()
+
+    return code

pwndbg/commands/context.py

@@ -5,6 +5,7 @@ from __future__ import division
 from __future__ import print_function
 from __future__ import unicode_literals
 
+import codecs
 import sys
 
 import gdb
@@ -15,6 +16,7 @@ import pwndbg.color
 import pwndbg.color.backtrace as B
 import pwndbg.color.context as C
 import pwndbg.color.memory as M
+import pwndbg.color.syntax_highlight as H
 import pwndbg.commands
 import pwndbg.commands.nearpc
 import pwndbg.commands.telescope
@@ -164,10 +166,28 @@ def context_disasm():
     return banner + result
 
 theme.Parameter('highlight-source', True, 'whether to highlight the closest source line')
+source_code_lines = pwndbg.config.Parameter('context-source-code-lines',
+                                             10,
+                                             'number of source code lines to print by the context command')
+theme.Parameter('code-prefix', '►', "prefix marker for 'context code' command")
+
+@pwndbg.memoize.reset_on_start
+def get_highlight_source(filename):
+    # Notice that the code is cached
+    with open(filename) as f:
+        source = f.read()
+
+    if pwndbg.config.syntax_highlight:
+        source = H.syntax_highlight(source, filename)
+
+    source_lines = source.splitlines()
+    source_lines = tuple(line.rstrip() for line in source_lines)
+    return source_lines
 
 
 def context_code():
     try:
+        # Compute the closest pc and line number
         symtab = gdb.selected_frame().find_sal().symtab
         linetable = symtab.linetable()
 
@@ -181,23 +201,56 @@ def context_code():
         if closest_line < 0:
             return []
 
-        source = gdb.execute('list %i' % closest_line, from_tty=False, to_string=True)
+        # Get the full source code
+        filename = symtab.fullname()
+        source = get_highlight_source(filename)
 
         # If it starts on line 1, it's not really using the
         # correct source code.
         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] = C.highlight(source_lines[i])
-                    break
-
-        banner = [pwndbg.ui.banner("source")]
-        banner.extend(source_lines)
+        n = int(source_code_lines)
+
+        # Compute the line range
+        start = max(closest_line - 1 - n//2, 0)
+        end = min(closest_line - 1 + n//2 + 1, len(source))
+        num_width = len(str(end))
+
+        # split the code
+        source = source[start:end]
+
+        # Compute the prefix_sign length
+        # TODO: remove this if the config setter can make sure everything is unicode.
+        # This is needed because the config value may be utf8 byte string.
+        # It is better to convert to unicode at setter of config and then we will not need this.
+        prefix_sign = pwndbg.config.code_prefix
+        value = prefix_sign.value
+        if isinstance(value, bytes):
+            value = codecs.decode(value, 'utf-8')
+        prefix_width = len(value)
+
+        # Convert config class to str to make format() work
+        prefix_sign = str(prefix_sign)
+
+        # Format the output
+        formatted_source = []
+        for line_number, code in enumerate(source, start=start + 1):
+            fmt = ' {prefix_sign:{prefix_width}} {line_number:>{num_width}} {code}'
+            if pwndbg.config.highlight_source and line_number == closest_line:
+                fmt = C.highlight(fmt)
+
+            line = fmt.format(
+                prefix_sign=C.prefix(prefix_sign) if line_number == closest_line else '',
+                prefix_width=prefix_width,
+                line_number=line_number,
+                num_width=num_width,
+                code=code
+            )
+            formatted_source.append(line)
+
+        banner = [pwndbg.ui.banner("Source (code)")]
+        banner.extend(formatted_source)
         return banner
     except:
         pass

requirements.txt

@@ -8,4 +8,5 @@ python-ptrace>=0.8
 ROPgadget
 six
 unicorn>=1.0.0
+pygments
 https://github.com/aquynh/capstone/archive/next.zip#subdirectory=bindings/python