From 5f730277cfa5628746b4144af86c25df2293f2dd Mon Sep 17 00:00:00 2001
From: Luke Imhoff <luke_imhoff@rapid7.com>
Date: Mon, 8 Dec 2014 14:31:00 -0600
Subject: [PATCH] Fix prompt coloring on Windows

MSP-11669

Set output stream for RbReadline (rl_outstream) to the
Rex::Ui::Text::Output::Stdio, which will use translate the ANSI color
escapes to set_color calls in Windows.
---
 lib/rex/ui/text/input/readline.rb | 33 +++++++++++-
 lib/rex/ui/text/output/stdio.rb   | 84 +++++++++++++++++++++++++------
 2 files changed, 102 insertions(+), 15 deletions(-)

diff --git a/lib/rex/ui/text/input/readline.rb b/lib/rex/ui/text/input/readline.rb
index b2b31242f3..c6839ee630 100644
--- a/lib/rex/ui/text/input/readline.rb
+++ b/lib/rex/ui/text/input/readline.rb
@@ -83,7 +83,7 @@ begin
         Thread.current.priority = -20
 
         output.prompting
-        line = ::Readline.readline(prompt, true)
+        line = readline_with_output(prompt, true)
         ::Readline::HISTORY.pop if (line and line.empty?)
       ensure
         Thread.current.priority = orig || 0
@@ -116,6 +116,37 @@ begin
     #
     attr_accessor :output
 
+    private
+
+    def readline_with_output(prompt, add_history=false)
+      # rb-readlines's Readline.readline hardcodes the input and output to $stdin and $stdout, which means setting
+      # `Readline.input` or `Readline.ouput` has no effect when running `Readline.readline` with rb-readline, so need
+      # 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.
+      if defined? RbReadline
+        RbReadline.rl_instream = fd
+        RbReadline.rl_outstream = output
+
+        begin
+          line = RbReadline.readline(prompt)
+        rescue ::Exception => exception
+          RbReadline.rl_cleanup_after_signal()
+          RbReadline.rl_deprep_terminal()
+
+          raise exception
+        end
+
+        if add_history && line
+          RbReadline.add_history(line)
+        end
+
+        line.try(:dup)
+      else
+        ::Readline.readline(prompt, true)
+      end
+    end
+
   end
 rescue LoadError
 end
diff --git a/lib/rex/ui/text/output/stdio.rb b/lib/rex/ui/text/output/stdio.rb
index 28eea33613..b55bd5e685 100644
--- a/lib/rex/ui/text/output/stdio.rb
+++ b/lib/rex/ui/text/output/stdio.rb
@@ -16,6 +16,76 @@ module Text
 #
 ###
 class Output::Stdio < Rex::Ui::Text::Output
+  #
+  # Attributes
+  #
+
+  # @!attribute io
+  #   The raw `IO` backing this Text output.  Defaults to `$stdout`
+  #
+  #   @return [#flush, #puts, #write]
+  attr_writer :io
+
+  #
+  # Constructor
+  #
+
+  # @param options [Hash{Symbol => IO}]
+  # @option options [IO]
+  def initialize(options={})
+    options.assert_valid_keys(:io)
+
+    super()
+
+    self.io = options[:io]
+  end
+
+  #
+  # Methods
+  #
+
+  def flush
+    io.flush
+  end
+
+  # IO to write to.
+  #
+  # @return [IO] Default to `$stdout`
+  def io
+    @io ||= $stdout
+  end
+
+  #
+  # Prints the supplied message to standard output.
+  #
+  def print_raw(msg = '')
+    if (Rex::Compat.is_windows and supports_color?)
+      WindowsConsoleColorSupport.new(io).write(msg)
+    else
+      io.print(msg)
+    end
+
+    io.flush
+
+    msg
+  end
+  alias_method :write, :print_raw
+
+  def puts(*args)
+    args.each do |argument|
+      line = argument.to_s
+      write(line)
+
+      unless line.ends_with? "\n"
+        # yes, this is output, but `IO#puts` uses `rb_default_rs`, which is
+        # [`$/`](https://github.com/ruby/ruby/blob/3af8e150aded9d162bfd41426aaaae0279e5a653/io.c#L12168-L12172),
+        # which is [`$INPUT_RECORD_SEPARATOR`](https://github.com/ruby/ruby/blob/3af8e150aded9d162bfd41426aaaae0279e5a653/lib/English.rb#L83)
+        write($INPUT_RECORD_SEPARATOR)
+      end
+    end
+
+    nil
+  end
 
   def supports_color?
     case config[:color]
@@ -31,20 +101,6 @@ class Output::Stdio < Rex::Ui::Text::Output
       return (term and term.match(/(?:vt10[03]|xterm(?:-color)?|linux|screen|rxvt)/i) != nil)
     end
   end
-
-  #
-  # Prints the supplied message to standard output.
-  #
-  def print_raw(msg = '')
-    if (Rex::Compat.is_windows and supports_color?)
-      WindowsConsoleColorSupport.new($stdout).write(msg)
-    else
-      $stdout.print(msg)
-    end
-    $stdout.flush
-
-    msg
-  end
 end
 
 end