Per section context output (#697)

* Configure context output per section

* banner respects width of target output

* Fixed context output help

* ui.banner optionaly force a width

* Allow python functions as context output

* Use is for StdOutput type comparison

Co-Authored-By: Disconnect3d <dominik.b.czarnota@gmail.com>

* Use list-function as initial value of default dict

Co-Authored-By: Disconnect3d <dominik.b.czarnota@gmail.com>

* Append final context linebreak only on stdout

* Documented context splitting feature

* Option to hide context section banners

* Option to set width per context section (currently only banner)

* Splitting screenshot

* Fixed empty lines when not clearing

* Fixed auto banner width (using stdin as before instead of stdout)

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

FEATURES.md

@@ -21,6 +21,54 @@ The output of the context may be redirected to a file (including other tty) by u
 
 ![](caps/context.png)  
 
+The context sections can be distibuted among different tty by using the `contextoutput` command.
+Example: `contextoutput stack /path/to/tty true`
+
+Python can be used to create a tmux layout when starting pwndbg and distributing the context among
+the splits.
+```python
+python
+import atexit
+import os
+from pwndbg.commands.context import contextoutput, output, clear_screen
+bt = os.popen('tmux split-window -P -F "#{pane_id}:#{pane_tty}" -d "cat -"').read().strip().split(":")
+st = os.popen(F'tmux split-window -h -t {bt[0]} -P -F '+'"#{pane_id}:#{pane_tty}" -d "cat -"').read().strip().split(":")
+re = os.popen(F'tmux split-window -h -t {st[0]} -P -F '+'"#{pane_id}:#{pane_tty}" -d "cat -"').read().strip().split(":")
+di = os.popen('tmux split-window -h -P -F "#{pane_id}:#{pane_tty}" -d "cat -"').read().strip().split(":")
+panes = dict(backtrace=bt, stack=st, regs=re, disasm=di)
+for sec, p in panes.items():
+  contextoutput(sec, p[1], True)
+contextoutput("legend", di[1], True)
+atexit.register(lambda: [os.popen(F"tmux kill-pane -t {p[0]}").read() for p in panes.values()])
+end
+```
+If you like it simple, try configuration with [splitmind](https://github.com/jerdna-regeiz/splitmind)
+
+![](caps/context_splitting.png)
+
+Note above example uses splitmind and following configuration:
+
+```python
+python
+import splitmind
+(splitmind.Mind()
+  .tell_splitter(show_titles=True)
+  .tell_splitter(set_title="Main")
+  .right(display="backtrace", size="25%")
+  .above(of="main", display="disasm", size="80%", banner="top")
+  .show("code", on="disasm", banner="none")
+  .right(cmd='tty; tail -f /dev/null', size="65%", clearing=False)
+  .tell_splitter(set_title='Input / Output')
+  .above(display="stack", size="75%")
+  .above(display="legend", size="25")
+  .show("regs", on="legend")
+  .below(of="backtrace", cmd="ipython", size="30%")
+).build(nobanner=True)
+end
+```
+
+
+
 ## Disassembly
 
 Pwndbg uses Capstone Engine to display disassembled instructions, but also leverages its introspection into the instruction to extract memory targets and condition codes.

pwndbg/commands/context.py

@@ -10,6 +10,7 @@ import ast
 import codecs
 import ctypes
 import sys
+from collections import defaultdict
 from io import open
 
 import gdb
@@ -49,6 +50,10 @@ config_context_sections = pwndbg.config.Parameter('context-sections',
                                                   'regs disasm code stack backtrace',
                                                   'which context sections are displayed (controls order)')
 
+# Storing output configuration per section
+outputs = {}
+output_settings = {}
+
 
 @pwndbg.config.Trigger([config_context_sections])
 def validate_context_sections():
@@ -70,17 +75,80 @@ def validate_context_sections():
 
 class StdOutput(object):
     """A context manager wrapper to give stdout"""
-    def __enter__(*args,**kwargs):
+    def __enter__(self):
         return sys.stdout
-    def __exit__(*args, **kwargs):
+    def __exit__(self, *args, **kwargs):
+        pass
+    def __hash__(self):
+        return hash(sys.stdout)
+    def __eq__(self, other):
+        return type(other) is StdOutput
+
+class FileOutput(object):
+    """A context manager wrapper to reopen files on enter"""
+    def __init__(self, *args):
+        self.args = args
+        self.handle = None
+    def __enter__(self):
+        self.handle = open(*self.args)
+        return self.handle
+    def __exit__(self, *args, **kwargs):
+        self.handle.close()
+    def __hash__(self):
+        return hash(self.args)
+    def __eq__(self, other):
+        return self.args == other.args
+
+class CallOutput(object):
+    """A context manager which calls a function on write"""
+    def __init__(self, func):
+        self.func = func
+    def __enter__(self):
+        return self
+    def __exit__(self, *args, **kwargs):
         pass
+    def __hash__(self):
+        return hash(self.func)
+    def __eq__(self, other):
+        return self.func == other.func
+    def write(self, data):
+        self.func(data)
+    def flush(self):
+        try:
+            return self.func.flush()
+        except AttributeError:
+            pass
+    def isatty(self):
+        try:
+            return self.func.isatty()
+        except AttributeError:
+            return False
 
-def output():
+
+def output(section):
     """Creates a context manager corresponding to configured context ouput"""
-    if not config_output or config_output == "stdout":
+    target = outputs.get(section, str(config_output))
+    if not target or target == "stdout":
         return StdOutput()
+    elif callable(target):
+        return CallOutput(target)
     else:
-        return open(str( config_output ), "w")
+        return FileOutput(target, "w")
+
+parser = argparse.ArgumentParser()
+parser.description = "Sets the output of a context section."
+parser.add_argument("section", type=str, help="The section which is to be configured. ('regs', 'disasm', 'code', 'stack', 'backtrace', and/or 'args')")
+parser.add_argument("path", type=str, help="The path to which the output is written")
+parser.add_argument("clearing", type=bool, help="Indicates weather to clear the output")
+parser.add_argument("banner", type=str, default="both", help="Where a banner should be placed: both, top , bottom, none")
+parser.add_argument("width", type=int, default=None, help="Sets a fixed width (used for banner). Set to None for auto")
+@pwndbg.commands.ArgparsedCommand(parser, aliases=['ctx-out'])
+def contextoutput(section, path, clearing, banner="both", width=None):
+    outputs[section] = path
+    output_settings[section] = dict(clearing=clearing,
+                                    width=width,
+                                    banner_top= banner in ["both", "top"],
+                                    banner_bottom= banner in ["both", "bottom"])
 
 # @pwndbg.events.stop
 
@@ -98,33 +166,46 @@ def context(subcontext=None):
     if subcontext is None:
         subcontext = []
     args = subcontext
-    
+
     if len(args) == 0:
         args = config_context_sections.split()
 
-    args = [a[0] for a in args]
-
-    result = [M.legend()] if args else []
+    sections = [("legend", lambda target=None, **kwargs: [M.legend()])] if args else []
+    sections += [(arg, context_sections.get(arg[0], None)) for arg in args]
 
-    for arg in args:
-        func = context_sections.get(arg, None)
+    result = defaultdict(list)
+    result_settings = defaultdict(dict)
+    for section, func in sections:
         if func:
-            result.extend(func())
-    if len(result) > 0:
-        result.append(pwndbg.ui.banner(""))
-    result.extend(context_signal())
-
-    with output() as out:
-        if config_clear_screen:
-            clear_screen(out)
-
-        for line in result:
-            out.write(line + '\n')
-        out.flush()
-
-
-def context_regs():
-    return [pwndbg.ui.banner("registers")] + get_regs()
+            target = output(section)
+            # Last section of an output decides about output settings
+            settings = output_settings.get(section, {})
+            result_settings[target].update(settings)
+            with target as out:
+                result[target].extend(func(target=out,
+                                           width=settings.get("width", None),
+                                           with_banner=settings.get("banner_top", True)))
+
+    for target, res in result.items():
+        settings = result_settings[target] 
+        if len(res) > 0 and settings.get("banner_bottom", True):
+            with target as out:
+                res.append(pwndbg.ui.banner("", target=out, 
+                                            width=settings.get("width", None)))
+
+    for target, lines in result.items():
+        with target as out:
+            if result_settings[target].get("clearing", config_clear_screen) and lines:
+                clear_screen(out)
+            out.write("\n".join(lines))
+            if out is sys.stdout:
+                out.write('\n')
+            out.flush()
+
+
+def context_regs(target=sys.stdout, with_banner=True, width=None):
+    banner = [pwndbg.ui.banner("registers", target=target, width=width)]
+    return banner + get_regs() if with_banner else get_regs()
 
 parser = argparse.ArgumentParser()
 parser.description = '''Print out all registers and enhance the information.'''
@@ -184,8 +265,8 @@ Unicorn emulation of code near the current instruction
 ''')
 code_lines = pwndbg.config.Parameter('context-code-lines', 10, 'number of additional lines to print in the code context')
 
-def context_disasm():
-    banner = [pwndbg.ui.banner("disasm")]
+def context_disasm(target=sys.stdout, with_banner=True, width=None):
+    banner = [pwndbg.ui.banner("disasm", target=target, width=width)]
     emulate = bool(pwndbg.config.emulate)
     result = pwndbg.commands.nearpc.nearpc(to_string=True, emulate=emulate, lines=code_lines // 2)
 
@@ -194,7 +275,7 @@ def context_disasm():
     while len(result) < code_lines + 1:
         result.append('')
 
-    return banner + result
+    return banner + result if with_banner else result
 
 theme.Parameter('highlight-source', True, 'whether to highlight the closest source line')
 source_code_lines = pwndbg.config.Parameter('context-source-code-lines',
@@ -271,12 +352,13 @@ def get_filename_and_formatted_source():
     return filename, formatted_source
 
 
-def context_code():
+def context_code(target=sys.stdout, with_banner=True, width=None):
     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
+        bannerline = [pwndbg.ui.banner("Source (code)", target=target, width=width)] if with_banner else []
+        return bannerline + ['In file: %s' % filename] + formatted_source
 
     # Try getting source from IDA Pro Hex-Rays Decompiler
     if not pwndbg.ida.available():
@@ -285,9 +367,10 @@ def context_code():
     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:
-        return [pwndbg.ui.banner("Hexrays pseudocode")] + code.splitlines()
+        bannerline = [pwndbg.ui.banner("Hexrays pseudocode", target=target, width=width)] if with_banner else []
+        return bannerline + code.splitlines()
     else:
         return []
 
@@ -295,8 +378,8 @@ def context_code():
 stack_lines = pwndbg.config.Parameter('context-stack-lines', 8, 'number of lines to print in the stack context')
 
 
-def context_stack():
-    result = [pwndbg.ui.banner("stack")]
+def context_stack(target=sys.stdout, with_banner=True, width=None):
+    result = [pwndbg.ui.banner("stack", target=target, width=width)] if with_banner else []
     telescope = pwndbg.commands.telescope.telescope(pwndbg.regs.sp, to_string=True, count=stack_lines)
     if telescope:
         result.extend(telescope)
@@ -305,11 +388,11 @@ def context_stack():
 backtrace_frame_label = theme.Parameter('backtrace-frame-label', 'f ', 'frame number label for backtrace')
 
 
-def context_backtrace(frame_count=10, with_banner=True):
+def context_backtrace(frame_count=10, with_banner=True, target=sys.stdout, width=None):
     result = []
 
     if with_banner:
-        result.append(pwndbg.ui.banner("backtrace"))
+        result.append(pwndbg.ui.banner("backtrace", target=target, width=width))
 
     this_frame    = gdb.selected_frame()
     newest_frame  = this_frame
@@ -354,7 +437,7 @@ def context_backtrace(frame_count=10, with_banner=True):
     return result
 
 
-def context_args(with_banner=True):
+def context_args(with_banner=True, target=sys.stdout, width=None):
     args = pwndbg.arguments.format_args(pwndbg.disasm.one())
 
     # early exit to skip section if no arg found
@@ -362,7 +445,7 @@ def context_args(with_banner=True):
         return []
 
     if with_banner:
-        args.insert(0, pwndbg.ui.banner("arguments"))
+        args.insert(0, pwndbg.ui.banner("arguments", target=target, width=width))
 
     return args
 

pwndbg/ui.py

@@ -37,10 +37,10 @@ def check_title_position():
               (title_position, ', '.join(valid_values))))
         title_position.revert_default()
 
-
-def banner(title):
+def banner(title, target=sys.stdin, width=None):
     title = title.upper()
-    _height, width = get_window_size()
+    if width is None: # auto width. In case of stdout, it's better to use stdin (b/c GdbOutputFile)
+        _height, width = get_window_size(target=target if target != sys.stdout else sys.stdin)
     if title:
         title = '%s%s%s' % (config.banner_title_surrounding_left, C.banner_title(title), config.banner_title_surrounding_right)
     if 'left' == title_position:
@@ -56,13 +56,13 @@ def addrsz(address):
     address = int(address) & pwndbg.arch.ptrmask
     return "%{}x".format(2*pwndbg.arch.ptrsize) % address
 
-def get_window_size():
+def get_window_size(target=sys.stdin):
     fallback = (int(os.environ.get('LINES', 20)), int(os.environ.get('COLUMNS', 80)))
-    if not sys.stdin.isatty:
+    if not target.isatty():
         return fallback
     try:
         # get terminal size and force ret buffer len of 4 bytes for safe unpacking by passing equally long arg
-        rows, cols = struct.unpack('hh', fcntl.ioctl(sys.stdin.fileno(), termios.TIOCGWINSZ, '1234'))
+        rows, cols = struct.unpack('hh', fcntl.ioctl(target.fileno(), termios.TIOCGWINSZ, '1234'))
     except:
         rows, cols = fallback
     return rows, cols