pretty output

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

lib/msf/ui/console/driver.rb

@@ -23,7 +23,7 @@ class Driver < Msf::Ui::Driver
 
 	include Msf::Ui::Console::Shell
 
-	def initialize(prompt = "msf")
+	def initialize(prompt = "%umsf", prompt_char = ">%c%b")
 		# Initialize attributes
 		self.framework        = Msf::Framework.new
 		self.dispatcher_stack = []
@@ -33,7 +33,7 @@ class Driver < Msf::Ui::Driver
 		enstack_dispatcher(CommandDispatcher::Core)
 
 		# Initialize the super
-		super(prompt)
+		super
 	end
 
 	# Run a single command line
@@ -42,6 +42,8 @@ class Driver < Msf::Ui::Driver
 		method    = arguments.shift
 		found     = false
 
+		reset_color if (supports_color?)
+
 		if (method)
 			entries = dispatcher_stack.length
 

lib/msf/ui/console/shell.rb

@@ -17,7 +17,7 @@ module Console
 ###
 module Shell
 
-	def initialize(prompt)
+	def initialize(prompt, prompt_char = '>')
 		# Initialize the input and output methods
 		self.input  = StdioInputMethod.new
 		self.output = StdioOutputMethod.new
@@ -32,6 +32,7 @@ module Shell
 
 		# Initialize the prompt
 		self.init_prompt = prompt
+		self.prompt_char = prompt_char
 
 		update_prompt
 
@@ -55,8 +56,57 @@ module Shell
 	end
 
 	# Change the input prompt
-	def update_prompt(prompt = '')
-		self.input.prompt = self.init_prompt + ' ' + prompt + '> '
+	def update_prompt(prompt = '', new_prompt_char = nil)
+		new_prompt = self.init_prompt + ' ' + prompt + prompt_char + ' '
+
+		new_prompt.gsub!(/%u/, colorize('underline'))
+		new_prompt.gsub!(/%b/, colorize('bold'))
+		new_prompt.gsub!(/%c/, colorize('clear'))
+		new_prompt.gsub!(/%red/, colorize('red'))
+		new_prompt.gsub!(/%grn/, colorize('green'))
+		new_prompt.gsub!(/%blu/, colorize('blue'))
+		new_prompt.gsub!(/%yel/, colorize('yellow'))
+		new_prompt.gsub!(/%cya/, colorize('cyan'))
+		new_prompt.gsub!(/%whi/, colorize('white'))
+		new_prompt.gsub!(/%mag/, colorize('magenta'))
+		new_prompt.gsub!(/%blk/, colorize('black'))
+		new_prompt.gsub!(/%dred/, colorize('dark', 'red'))
+		new_prompt.gsub!(/%dgrn/, colorize('dark', 'green'))
+		new_prompt.gsub!(/%dblu/, colorize('dark', 'blue'))
+		new_prompt.gsub!(/%dyel/, colorize('dark', 'yellow'))
+		new_prompt.gsub!(/%dcya/, colorize('dark', 'cyan'))
+		new_prompt.gsub!(/%dwhi/, colorize('dark', 'white'))
+		new_prompt.gsub!(/%dmag/, colorize('dark', 'magenta'))
+
+		self.input.prompt = new_prompt
+		self.prompt_char  = new_prompt_char if (new_prompt_char)
+	end
+
+	#
+	# Color checks
+	#
+	
+	#
+	# Checks to see whether or not colors are supported on this shell
+	# console
+	#
+	def supports_color?
+		return (ENV['TERM'].match(/(?:vt10[03]|xterm(?:-color)?|linux|screen)/i) != nil)
+	end
+
+	#
+	# Resets coloring so that it's back to normal
+	#
+	def reset_color
+		output.print(colorize('clear'))
+	end
+
+	#
+	# Returns colorized text if it's supported, otherwise an empty string
+	#
+	def colorize(*color)
+		#return (supports_color? == false) ? '' : Rex::Ui::Text::Color.ansi(color)
+		return Rex::Ui::Text::Color.ansi(*color)
 	end
 
 	#
@@ -85,11 +135,12 @@ protected
 	def parse_line(line)
 		line.gsub!("(\r|\n)", '')
 
-		args = line.split(' ')
+		args = Rex::Parser::Arguments.from_s(line)
 	end
 
 
 	attr_accessor :input, :output, :stop_flag, :init_prompt
+	attr_accessor :prompt_char
 
 end
 

lib/rex.rb

@@ -25,4 +25,5 @@ require 'rex/socket/comm/local'
 require 'rex/parser/arguments'
 
 # Ui
+require 'rex/ui/text/color'
 require 'rex/ui/text/table'

lib/rex.rb.ts.rb

