Land #10477, console prompt fixes and refactor

4.x
William Vu 2018-09-12 18:00:48 -05:00 committed by Metasploit
parent 90c31b96a7
commit 0f7b8a2453
No known key found for this signature in database
GPG Key ID: CDFB5FA52007B954
5 changed files with 73 additions and 122 deletions

View File

@ -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

View File

@ -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

View File

@ -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.
#

View File

@ -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

View File

@ -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: