From 399eb3137c557c7edd481abdf9700d099aa849b7 Mon Sep 17 00:00:00 2001
From: Benedikt Werner <1benediktwerner@gmail.com>
Date: Tue, 24 Apr 2018 08:39:23 +0200
Subject: [PATCH] Improve behavior without IDA Pro (#442)

* Improve behavior without IDA Pro

* Fix import order

* Improved IDA Pro behaviour more

* Added only_after_first_prompt decorator

* Removed newline after import

* Added documentation

* Improved docstring
---
 DEVELOPING.md        |  3 ++-
 pwndbg/config.py     |  5 +++++
 pwndbg/decorators.py | 27 +++++++++++++++++++++++++++
 pwndbg/ida.py        | 37 ++++++++++++++++++++++++++++++-------
 pwndbg/prompt.py     |  3 +++
 5 files changed, 67 insertions(+), 8 deletions(-)
 create mode 100644 pwndbg/decorators.py

diff --git a/DEVELOPING.md b/DEVELOPING.md
index 3713c7bdf..a7923b703 100644
--- a/DEVELOPING.md
+++ b/DEVELOPING.md
@@ -18,6 +18,8 @@ Feel free to update the list below!
 
 * We have a caching mechanism (["memoization"](https://en.wikipedia.org/wiki/Memoization)) which we use through Python's decorators - those are defined in `pwndbg/memoize.py` - just check its usages
 
+* To block a function before the first prompt was displayed use the `pwndbg.decorators.only_after_first_prompt` decorator.
+
 * Memory accesses should be done through `pwndbg/memory.py` functions
 
 * Process properties can be retrieved thx to `pwndbg/proc.py` - e.g. using `pwndbg.proc.pid` will give us current process pid
@@ -35,4 +37,3 @@ Feel free to update the list below!
 Our tests are written using [pytest](https://docs.pytest.org/en/latest/). It uses some magic so that Python's `assert` can be used for asserting things in tests and it injects dependencies which are called fixtures, into test functions.
 
 The fixtures should be defined in [tests/conftest.py](tests/conftest.py). If you need help with writing tests, feel free to reach out on gitub issues/pr or on our irc channel on freenode.
-
diff --git a/pwndbg/config.py b/pwndbg/config.py
index 0029aac52..85bf88ceb 100644
--- a/pwndbg/config.py
+++ b/pwndbg/config.py
@@ -32,6 +32,8 @@ from functools import total_ordering
 import gdb
 import six
 
+import pwndbg.decorators
+
 TYPES = collections.OrderedDict()
 
 # The value is a plain boolean.
@@ -165,6 +167,9 @@ class Parameter(gdb.Parameter):
         for trigger in triggers[self.name]:
             trigger()
 
+        if not pwndbg.decorators.first_prompt:
+            # Remove the newline that gdb adds automatically
+            return '\b'
         return 'Set %s to %r' % (self.docstring, self.value)
 
     def get_show_string(self, svalue):
diff --git a/pwndbg/decorators.py b/pwndbg/decorators.py
new file mode 100644
index 000000000..ab1643bd8
--- /dev/null
+++ b/pwndbg/decorators.py
@@ -0,0 +1,27 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import functools
+
+first_prompt = False
+
+
+def only_after_first_prompt(value_before=None):
+    """
+    Decorator to prevent a function from running before the first prompt was displayed.
+    The 'value_before' parameter can be used to specify the value that is
+    returned if the function is called before the first prompt was displayed.
+    """
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            if first_prompt:
+                return func(*args, **kwargs)
+            else:
+                return value_before
+        return wrapper
+    return decorator
diff --git a/pwndbg/ida.py b/pwndbg/ida.py
index 90abf8a17..f2f3152d6 100644
--- a/pwndbg/ida.py
+++ b/pwndbg/ida.py
@@ -14,6 +14,7 @@ import errno
 import functools
 import socket
 import sys
+import time
 import traceback
 
 import gdb
@@ -21,6 +22,7 @@ import gdb
 import pwndbg.arch
 import pwndbg.compat
 import pwndbg.config
+import pwndbg.decorators
 import pwndbg.elf
 import pwndbg.events
 import pwndbg.memoize
@@ -37,6 +39,7 @@ except:
 ida_rpc_host = pwndbg.config.Parameter('ida-rpc-host', '127.0.0.1', 'ida xmlrpc server address')
 ida_rpc_port = pwndbg.config.Parameter('ida-rpc-port', 8888, 'ida xmlrpc server port')
 ida_enabled = pwndbg.config.Parameter('ida-enabled', True, 'whether to enable ida integration')
+ida_timeout = pwndbg.config.Parameter('ida-timeout', 2, 'time to wait for ida xmlrpc in seconds')
 
 xmlrpclib.Marshaller.dispatch[int] = lambda _, v, w: w("<value><i8>%d</i8></value>" % v)
 
@@ -50,14 +53,26 @@ _ida = None
 # to avoid printing the same exception multiple times, we store the last exception here
 _ida_last_exception = None
 
+# to avoid checking the connection multiple times with no delay, we store the last time we checked it
+_ida_last_connection_check = 0
 
-@pwndbg.config.Trigger([ida_rpc_host, ida_rpc_port])
+
+@pwndbg.decorators.only_after_first_prompt()
+@pwndbg.config.Trigger([ida_rpc_host, ida_rpc_port, ida_timeout])
 def init_ida_rpc_client():
-    global _ida, _ida_last_exception
+    global _ida, _ida_last_exception, _ida_last_connection_check
+
+    if not ida_enabled:
+        return
+
+    now = time.time()
+    if _ida is None and (now - _ida_last_connection_check) < int(ida_timeout) + 5:
+        return
+
     addr = 'http://{host}:{port}'.format(host=ida_rpc_host, port=ida_rpc_port)
 
     _ida = xmlrpclib.ServerProxy(addr)
-    socket.setdefaulttimeout(3)
+    socket.setdefaulttimeout(int(ida_timeout))
 
     exception = None # (type, value, traceback)
     try:
@@ -76,10 +91,20 @@ def init_ida_rpc_client():
 
     if exception:
         if not isinstance(_ida_last_exception, exception[0]) or _ida_last_exception.args != exception[1].args:
-            print(message.error("[!] Ida Pro xmlrpc error"))
-            traceback.print_exception(*exception)
+            if hasattr(pwndbg.config, "exception_verbose") and pwndbg.config.exception_verbose:
+                print(message.error("[!] Ida Pro xmlrpc error"))
+                traceback.print_exception(*exception)
+            else:
+                exc_type, exc_value, _ = exception
+                print(message.error('Failed to connect to IDA Pro ({}: {})'.format(exc_type.__qualname__, exc_value)))
+                if exc_type is socket.timeout:
+                    print(message.notice('To increase the time to wait for IDA Pro use `') + message.hint('set ida-timeout <new-timeout-in-seconds>') + message.notice('`'))
+                else:
+                    print(message.notice('For more info invoke `') + message.hint('set exception-verbose on') + message.notice('`'))
+                print(message.notice('To disable IDA Pro integration invoke `') + message.hint('set ida-enabled off') + message.notice('`'))
 
     _ida_last_exception = exception and exception[1]
+    _ida_last_connection_check = now
 
 
 class withIDA(object):
@@ -88,8 +113,6 @@ class withIDA(object):
         functools.update_wrapper(self, fn)
 
     def __call__(self, *args, **kwargs):
-        if not ida_enabled:
-            return None
         if _ida is None:
             init_ida_rpc_client()
         if _ida is not None:
diff --git a/pwndbg/prompt.py b/pwndbg/prompt.py
index 57a528529..946b68126 100644
--- a/pwndbg/prompt.py
+++ b/pwndbg/prompt.py
@@ -7,6 +7,7 @@ from __future__ import unicode_literals
 
 import gdb
 
+import pwndbg.decorators
 import pwndbg.events
 import pwndbg.gdbutils
 import pwndbg.memoize
@@ -28,6 +29,8 @@ cur = (gdb.selected_inferior(), gdb.selected_thread())
 
 def prompt_hook(*a):
     global cur
+    pwndbg.decorators.first_prompt = True
+
     new = (gdb.selected_inferior(), gdb.selected_thread())
 
     if cur != new: