Merge custom prompt code

git-svn-id: file:///home/svn/framework3/trunk@13272 4d416f70-5f16-0410-b530-b9f4589650da

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
 

lib/rex/ui/text/dispatcher_shell.rb

@@ -78,8 +78,8 @@ module DispatcherShell
 		#
 		# Wraps shell.update_prompt
 		#
-		def update_prompt(prompt=nil)
+		def update_prompt(prompt=nil, prompt_char = nil, mode = false)
-			shell.update_prompt(prompt)
+			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

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

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,14 +178,21 @@ 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