Fix IDA Pro decompiled code not being displayed (#328)

* Fix withHexrays decorator not returning wrapper function

* IDA xmlrpc: add cfuncptr_t marshaller & better errors

* IDA xmlrpc server: add shutdown() which can be used for dev

* Small refactor of context.py

* Fix context Hexrays decompiled code display

ida_script.py

@@ -6,6 +6,7 @@ import datetime
 import threading
 import xmlrpclib
 from SimpleXMLRPCServer import SimpleXMLRPCServer
+from xml.sax.saxutils import escape
 
 import idaapi
 import idautils
@@ -25,8 +26,28 @@ if idaapi.IDA_SDK_VERSION >= 700:
 else:
     idc.SaveBase(idc.GetIdbPath() + '.' + dt)
 
-xmlrpclib.Marshaller.dispatch[type(0L)] = lambda _, v, w: w("<value><i8>%d</i8></value>" % v)
-xmlrpclib.Marshaller.dispatch[type(0)] = lambda _, v, w: w("<value><i8>%d</i8></value>" % v)
+
+DEBUG_MARSHALLING = False
+
+def create_marshaller(use_format=None, just_to_str=False):
+    assert use_format or just_to_str, 'Either pass format to use or make it converting the value to str.'
+
+    def wrapper(_marshaller, value, appender):
+        if use_format:
+            marshalled = use_format % value
+        elif just_to_str:
+            marshalled = '<value><string>%s</string></value>' % escape(str(value))
+
+        if DEBUG_MARSHALLING:
+            print("Marshalled: '%s'" % marshalled)
+
+        appender(marshalled)
+
+    return wrapper
+
+xmlrpclib.Marshaller.dispatch[type(0L)] = create_marshaller("<value><i8>%d</i8></value>")
+xmlrpclib.Marshaller.dispatch[type(0)] = create_marshaller("<value><i8>%d</i8></value>")
+xmlrpclib.Marshaller.dispatch[idaapi.cfuncptr_t] = create_marshaller(just_to_str=True)
 
 host = '127.0.0.1'
 port = 8888
@@ -47,23 +68,30 @@ mutex = threading.Condition()
 
 def wrap(f):
     def wrapper(*a, **kw):
-        try:
-            rv = []
-
-            def work():
-                rv.append(f(*a, **kw))
-
-            with mutex:
-                flags = idaapi.MFF_WRITE
-                if f == idc.SetColor:
-                    flags |= idaapi.MFF_NOWAIT
-                    rv.append(None)
-                idaapi.execute_sync(work, flags)
-            return rv[0]
-        except:
-            import traceback
-            traceback.print_exc()
-            raise
+        rv = []
+        error = []
+
+        def work():
+            try:
+                result = f(*a, **kw)
+                rv.append(result)
+            except Exception as e:
+                error.append(e)
+
+        with mutex:
+            flags = idaapi.MFF_WRITE
+            if f == idc.SetColor:
+                flags |= idaapi.MFF_NOWAIT
+                rv.append(None)
+            idaapi.execute_sync(work, flags)
+
+        if error:
+            msg = 'Failed on calling {}.{} with args: {}, kwargs: {}\nException: {}' \
+                .format(f.__module__, f.__name__, a, kw, str(error[0]))
+            print('[!!!] ERROR:', msg)
+            raise error[0]
+
+        return rv[0]
 
     return wrapper
 
@@ -81,8 +109,18 @@ register_module(idaapi)
 server.register_function(lambda a: eval(a, globals(), locals()), 'eval')
 server.register_introspection_functions()
 
-print('Ida Pro xmlrpc hosted on http://%s:%s' % (host, port))
+print('IDA Pro xmlrpc hosted on http://%s:%s' % (host, port))
+print('Call `shutdown()` to shutdown the IDA Pro xmlrpc server.')
 
 thread = threading.Thread(target=server.serve_forever)
 thread.daemon = True
 thread.start()
+
+
+def shutdown():
+    global server
+    global thread
+    server.shutdown()
+    server.server_close()
+    del server
+    del thread

pwndbg/commands/context.py

@@ -71,11 +71,10 @@ def context(*args):
         sys.stdout.write(line + '\n')
     sys.stdout.flush()
 
+
 def context_regs():
-    result = []
-    result.append(pwndbg.ui.banner("registers"))
-    result.extend(get_regs())
-    return result
+    return [pwndbg.ui.banner("registers")] + get_regs()
+
 
 @pwndbg.commands.Command
 @pwndbg.commands.OnlyWhenRunning
@@ -86,6 +85,7 @@ def regs(*regs):
 pwndbg.config.Parameter('show-flags', False, 'whether to show flags registers')
 pwndbg.config.Parameter('show-retaddr-reg', False, 'whether to show return address register')
 
+
 def get_regs(*regs):
     result = []
 
@@ -164,6 +164,7 @@ def context_disasm():
 
 theme.Parameter('highlight-source', True, 'whether to highlight the closest source line')
 
+
 def context_code():
     try:
         symtab = gdb.selected_frame().find_sal().symtab
@@ -172,7 +173,7 @@ def context_code():
         closest_pc = -1
         closest_line = -1
         for line in linetable:
-            if line.pc <= pwndbg.regs.pc and line.pc > closest_pc:
+            if closest_pc < line.pc <= pwndbg.regs.pc:
                 closest_line = line.line
                 closest_pc   = line.pc
 
@@ -203,21 +204,22 @@ def context_code():
     if not pwndbg.ida.available():
         return []
 
-    try:
-        name = pwndbg.ida.GetFunctionName(pwndbg.regs.pc)
-        addr = pwndbg.ida.LocByName(name)
-        lines = pwndbg.ida.decompile(addr)
-        return lines.splitlines()
-    except:
-        pass
+    name = pwndbg.ida.GetFunctionName(pwndbg.regs.pc)
+    addr = pwndbg.ida.LocByName(name)
+    # May be None when decompilation failed or user loaded wrong binary in IDA
+    code = pwndbg.ida.decompile(addr)
+
+    if code:
+        return [pwndbg.ui.banner("Hexrays pseudocode")] + code.splitlines()
+    else:
+        return []
 
-    return []
 
 stack_lines = pwndbg.config.Parameter('context-stack-lines', 8, 'number of lines to print in the stack context')
 
+
 def context_stack():
-    result = []
-    result.append(pwndbg.ui.banner("stack"))
+    result = [pwndbg.ui.banner("stack")]
     telescope = pwndbg.commands.telescope.telescope(pwndbg.regs.sp, to_string=True, count=stack_lines)
     if telescope:
         result.extend(telescope)
@@ -225,6 +227,7 @@ 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):
     result = []
 
