From 0f7b8a24536378fbc59b83817d1166021366b46a Mon Sep 17 00:00:00 2001 From: William Vu Date: Wed, 12 Sep 2018 18:00:48 -0500 Subject: [PATCH] Land #10477, console prompt fixes and refactor --- lib/msf/ui/console/command_dispatcher/core.rb | 34 +---- .../ui/console/command_dispatcher/modules.rb | 10 -- lib/msf/ui/console/driver.rb | 19 ++- lib/rex/ui/text/dispatcher_shell.rb | 4 +- lib/rex/ui/text/shell.rb | 128 +++++++----------- 5 files changed, 73 insertions(+), 122 deletions(-) diff --git a/lib/msf/ui/console/command_dispatcher/core.rb b/lib/msf/ui/console/command_dispatcher/core.rb index 62297d8f42..c43f2d0d1e 100644 --- a/lib/msf/ui/console/command_dispatcher/core.rb +++ b/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 - orig_driver = self.driver - # 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 + # redirect output after saving the old one and getting a new output buffer to use for redirect + orig_output = driver.output # 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) - # 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 + driver.init_ui(driver.input, orig_output) # dump the command's output so we can grep it cmd_output = temp_output.dump_buffer diff --git a/lib/msf/ui/console/command_dispatcher/modules.rb b/lib/msf/ui/console/command_dispatcher/modules.rb index d8bfe870a5..d68d7b86c5 100644 --- a/lib/msf/ui/console/command_dispatcher/modules.rb +++ b/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 diff --git a/lib/msf/ui/console/driver.rb b/lib/msf/ui/console/driver.rb index 569893ae26..c505222a44 100644 --- a/lib/msf/ui/console/driver.rb +++ b/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. # diff --git a/lib/rex/ui/text/dispatcher_shell.rb b/lib/rex/ui/text/dispatcher_shell.rb index 55ae2a61f8..c91d02bfeb 100644 --- a/lib/rex/ui/text/dispatcher_shell.rb +++ b/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) - shell.update_prompt(prompt, prompt_char, mode) + def update_prompt(*args) + shell.update_prompt(*args) end def cmd_help_help diff --git a/lib/rex/ui/text/shell.rb b/lib/rex/ui/text/shell.rb index b7256f6854..14cb4c3336 100644 --- a/lib/rex/ui/text/shell.rb +++ b/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 - - 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 + update_prompt 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(prompt = nil, new_prompt_char = nil, mode = false) + def update_prompt(new_prompt = self.prompt, new_prompt_char = self.prompt_char) if (self.input) - if prompt - 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 + p = new_prompt + ' ' + new_prompt_char + ' ' # 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.prompt_char = new_prompt_char if (new_prompt_char) + self.input.prompt = self.output.update_prompt(format_prompt(p)) 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] - update_prompt p, ' ', true + old_p = [self.prompt, self.prompt_char] + 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_accessor :prompt # :nodoc: - attr_accessor :prompt_char, :tab_complete_proc # :nodoc: + attr_writer :prompt, :prompt_char # :nodoc: + attr_accessor :stop_flag, :cont_prompt # :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: