Land rapid7#9516, Support Bash-Style Continuation Lines

Author: busterb

lib/rex/ui/text/input/readline.rb

@@ -88,7 +88,7 @@ def gets()
     # down other background threads. This is important when there are many active
     # background jobs, such as when the user is running Karmetasploit
     #
-    def pgets()
+    def pgets
 
       line = nil
       orig = Thread.current.priority

lib/rex/ui/text/shell.rb

@@ -22,7 +22,7 @@ module Shell
   module InputShell
     attr_accessor :prompt, :output
 
-    def pgets()
+    def pgets
 
       output.print(prompt)
       output.flush
@@ -46,6 +46,8 @@ def initialize(prompt, prompt_char = '>', histfile = nil, framework = nil)
 
     # Initialize the prompt
     self.init_prompt = prompt
+    self.cont_prompt = ' > '
+    self.cont_flag = false
     self.prompt_char = prompt_char
 
     self.histfile = histfile
@@ -56,7 +58,8 @@ def initialize(prompt, prompt_char = '>', histfile = nil, framework = nil)
 
   def init_tab_complete
     if (self.input and self.input.supports_readline)
-      self.input = Input::Readline.new(lambda { |str| tab_complete(str) })
+      # Unless cont_flag because there's no tab complete for continuation lines
+      self.input = Input::Readline.new(lambda { |str| tab_complete(str) unless cont_flag })
       if Readline::HISTORY.length == 0 and histfile and File.exist?(histfile)
         File.readlines(histfile).each { |e|
           Readline::HISTORY << e.chomp
@@ -185,17 +188,14 @@ def run(&block)
           self.init_prompt = input.prompt
         end
 
-        output.input = input
-        line = input.pgets()
-        output.input = nil
-        log_output(input.prompt)
+        line = get_input_line
 
         # If a block was passed in, pass the line to it.  If it returns true,
         # break out of the shell loop.
         if (block)
           break if (line == nil or block.call(line))
         elsif(input.eof? or line == nil)
-        # If you have sessions active, this will give you a shot to exit gravefully
+        # If you have sessions active, this will give you a shot to exit gracefully
         # If you really are ambitious, 2 eofs will kick this out
           self.stop_count += 1
           next if(self.stop_count > 1)
@@ -350,6 +350,40 @@ def print(msg='')
 
 protected
 
+  #
+  # Get a single line of input, following continuation directives as necessary.
+  #
+  def get_input_line
+    line = "\\\n"
+    prompt_needs_reset = false
+
+    self.cont_flag = false
+    while line =~ /(^|[^\\])\\\s*$/
+      # Strip \ and all the trailing whitespace
+      line.sub!(/\\\s*/, '')
+
+      if line.length > 0
+        # Using update_prompt will overwrite the primary prompt
+        input.prompt = output.update_prompt(self.cont_prompt)
+        self.cont_flag = true
+        prompt_needs_reset = true
+      end
+
+      output.input = input
+      line << input.pgets
+      output.input = nil
+      log_output(input.prompt)
+    end
+    self.cont_flag = false
+
+    if prompt_needs_reset
+      # The continuation prompt was used so reset the prompt
+      update_prompt
+    end
+
+    line
+  end
+
   #
   # Parse a line into an array of arguments.
   #
@@ -389,12 +423,17 @@ def log_output(buf)
   end
 
   attr_writer   :input, :output # :nodoc:
-  attr_accessor :stop_flag, :init_prompt # :nodoc:
+  attr_accessor :stop_flag, :init_prompt, :cont_prompt # :nodoc:
   attr_accessor :prompt # :nodoc:
   attr_accessor :prompt_char, :tab_complete_proc # :nodoc:
   attr_accessor :histfile # :nodoc:
   attr_accessor :hist_last_saved # the number of history lines when last saved/loaded
   attr_accessor :log_source, :stop_count # :nodoc:
+  attr_reader   :cont_flag # :nodoc:
+
+private
+
+  attr_writer   :cont_flag # :nodoc:
 
 end