@@ -273,6 +276,7 @@ def context_backtrace(frame_count=10, with_banner=True):
         i    += 1
     return result
 
+
 def context_args():
     result = []
 
@@ -292,6 +296,7 @@ def context_args():
 
 last_signal = []
 
+
 def save_signal(signal):
     global last_signal
     last_signal = result = []
@@ -321,6 +326,7 @@ gdb.events.cont.connect(save_signal)
 gdb.events.stop.connect(save_signal)
 gdb.events.exited.connect(save_signal)
 
+
 def context_signal():
     return last_signal
 

pwndbg/ida.py

@@ -74,6 +74,7 @@ class withIDA(object):
             return self.fn(*args, **kwargs)
         return None
 
+
 def withHexrays(f):
     @withIDA
     @functools.wraps(f)
@@ -81,6 +82,9 @@ def withHexrays(f):
         if _ida.init_hexrays_plugin():
             return f(*a, **kw)
 
+    return wrapper
+
+
 def takes_address(function):
     @functools.wraps(function)
     def wrapper(address, *args, **kwargs):
@@ -332,7 +336,13 @@ def has_cached_cfunc(addr):
 @takes_address
 @pwndbg.memoize.reset_on_stop
 def decompile(addr):
-    return _ida.decompile(addr)
+    try:
+        return _ida.decompile(addr)
+    except xmlrpclib.Fault as f:
+        if str(f) == '''<Fault 1: "<class 'idaapi.DecompilationFailure'>:Decompilation failed: ">''':
+            print('Returning an empty string')
+            return None
+        raise
 
 
 @withIDA