From 491a3a31627bf17fce10839e0862914aba78c9a0 Mon Sep 17 00:00:00 2001
From: Dylan Davis <dylan.davis@risksense.com>
Date: Wed, 16 Nov 2016 20:43:07 -0700
Subject: [PATCH] Prevent the input prompt from being mangled by asynchronous
 prints.

---
 lib/rex/ui/text/input/readline.rb | 19 +++++++++++++++++--
 lib/rex/ui/text/output.rb         |  8 +++++++-
 lib/rex/ui/text/shell.rb          |  2 ++
 3 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/lib/rex/ui/text/input/readline.rb b/lib/rex/ui/text/input/readline.rb
index c6839ee630..a2aa670b56 100644
--- a/lib/rex/ui/text/input/readline.rb
+++ b/lib/rex/ui/text/input/readline.rb
@@ -40,6 +40,20 @@ begin
       ::Readline.completion_proc = tab_complete_proc || @rl_saved_proc
     end
 
+
+    #
+    # Retrieve the line buffer
+    #
+    def line_buffer
+      if defined? RbReadline
+        RbReadline.rl_line_buffer
+      else
+        ::Readline.line_buffer
+      end
+    end
+
+    attr_accessor :prompt
+
     #
     # Whether or not the input medium supports readline.
     #
@@ -124,12 +138,13 @@ begin
       # to reimplement []`Readline.readline`](https://github.com/luislavena/rb-readline/blob/ce4908dae45dbcae90a6e42e3710b8c3a1f2cd64/lib/readline.rb#L36-L58)
       # for rb-readline to support setting input and output.  Output needs to be set so that colorization works for the
       # prompt on Windows.
+      self.prompt = prompt
       if defined? RbReadline
         RbReadline.rl_instream = fd
         RbReadline.rl_outstream = output
 
         begin
-          line = RbReadline.readline(prompt)
+          line = RbReadline.readline("\001\r\033[K\002" + prompt)
         rescue ::Exception => exception
           RbReadline.rl_cleanup_after_signal()
           RbReadline.rl_deprep_terminal()
@@ -143,7 +158,7 @@ begin
 
         line.try(:dup)
       else
-        ::Readline.readline(prompt, true)
+        ::Readline.readline("\001\r\033[K\002" + prompt, true)
       end
     end
 
diff --git a/lib/rex/ui/text/output.rb b/lib/rex/ui/text/output.rb
index d48eda6d2e..28dc437061 100644
--- a/lib/rex/ui/text/output.rb
+++ b/lib/rex/ui/text/output.rb
@@ -29,6 +29,7 @@ class Output < Rex::Ui::Output
     super
   end
   attr_reader :config
+  attr_accessor :input
 
   def disable_color
     @config[:color] = false
@@ -60,7 +61,12 @@ class Output < Rex::Ui::Output
   end
 
   def print_line(msg = '')
-    print(msg + "\n")
+    print("\r\033[K" + msg + "\n")
+    if input and input.prompt
+      print("\r\033[K")
+      print(input.prompt)
+      print(input.line_buffer)
+    end
   end
 
   def print_warning(msg = '')
diff --git a/lib/rex/ui/text/shell.rb b/lib/rex/ui/text/shell.rb
index 63089bc194..373c6a9eb0 100644
--- a/lib/rex/ui/text/shell.rb
+++ b/lib/rex/ui/text/shell.rb
@@ -184,7 +184,9 @@ module Shell
           self.init_prompt = input.prompt
         end
 
+        output.input = input
         line = input.pgets()
+        output.input = nil
         log_output(input.prompt)
 
         # If a block was passed in, pass the line to it.  If it returns true,