Adds $rebase(offset) function (#374)

Adds `$rebase(offset)` gdbfunction that can be used to set up a breakpoint
over an offset from program image base.

Also changed a bit the pwndbg banner displayed at startup.

pwndbg/__init__.py

@@ -62,6 +62,7 @@ import pwndbg.disasm.x86
 import pwndbg.dt
 import pwndbg.elf
 import pwndbg.exception
+import pwndbg.gdbutils.functions
 import pwndbg.heap
 import pwndbg.inthook
 import pwndbg.memory

pwndbg/commands/ida.py

@@ -15,6 +15,7 @@ import pwndbg.commands
 import pwndbg.commands.context
 import pwndbg.ida
 import pwndbg.regs
+from pwndbg.gdbutils.functions import GdbFunction
 
 
 @pwndbg.commands.ParsedCommand
@@ -117,21 +118,14 @@ def save_ida():
 save_ida()
 
 
-class ida(gdb.Function):
-    """Evaluate ida.LocByName() on the supplied value.
-    """
-
-    def __init__(self):
-        super(ida, self).__init__('ida')
-
-    def invoke(self, name):
-        name = name.string()
-        result = pwndbg.ida.LocByName(name)
-
-        if 0xffffe000 <= result <= 0xffffffff or 0xffffffffffffe000 <= result <= 0xffffffffffffffff:
-            raise ValueError("ida.LocByName(%r) == BADADDR" % name)
+@GdbFunction()
+def ida(name):
 
-        return result
+    """Evaluate ida.LocByName() on the supplied value."""
+    name = name.string()
+    result = pwndbg.ida.LocByName(name)
 
+    if 0xffffe000 <= result <= 0xffffffff or 0xffffffffffffe000 <= result <= 0xffffffffffffffff:
+        raise ValueError("ida.LocByName(%r) == BADADDR" % name)
 
-ida()
+    return result

pwndbg/gdbutils/__init__.py

@@ -0,0 +1,10 @@
+# -*- coding: utf-8 -*-
+"""
+Put all new things related to gdb in this module.
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import pwndbg.gdbutils.functions

pwndbg/gdbutils/functions.py

@@ -0,0 +1,54 @@
+# -*- coding: utf-8 -*-
+"""
+Put all functions defined for gdb in here.
+
+This file might be changed into a module in the future.
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import functools
+
+import gdb
+
+import pwndbg.proc
+
+functions = []
+
+
+def GdbFunction(only_when_running=False):
+    return functools.partial(_GdbFunction, only_when_running=only_when_running)
+
+
+class _GdbFunction(gdb.Function):
+    def __init__(self, func, only_when_running):
+        self.name = func.__name__
+        self.func = func
+        self.only_when_running = only_when_running
+
+        functions.append(self)
+
+        super(_GdbFunction, self).__init__(self.name)
+
+        functools.update_wrapper(self, func)
+        self.__doc__ = func.__doc__
+
+    def invoke(self, *args):
+        if self.only_when_running and not pwndbg.proc.alive:
+            # Returning empty string is a workaround that we can't stop e.g. `break *$rebase(offset)`
+            # Thx to that, gdb will print out 'evaluation of this expression requires the target program to be active'
+            return ''
+
+        return self.func(*args)
+
+    def __call__(self, *args):
+        return self.invoke(*args)
+
+
+@GdbFunction(only_when_running=True)
+def rebase(addr):
+    """Return rebased address."""
+    base = pwndbg.elf.exe().address
+    return base + int(addr)

pwndbg/prompt.py

@@ -7,12 +7,21 @@ from __future__ import unicode_literals
 
 import gdb
 
+import pwndbg.color as C
 import pwndbg.events
+import pwndbg.gdbutils
 import pwndbg.memoize
 
-hint_msg = 'Loaded %i commands. Type pwndbg [filter] for a list.' % len(pwndbg.commands.commands)
+funcs_list_str = ', '.join(C.purple('$' + f.name) for f in pwndbg.gdbutils.functions.functions)
+
+hint_lines = (
+    'loaded %i commands. Type %s for a list.' % (len(pwndbg.commands.commands), C.purple('pwndbg [filter]')),
+    'created %s gdb functions (can be used with print/break)' % funcs_list_str
+)
+
+for line in hint_lines:
+    print(C.light_red(pwndbg.color.bold('pwndbg: ') + line))
 
-print(pwndbg.color.red(hint_msg))
 cur = (gdb.selected_inferior(), gdb.selected_thread())
 
 

tests/testLoadsWithoutCrashing.py

@@ -3,9 +3,19 @@
 from __future__ import absolute_import
 from __future__ import unicode_literals
 
+import re
+
 from . import common
 
 
+def escape_ansi(line):
+    # via https://stackoverflow.com/questions/14693701/how-can-i-remove-the-ansi-escape-sequences-from-a-string-in-python
+    ansi_escape = re.compile(r'(\x9B|\x1B\[)[0-?]*[ -/]*[@-~]')
+    return ansi_escape.sub('', line)
+
+
 def test_loads_wivout_crashing_bruv():
-    output = common.run_gdb_with_script()
-    assert 'Type pwndbg [filter] for a list.' in output, output
+    output = escape_ansi(common.run_gdb_with_script())
+
+    assert ('pwndbg: loaded 156 commands. Type pwndbg [filter] for a list.\n'
+            'pwndbg: created $rebase, $ida gdb functions (can be used with print/break)') in output, output