Fix source code display (#638)

* Fix source code display

Fixes and simplifies source code display in `context code`. The
reasoning below.

---

It turns out that we determine the source code line for known code in a
tricky and buggy way. It is wrong when we are deep inside many calls and
we go to upper frame via e.g. `up` and then display the `context code`.

And this also occurs after applying the PR #637 before/after this fix.

We also do `except` all exceptions happening in the source code
retrieval which makes it hard to spot potential bugs in this code path.

This commit removes the `except` path and checks for particular edge
cases where we wouldn't have sources available.

---

Also note the `FileNotFoundError` path - in this case the debug symbols
have a source file path there, but it hasn't been found by us. It might
be because the program (or anything else) removed it OR maybe we debug
remotely etc (not sure if this is the case).

This being said, we could potentially inform the users about that BUT...
GDB already does that by itself - maybe in a bit vague way, as it does
it before us displaying the context e.g.:

```
pwndbg> r
Starting program: /home/dc/pwndbg/a.out

Program received signal SIGSEGV, Segmentation fault.
0x000055555555460a in foo () at ./a.c:2
2	./a.c: No such file or directory.
LEGEND: STACK | HEAP | CODE | DATA | RWX | RODATA
// (the rest of the context here)
```

But I am fine with that for now - and I am not sure if we could even
suppress this warning (oh, maybe this string comes to us into prompt
hook/event? I don't know)

* Fix wrong code line if source startswith newlines

For some reason the `pygments.hightlight` remove newlines at start.
We need to preserve them in order to keep proper code lines.

* Use lexer's stripnl=False

pwndbg/color/syntax_highlight.py

@@ -4,6 +4,7 @@ from __future__ import division
 from __future__ import print_function
 from __future__ import unicode_literals
 
+import itertools
 import os.path
 import re
 
@@ -64,13 +65,14 @@ def syntax_highlight(code, filename='.asm'):
 
     if not lexer:
         try:
-            lexer = pygments.lexers.guess_lexer_for_filename(filename, code)
+            lexer = pygments.lexers.guess_lexer_for_filename(filename, code, stripnl=False)
         except pygments.util.ClassNotFound:
             # no lexer for this file or invalid style
             pass
 
     if lexer:
         lexer_cache[filename] = lexer
+
         code = pygments.highlight(code, lexer, formatter).rstrip()
 
     return code

pwndbg/commands/context.py

@@ -232,75 +232,75 @@ def get_highlight_source(filename):
     source_lines = tuple(line.rstrip() for line in source_lines)
     return source_lines
 
+def get_filename_and_formatted_source():
+    """
+    Returns formatted, lines limited and highlighted source as list
+    or if it isn't there - an empty list
+    """
+    sal = gdb.selected_frame().find_sal()  # gdb.Symtab_and_line
+
+    # Check if source code is available
+    if sal.symtab is None:
+        return '', []
+
+    # Get the full source code
+    closest_line = sal.line
+    filename = sal.symtab.fullname()
 
-def context_code():
     try:
-        # Compute the closest pc and line number
-        symtab = gdb.selected_frame().find_sal().symtab
-        linetable = symtab.linetable()
-
-        closest_pc = -1
-        closest_line = -1
-        for line in linetable:
-            real_address = ctypes.c_uint64(line.pc).value
-            # print("line is %d, address is %s" % (line.line, hex(real_address)))
-            if closest_pc < real_address <= pwndbg.regs.pc:
-                closest_line = line.line
-                closest_pc   = real_address
-
-        if closest_line < 0:
-            return []
-
-        # Get the full source code
-        filename = symtab.fullname()
         source = get_highlight_source(filename)
+    except FileNotFoundError:
+        return '', []
 
-        # If it starts on line 1, it's not really using the
-        # correct source code.
-        if not source or closest_line <= 1:
-            return []
-
-        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
-        prefix_sign = pwndbg.config.code_prefix
-        prefix_width = len(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)"), 'In file: %s' % filename]
-        banner.extend(formatted_source)
-        return banner
-    except:
-        pass
+    if not source:
+        return '', []
+
+    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
+    prefix_sign = pwndbg.config.code_prefix
+    prefix_width = len(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)
+
+    return filename, formatted_source
+
+
+def context_code():
+    filename, formatted_source = get_filename_and_formatted_source()
+
+    # Try getting source from files
+    if formatted_source:
+        return [pwndbg.ui.banner("Source (code)"), 'In file: %s' % filename] + formatted_source
+
+    # Try getting source from IDA Pro Hex-Rays Decompiler
     if not pwndbg.ida.available():
         return []
 
-    # May be None when decompilation failed or user loaded wrong binary in IDA
     n = int(int(int(source_code_lines) / 2)) # int twice to make it a real int instead of inthook
+    # May be None when decompilation failed or user loaded wrong binary in IDA
     code = pwndbg.ida.decompile_context(pwndbg.regs.pc, n)
     
     if code: