Land #10477, console prompt fixes and refactor

lib/msf/ui/console/command_dispatcher/core.rb

@@ -160,7 +160,6 @@ class Core
       cmd_color_help
       return
     end
-    driver.update_prompt
   end
 
   #
@@ -1088,19 +1087,10 @@ class Core
       msg = "Spooling to file #{args[0]}..."
     end
 
-    # Restore color and prompt
+    # Restore color
     driver.output.config[:color] = color
-    prompt = framework.datastore['Prompt'] || Msf::Ui::Console::Driver::DefaultPrompt
-    if active_module
-      # intentionally += and not << because we don't want to modify
-      # datastore or the constant DefaultPrompt
-      prompt += " #{active_module.type}(%bld%red#{active_module.promptname}%clr)"
-    end
-    prompt_char = framework.datastore['PromptChar'] || Msf::Ui::Console::Driver::DefaultPromptChar
-    driver.update_prompt("#{prompt} ", prompt_char, true)
 
     print_status(msg)
-    return
   end
 
   def cmd_sessions_help
@@ -2025,31 +2015,19 @@ class Core
 
     rx = Regexp.new(pattern, match_mods[:insensitive])
 
-    # get a ref to the current console driver
+    # redirect output after saving the old one and getting a new output buffer to use for redirect
-    orig_driver = self.driver
+    orig_output = driver.output
-    # redirect output after saving the old ones and getting a new output buffer to use for redirect
-    orig_driver_output = orig_driver.output
-    orig_driver_input = orig_driver.input
 
     # we use a rex buffer but add a write method to the instance, which is
     # required in order to be valid $stdout
     temp_output = Rex::Ui::Text::Output::Buffer.new
     temp_output.extend Rex::Ui::Text::Output::Buffer::Stdout
 
-    orig_driver.init_ui(orig_driver_input,temp_output)
+    driver.init_ui(driver.input, temp_output)
     # run the desired command to be grepped
-    orig_driver.run_single(cmd)
+    driver.run_single(cmd)
     # restore original output
-    orig_driver.init_ui(orig_driver_input,orig_driver_output)
+    driver.init_ui(driver.input, orig_output)
-    # restore the prompt so we don't get "msf >  >".
-    prompt = framework.datastore['Prompt'] || Msf::Ui::Console::Driver::DefaultPrompt
-    prompt_char = framework.datastore['PromptChar'] || Msf::Ui::Console::Driver::DefaultPromptChar
-    mod = active_module
-    if mod # if there is an active module, give them the fanciness they have come to expect
-      driver.update_prompt("#{prompt} #{mod.type}(%bld%red#{mod.promptname}%clr) ", prompt_char, true)
-    else
-      driver.update_prompt("#{prompt} ", prompt_char, true)
-    end
 
     # dump the command's output so we can grep it
     cmd_output = temp_output.dump_buffer

lib/msf/ui/console/command_dispatcher/modules.rb

@@ -620,11 +620,6 @@ module Msf
             end
 
             mod.init_ui(driver.input, driver.output)
-
-            # Update the command prompt
-            prompt = framework.datastore['Prompt'] || Msf::Ui::Console::Driver::DefaultPrompt
-            prompt_char = framework.datastore['PromptChar'] || Msf::Ui::Console::Driver::DefaultPromptChar
-            driver.update_prompt("#{prompt} #{mod.type}(%bld%red#{mod.promptname}%clr) ", prompt_char, true)
           end
 
           #
@@ -807,11 +802,6 @@ module Msf
 
               # Destack the current dispatcher
               driver.destack_dispatcher
-
-              # Restore the prompt
-              prompt = framework.datastore['Prompt'] || Msf::Ui::Console::Driver::DefaultPrompt
-              prompt_char = framework.datastore['PromptChar'] || Msf::Ui::Console::Driver::DefaultPromptChar
-              driver.update_prompt("#{prompt} ", prompt_char, true)
             end
           end
 

lib/msf/ui/console/driver.rb

