From fcc96422506140521bfcdb46c131a5070855484c Mon Sep 17 00:00:00 2001
From: Arusekk <arek_koz@o2.pl>
Date: Fri, 6 Dec 2019 10:26:12 +0100
Subject: [PATCH] Fix ugly raw_input()-related stuff

Closes #1373 #1379 #1380
---
 docs/source/term.rst    |  6 +++
 pwnlib/term/readline.py | 81 +++++++++++++++++++++++++++--------------
 pwnlib/term/term.py     |  2 +-
 pwnlib/ui.py            |  1 +
 4 files changed, 61 insertions(+), 29 deletions(-)

diff --git a/docs/source/term.rst b/docs/source/term.rst
index bee06f8b1..369dc0d89 100644
--- a/docs/source/term.rst
+++ b/docs/source/term.rst
@@ -3,3 +3,9 @@
 
 .. automodule:: pwnlib.term
    :members:
+
+Term Modules
+-------------------
+
+.. toctree::
+    term/readline
diff --git a/pwnlib/term/readline.py b/pwnlib/term/readline.py
index 4bee8ed7a..a535fe842 100644
--- a/pwnlib/term/readline.py
+++ b/pwnlib/term/readline.py
@@ -3,6 +3,9 @@
 from __future__ import absolute_import
 from __future__ import division
 
+import six
+import sys
+
 from pwnlib.term import keyconsts as kc
 from pwnlib.term import keymap as km
 from pwnlib.term import term
@@ -364,13 +367,18 @@ def go_end(*_):
     '<any>'       : handle_keypress,
     })
 
-def readline(_size = None, prompt = '', float = True, priority = 10):
+def readline(_size=None, prompt='', float=True, priority=10):
     # The argument  _size is unused, but is there for compatibility
     # with the existing readline
 
     global buffer_handle, prompt_handle, suggest_handle, eof, \
         show_suggestions
 
+    # XXX circular imports
+    from pwnlib.term import term_mode
+    if not term_mode:
+        print(prompt, end='', flush=True)
+        return sys.stdin.readline().rstrip('\n')
     show_suggestions = False
     eof = False
     if prompt:
@@ -413,7 +421,46 @@ def readline(_size = None, prompt = '', float = True, priority = 10):
         if shutdown_hook:
             shutdown_hook()
 
+def raw_input(prompt='', float=True):
+    r"""raw_input(prompt='', float=True)
+
+    Replacement for the built-in ``raw_input`` using ``pwnlib`` readline
+    implementation.
+
+    Arguments:
+        prompt(str): The prompt to show to the user.
+        float(bool): If set to `True`, prompt and input will float to the
+                     bottom of the screen when `term.term_mode` is enabled.
+    """
+    return readline(None, prompt, float)
+
+def eval_input(prompt='', float=True):
+    """eval_input(prompt='', float=True)
+
+    Replacement for the built-in python 2 - style ``input`` using
+    ``pwnlib`` readline implementation, and `pwnlib.util.safeeval.expr`
+    instead of ``eval`` (!).
+
+    Arguments:
+        prompt(str): The prompt to show to the user.
+        float(bool): If set to ``True``, prompt and input will float to the
+                     bottom of the screen when `term.term_mode` is enabled.
+
+    Example:
+
+        >>> try:
+        ...     saved_stdin = sys.stdin
+        ...     sys.stdin = io.StringIO("{'a':20}")
+        ...     eval_input("Favorite object? ")['a']
+        ... finally:
+        ...     sys.stdin = saved_stdin
+        Favorite object? 20
+    """
+    from pwnlib.util import safeeval
+    return safeeval.const(readline(None, prompt, float))
+
 def init():
+    global safeeval
     # defer imports until initialization
     import sys
     from six.moves import builtins
@@ -428,30 +475,8 @@ def __getattr__(self, k):
             return self._fd.__getattribute__(k)
     sys.stdin = Wrapper(sys.stdin)
 
-    def raw_input(prompt = '', float = True):
-        """raw_input(prompt = '', float = True)
-
-        Replacement for the built-in `raw_input` using ``pwnlib``s readline
-        implementation.
-
-        Arguments:
-            prompt(str): The prompt to show to the user.
-            float(bool): If set to `True`, prompt and input will float to the
-                         bottom of the screen when `term.term_mode` is enabled.
-        """
-        return readline(None, prompt, float)
-    builtins.raw_input = raw_input
-
-    def input(prompt = '', float = True):
-        """input(prompt = '', float = True)
-
-        Replacement for the built-in `input` using ``pwnlib``s readline
-        implementation, and `pwnlib.util.safeeval.expr` instead of `eval` (!).
-
-        Arguments:
-            prompt(str): The prompt to show to the user.
-            float(bool): If set to `True`, prompt and input will float to the
-                         bottom of the screen when `term.term_mode` is enabled.
-        """
-        return safeeval.const(readline(None, prompt, float))
-    builtins.input = input
+    if six.PY2:
+        builtins.raw_input = raw_input
+        builtins.input = eval_input
+    else:
+        builtins.input = raw_input
diff --git a/pwnlib/term/term.py b/pwnlib/term/term.py
index 6e91d6c2c..94cb05198 100644
--- a/pwnlib/term/term.py
+++ b/pwnlib/term/term.py
@@ -487,7 +487,7 @@ def output(s = '', float = False, priority = 10, frozen = False,
         else:
             is_floating = False
             i = len(cells) - 1
-            while cells[i].float and i > 0:
+            while i > 0 and cells[i].float:
                 i -= 1
         # put('xx %d\n' % i)
         cell = Cell()
diff --git a/pwnlib/ui.py b/pwnlib/ui.py
index fad03971a..7df35c7e0 100644
--- a/pwnlib/ui.py
+++ b/pwnlib/ui.py
@@ -7,6 +7,7 @@
 
 from pwnlib import term
 from pwnlib.log import getLogger
+from pwnlib.term.readline import raw_input
 
 log = getLogger(__name__)