From 1902b928787678af499ff9dbd5e3b00c64c287ac Mon Sep 17 00:00:00 2001 From: David Rude Date: Thu, 21 Jul 2011 06:14:25 +0000 Subject: [PATCH] Merge custom prompt code git-svn-id: file:///home/svn/framework3/trunk@13272 4d416f70-5f16-0410-b530-b9f4589650da --- lib/msf/ui/console/driver.rb | 22 ++++++------ lib/rex/ui/text/dispatcher_shell.rb | 6 ++-- lib/rex/ui/text/input/readline.rb | 56 +++++++++++++++++++++++++++-- lib/rex/ui/text/shell.rb | 17 ++++++--- 4 files changed, 81 insertions(+), 20 deletions(-) diff --git a/lib/msf/ui/console/driver.rb b/lib/msf/ui/console/driver.rb index f0b945c662..cd8fe28757 100644 --- a/lib/msf/ui/console/driver.rb +++ b/lib/msf/ui/console/driver.rb @@ -72,8 +72,16 @@ class Driver < Msf::Ui::Driver histfile = opts['HistFile'] || Msf::Config.history_file + # Initialize attributes + self.framework = opts['Framework'] || Msf::Simple::Framework.create(opts) + + if self.framework.datastore['Prompt'] + prompt = self.framework.datastore['Prompt'] + prompt_char = self.framework.datastore['Prompt_Char'] || DefaultPromptChar + end + # Call the parent - super(prompt, prompt_char, histfile) + super(prompt, prompt_char, histfile, framework) # Temporarily disable output self.disable_output = true @@ -81,9 +89,6 @@ class Driver < Msf::Ui::Driver # Load pre-configuration load_preconfig - # Initialize attributes - self.framework = opts['Framework'] || Msf::Simple::Framework.create(opts) - # Initialize the user interface to use a different input and output # handle if one is supplied input = opts['LocalInput'] @@ -99,12 +104,6 @@ class Driver < Msf::Ui::Driver output = Rex::Ui::Text::Output::Stdio.new end - # Verify console compatibility on Windows (require Console2) - if Rex::Compat.is_windows and not Rex::Compat.win32_console2_verify and input.kind_of?(Rex::Ui::Text::Input::Stdio) - $stdout.puts "Error: The Metasploit Framework is not compatible with this console" - exit(1) - end - init_ui(input, output) init_tab_complete @@ -390,6 +389,9 @@ 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['Prompt_Char'] || DefaultPromptChar, true) + when "prompt_char" end end diff --git a/lib/rex/ui/text/dispatcher_shell.rb b/lib/rex/ui/text/dispatcher_shell.rb index 17614fdaad..c444426cd4 100644 --- a/lib/rex/ui/text/dispatcher_shell.rb +++ b/lib/rex/ui/text/dispatcher_shell.rb @@ -78,8 +78,8 @@ module DispatcherShell # # Wraps shell.update_prompt # - def update_prompt(prompt=nil) - shell.update_prompt(prompt) + def update_prompt(prompt=nil, prompt_char = nil, mode = false) + shell.update_prompt(prompt, prompt_char, mode) end # @@ -204,7 +204,7 @@ module DispatcherShell # # Initialize the dispatcher shell. # - def initialize(prompt, prompt_char = '>', histfile = nil) + def initialize(prompt, prompt_char = '>', histfile = nil, framework = nil) super # Initialze the dispatcher array diff --git a/lib/rex/ui/text/input/readline.rb b/lib/rex/ui/text/input/readline.rb index d576aafd90..4a96db6cfd 100644 --- a/lib/rex/ui/text/input/readline.rb +++ b/lib/rex/ui/text/input/readline.rb @@ -18,8 +18,6 @@ begin # Initializes the readline-aware Input instance for text. # def initialize(tab_complete_proc = nil) - - if(not Object.const_defined?('Readline')) begin require 'readline' @@ -79,14 +77,66 @@ begin # 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(framework = nil) line = nil orig = Thread.current.priority begin Thread.current.priority = -20 + output.prompting + if framework + if prompt.include?("%T") + t = Time.now + if framework.datastore['Time_Fmt'] + t = t.strftime(framework.datastore['Time_Fmt']) + end + prompt.gsub!(/%T/, t.to_s) + end + + if 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 + + prompt.gsub!(/%H/, hostname.chomp) + end + + if 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 + + prompt.gsub!(/%U/, user.chomp) + end + + prompt.gsub!(/%S/, framework.sessions.count.to_s) + prompt.gsub!(/%J/, framework.jobs.count.to_s) + prompt.gsub!(/%L/, Rex::Socket.source_address("50.50.50.50")) + prompt.gsub!(/%D/, ::Dir.getwd) + end + line = ::Readline.readline(prompt, true) ::Readline::HISTORY.pop if (line and line.empty?) ensure diff --git a/lib/rex/ui/text/shell.rb b/lib/rex/ui/text/shell.rb index 9297b31a99..cd455b2a8c 100644 --- a/lib/rex/ui/text/shell.rb +++ b/lib/rex/ui/text/shell.rb @@ -36,7 +36,7 @@ module Shell # # Initializes a shell that has a prompt and can be interacted with. # - def initialize(prompt, prompt_char = '>', histfile = nil) + def initialize(prompt, prompt_char = '>', histfile = nil, framework = nil) # Set the stop flag to false self.stop_flag = false self.disable_output = false @@ -48,6 +48,8 @@ module Shell self.histfile = histfile self.hist_last_saved = 0 + + self.framework = framework end def init_tab_complete @@ -124,7 +126,7 @@ module Shell break if (self.stop_flag or self.stop_count > 1) init_tab_complete - line = input.pgets + line = input.pgets(self.framework) log_output(input.prompt) # If a block was passed in, pass the line to it. If it returns true, @@ -176,13 +178,20 @@ module Shell # # Change the input prompt. # - def update_prompt(prompt = nil, new_prompt_char = nil) + # 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) if (self.input) - if (prompt) + 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 # Save the prompt before any substitutions self.prompt = new_prompt