@@ -416,10 +416,6 @@ class Driver < Msf::Ui::Driver
         handle_console_logging(val) if (glob)
       when "loglevel"
         handle_loglevel(val) if (glob)
-      when "prompt"
-        update_prompt(val, framework.datastore['PromptChar'] || DefaultPromptChar, true)
-      when "promptchar"
-        update_prompt(framework.datastore['Prompt'] || DefaultPrompt, val, true)
     end
   end
 
@@ -438,6 +434,21 @@ class Driver < Msf::Ui::Driver
     end
   end
 
+  #
+  # Proxies to shell.rb's update prompt with our own extras
+  #
+  def update_prompt(*args)
+    if args.empty?
+      pchar = framework.datastore['PromptChar'] || DefaultPromptChar
+      p = framework.datastore['Prompt'] || DefaultPrompt
+      p = "#{p} #{active_module.type}(%bld%red#{active_module.promptname}%clr)" if active_module
+      super(p, pchar)
+    else
+      # Don't squash calls from within lib/rex/ui/text/shell.rb
+      super(*args)
+    end
+  end
+
   #
   # The framework instance associated with this driver.
   #

lib/rex/ui/text/dispatcher_shell.rb

@@ -131,8 +131,8 @@ module DispatcherShell
     #
     # Wraps shell.update_prompt
     #
-    def update_prompt(prompt=nil, prompt_char = nil, mode = false)
+    def update_prompt(*args)
-      shell.update_prompt(prompt, prompt_char, mode)
+      shell.update_prompt(*args)
     end
 
     def cmd_help_help

lib/rex/ui/text/shell.rb

@@ -45,7 +45,6 @@ module Shell
     self.stop_count     = 0
 
     # Initialize the prompt
-    self.init_prompt = prompt
     self.cont_prompt = ' > '
     self.cont_flag = false
     self.prompt_char = prompt_char
@@ -67,7 +66,6 @@ module Shell
         self.hist_last_saved = Readline::HISTORY.length
       end
       self.input.output = self.output
-      update_prompt(input.prompt)
     end
   end
 
@@ -88,7 +86,6 @@ module Shell
 
       self.input.output = self.output
     end
-    update_prompt('')
   end
 
   #
@@ -131,62 +128,7 @@ module Shell
         break if self.stop_flag || self.stop_count > 1
 
         init_tab_complete
-
+        update_prompt
-        if framework
-          if input.prompt.include?("%T")
-            t = Time.now
-            # This %T is the strftime shorthand for %H:%M:%S
-            format = framework.datastore['PromptTimeFormat'] || "%T"
-            t = t.strftime(format)
-            # This %T is the marker in the prompt where we need to place the time
-            input.prompt.gsub!(/%T/, t.to_s)
-          end
-
-          if input.prompt.include?("%H")
-            hostname = ENV['HOSTNAME']
-            if hostname.nil?
-              hostname = `hostname`.split('.')[0]
-            end
-
-            # check if hostname is still nil
-            if hostname.nil?
-              hostname = ENV['COMPUTERNAME']
-            end
-
-            if hostname.nil?
-              hostname = 'unknown'
-            end
-
-            input.prompt.gsub!(/%H/, hostname.chomp)
-          end
-
-          if input.prompt.include?("%U")
-            user = ENV['USER']
-            if user.nil?
-              user = `whoami`
-            end
-
-            # check if username is still nil
-            if user.nil?
-              user = ENV['USERNAME']
-            end
-
-            if user.nil?
-              user = 'unknown'
-            end
-
-            input.prompt.gsub!(/%U/, user.chomp)
-          end
-
-          input.prompt.gsub!(/%S/, framework.sessions.length.to_s)
-          input.prompt.gsub!(/%J/, framework.jobs.length.to_s)
-          input.prompt.gsub!(/%L/, Rex::Socket.source_address("50.50.50.50"))
-          input.prompt.gsub!(/%D/, ::Dir.getwd)
-          if framework.db.active
-            input.prompt.gsub!(/%W/, framework.db.workspace.name)
-          end
-          self.init_prompt = input.prompt
-        end
 
         line = get_input_line
 
