Enhancements to IDA script to make it crash less (never?)

ida_script.py

@@ -1,7 +1,7 @@
 import idaapi
 import idautils
 import idc
-
+import functools
 import datetime
 import threading
 import xmlrpclib
@@ -13,8 +13,7 @@ idc.SaveBase(idc.GetIdbPath() + '.' + datetime.datetime.now().isoformat())
 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)
 
-port = 8888
-
+port       = 8888
 orig_LineA = idc.LineA
 
 def LineA(*a,**kw):
@@ -25,10 +24,25 @@ def LineA(*a,**kw):
 
 idc.LineA = LineA
 
+mutex = threading.Condition()
+
+def wrap(f):
+    def wrapper(*a, **kw):
+        try:
+            rv = []
+            def work(): rv.append(f(*a,**kw))
+            with mutex: idaapi.execute_sync(work, idaapi.MFF_WRITE)
+            return rv[0]
+        except:
+            import traceback
+            traceback.print_exc()
+            raise
+    return wrapper
+
 def register_module(module):
     for name, function in module.__dict__.items():
         if hasattr(function, '__call__'):
-            server.register_function(function, name)
+            server.register_function(wrap(function), name)
 
 server = SimpleXMLRPCServer(('127.0.0.1', port), logRequests=True, allow_none=True)
 register_module(idc)

pwndbg/ida.py

@@ -5,8 +5,10 @@ Talks to an XMLRPC server running inside of an active IDA Pro instance,
 in order to query it about the database.  Allows symbol resolution and
 interactive debugging.
 """
+import errno
 import functools
 import socket
+import traceback
 from contextlib import closing
 
 import gdb
@@ -35,20 +37,25 @@ _ida = None
 
 xmlrpclib.Marshaller.dispatch[type(0)] = lambda _, v, w: w("<value><i8>%d</i8></value>" % v)
 
-
 def setPort(port):
     global _ida
     _ida = xmlrpclib.ServerProxy('http://localhost:%s' % port)
-    try:     _ida.here()
-    except:  _ida = None
-
-setPort(8888)
+    try:
+        _ida.here()
+    except socket.error as e:
+        if e.errno != errno.ECONNREFUSED:
+            traceback.print_exc()
+        _ida = None
 
 class withIDA(object):
     def __init__(self, fn):
         self.fn = fn
         functools.update_wrapper(self, fn)
     def __call__(self, *args, **kwargs):
+        # import pdb
+        # pdb.set_trace()
+        if _ida is None:
+            setPort(8888)
         if _ida is not None:
             return self.fn(*args, **kwargs)
         return None