@@ -23,6 +23,7 @@ require 'rex/socket/comm/local.rb.ut'
 
 require 'rex/parser/arguments.rb.ut'
 
+require 'rex/ui/text/color.rb.ut'
 require 'rex/ui/text/table.rb.ut'
 
 class Rex::TestSuite
@@ -54,6 +55,7 @@ class Rex::TestSuite
 		suite << Rex::Parser::Arguments::UnitTest.suite
 
 		# Ui
+		suite << Rex::Ui::Color::Table::UnitTest.suite
 		suite << Rex::Ui::Text::Table::UnitTest.suite
 
 		return suite;

lib/rex/parser/arguments.rb

@@ -33,14 +33,18 @@ class Arguments
 	# Takes a string and converts it into an array of arguments
 	#
 	def self.from_s(str)
-		sub = str.gsub(/[^\\?]"(.+?)[^\\?]"/) { |s| s.gsub(/\s/, '__SEP__') }.gsub(/"/, '')
+		# I am the king of sucking at regex.
+		sub  = str.gsub(/["](.+?)[^\\]["]/) { |s| s.gsub(/\s/, '__SEP__') }
+		sub.gsub!(/[^\\](["])/) { |s| s[0..0] }
+		sub.gsub!(/(\\")/, '"')
 
-		# Uneven groupings?
-		if (group.length % 2 != 0)
-			raise ArgumentParseError.new, "Uneven groupings", caller
-		end
+		args = []
 
+		sub.split(/\s/).each { |arg|
+			args << arg.gsub(/__SEP__/, ' ')
+		}
 
+		return args
 	end
 
 	#

lib/rex/parser/arguments.rb.ut.rb

@@ -51,4 +51,17 @@ class Rex::Parser::Arguments::UnitTest < Test::Unit::TestCase
 		assert_equal(none, "none")
 	end
 
+	def test_from_s
+		args = Rex::Parser::Arguments.from_s(
+			"this is a test \"of the emergency pimping\" system \\\"buh lee dat\\\" yup")
+
+		assert_equal(args[0], "this")
+		assert_equal(args[3], "test")
+		assert_equal(args[4], "of the emergency pimping")
+		assert_equal(args[5], "system")
+		assert_equal(args[6], "\"buh")
+		assert_equal(args[8], "dat\"")
+		assert_equal(args[9], "yup")
+	end
+
 end

lib/rex/ui/text/color.rb

@@ -0,0 +1,50 @@
+module Rex
+module Ui
+module Text
+
+###
+#
+# Color
+# -----
+#
+# This module provides an interface to getting ANSI color codes.
+# It's taken mostly from perl's Term::ANSIColor by Russ Allbery
+# <rra@stanford.edu> and Zenin <zenin@best.com>.
+#
+###
+module Color
+
+	AnsiAttributes =
+		{
+			'clear'      => 0,
+			'reset'      => 0,
+			'bold'       => 1,
+			'dark'       => 2,
+			'underline'  => 4,
+			'underscore' => 4,
+			'blink'      => 5,
+			'reverse'    => 7,
+			'concealed'  => 8,
+			'black'      => 30,   'on_black'   => 40,
+			'red'        => 31,   'on_red'     => 41,
+			'green'      => 32,   'on_green'   => 42,
+			'yellow'     => 33,   'on_yellow'  => 43,
+			'blue'       => 34,   'on_blue'    => 44,
+			'magenta'    => 35,   'on_magenta' => 45,
+			'cyan'       => 36,   'on_cyan'    => 46,
+			'white'      => 37,   'on_white'   => 47
+		}
+
+	#
+	# Return a string with ANSI codes substituted.  Derived from code
+	# written by The FaerieMUD Consortium.
+	#
+	def self.ansi(*attrs)
+		attr = attrs.collect {|a| AnsiAttributes[a] ? AnsiAttributes[a] : nil}.compact.join(';')
+		attr = "\e[%sm" % attr if (attr.empty? == false)
+		return attr
+	end
+
+end 
+
+end end end

lib/rex/ui/text/color.rb.ut.rb

@@ -0,0 +1,18 @@
+#!/usr/bin/ruby
+
+$:.unshift(File.join(File.dirname(__FILE__), '..', '..', '..'))
+
+require 'test/unit'
+require 'rex/ui/text/color'
+
+class Rex::Ui::Text::Color::UnitTest < Test::Unit::TestCase
+
+	def test_color
+		color  = Rex::Ui::Text::Color.ansi('bold', 'red')
+		color += 'hey sup'
+		color += Rex::Ui::Text::Color.ansi('clear')
+
+		assert_equal("\e[1;31mhey sup\e[0m", color)
+	end
+
+end