@@ -242,26 +184,17 @@ module Shell
   #
   # prompt - the actual prompt
   # new_prompt_char the char to append to the prompt
-  # mode - append or not to append - false = append true = make a new prompt
+  def update_prompt(new_prompt = self.prompt, new_prompt_char = self.prompt_char)
-  def update_prompt(prompt = nil, new_prompt_char = nil, mode = false)
     if (self.input)
-      if prompt
+      p = new_prompt + ' ' + new_prompt_char + ' '
-        new_prompt = self.init_prompt + ' ' + prompt + prompt_char + ' '
-      else
-        new_prompt = self.prompt || ''
-      end
-
-      if mode
-        new_prompt = prompt + (new_prompt_char || prompt_char) + ' '
-      end
 
       # Save the prompt before any substitutions
       self.prompt = new_prompt
+      self.prompt_char  = new_prompt_char
 
       # Set the actual prompt to the saved prompt with any substitutions
       # or updates from our output driver, be they color or whatever
-      self.input.prompt = self.output.update_prompt(new_prompt)
+      self.input.prompt = self.output.update_prompt(format_prompt(p))
-      self.prompt_char  = new_prompt_char if (new_prompt_char)
     end
   end
 
@@ -345,6 +278,7 @@ module Shell
   #
   attr_reader   :output
 
+  attr_reader   :prompt, :prompt_char
   attr_accessor :on_command_proc
   attr_accessor :on_print_proc
   attr_accessor :framework
@@ -434,17 +368,55 @@ protected
   #
   def prompt_yesno(query)
     p = "#{query} [y/N]"
-    old_p = [self.prompt.sub(/#{Regexp.escape(self.prompt_char)} $/, ''), self.prompt_char]
+    old_p = [self.prompt, self.prompt_char]
-    update_prompt p, ' ', true
+    update_prompt p, ' '
     /^y/i === get_input_line
   ensure
-    update_prompt *old_p, true
+    update_prompt *old_p
+  end
+
+  #
+  # Handle prompt substitutions
+  #
+  def format_prompt(str)
+    if framework
+      if str.include?("%T")
+        t = Time.now
+        # This %T is the strftime shorthand for %H:%M:%S
+        format = framework.datastore['PromptTimeFormat'] || "%T"
+        t = t.strftime(format)
+        # This %T is the marker in the prompt where we need to place the time
+        str.gsub!(/%T/, t.to_s)
+      end
+
+      if str.include?("%H")
+        hostname = ENV['HOSTNAME'] || `hostname`.split('.')[0] ||
+          ENV['COMPUTERNAME'] || 'unknown'
+
+        str.gsub!(/%H/, hostname.chomp)
+      end
+
+      if str.include?("%U")
+        user = ENV['USER'] || `whoami` || ENV['USERNAME'] || 'unknown'
+        str.gsub!(/%U/, user.chomp)
+      end
+
+      str.gsub!(/%S/, framework.sessions.length.to_s)
+      str.gsub!(/%J/, framework.jobs.length.to_s)
+      str.gsub!(/%L/, Rex::Socket.source_address("50.50.50.50"))
+      str.gsub!(/%D/, ::Dir.getwd)
+      if framework.db.active
+        str.gsub!(/%W/, framework.db.workspace.name)
+      end
+    end
+
+    str
   end
 
   attr_writer   :input, :output # :nodoc:
-  attr_accessor :stop_flag, :init_prompt, :cont_prompt # :nodoc:
+  attr_writer   :prompt, :prompt_char # :nodoc:
-  attr_accessor :prompt # :nodoc:
+  attr_accessor :stop_flag, :cont_prompt # :nodoc:
-  attr_accessor :prompt_char, :tab_complete_proc # :nodoc:
+  attr_accessor :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: