variable validation, fixed tab complete

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

lib/msf/base/simple/framework.rb

@@ -27,23 +27,32 @@ module Framework
 		}
 
 	#
-	# Create a simplified instance of the framework
+	# Create a simplified instance of the framework.  This routine takes a hash
+	# of parameters as an argument.  This hash can contain:
 	#
-	def self.create
+	#   OnCreateProc => A callback procedure that is called once the framework
+	#   instance is created.
+	#
+	def self.create(opts = {})
 		framework = Msf::Framework.new
 
-		return simplify(framework)
+		return simplify(framework, opts)
 	end
 
 	#
 	# Extends a framework object that may already exist
 	#
-	def self.simplify(framework)
+	def self.simplify(framework, opts)
 		framework.extend(Msf::Simple::Framework)
 
 		# Initialize the simplified framework
 		framework.init_simplified()
 
+		# Call the creation procedure if one was supplied
+		if (opts['OnCreateProc'])
+			opts['OnCreateProc'].call(framework)
+		end
+
 		# Initialize configuration and logging
 		Msf::Config.init
 		Msf::Logging.init

lib/msf/core/module_manager.rb

@@ -54,6 +54,13 @@ class ModuleSet < Hash
 		return instance
 	end
 
+	#
+	# Checks to see if the supplied module name is valid.
+	#
+	def valid?(name)
+		(self[name]) ? true : false
+	end
+
 	#
 	# Enumerates each module class in the set
 	#
@@ -63,6 +70,10 @@ class ModuleSet < Hash
 		each_module_list(mod_sorted, opts, &block)
 	end
 
+	#
+	# Enumerates each module class in the set based on their relative ranking
+	# to one another.  Modules that are ranked higher are shown first.
+	#
 	def each_module_ranked(opts = {}, &block)
 		mod_ranked = rank_modules if (mod_ranked == nil)
 

lib/msf/ui/console/command_dispatcher/core.rb

@@ -353,6 +353,12 @@ class Core
 		name  = args[0]
 		value = args[1]
 
+		# If the driver indicates that the value is not valid, bust out.
+		if (driver.on_variable_set(name, value) == false)
+			print_error("The value specified for #{name} is not valid.")
+			return true
+		end
+
 		datastore[name] = value
 
 		print_line("#{name} => #{value}")
@@ -440,6 +446,11 @@ class Core
 		end
 
 		while ((val = args.shift))
+			if (driver.on_variable_unset(val) == false)
+				print_error("The variable #{val} cannot be unset at this time.")
+				next
+			end
+
 			print_line("Unsetting #{val}...")
 
 			datastore.delete(val)
@@ -527,6 +538,13 @@ class Core
 		return true
 	end
 
+	#
+	# Internal routine to recalculate tab complete.
+	#
+	def cmd__recalculate_tc(*args)
+		recalculate_tab_complete
+	end
+
 protected
 
 	#

lib/msf/ui/console/driver.rb

@@ -46,7 +46,7 @@ class Driver < Msf::Ui::Driver
 
 		# Register event handlers
 		register_event_handlers
-		
+
 		# Process things before we actually display the prompt and get rocking
 		on_startup
 
@@ -112,10 +112,35 @@ class Driver < Msf::Ui::Driver
 	# displayed, scripts can be processed, and other fun can be had.
 	#
 	def on_startup
+		# Recalculate tab completion
+		run_single("_recalculate_tc")
+
 		# Build the banner message
 		run_single("banner")
 	end
 
+	#
+	# Called when a variable is set to a specific value.  This allows the
+	# console to do extra processing, such as enabling logging or doing
+	# some other kind of task.  If this routine returns false it will indicate
+	# that the variable is not being set to a valid value.
+	#
+	def on_variable_set(var, val)
+		case var.downcase
+			when "payload"
+				if (framework.modules.valid?(val) == false)
+					return false
+				end
+		end
+	end
+
+	#
+	# Called when a variable is unset.  If this routine returns false it is an
+	# indication that the variable should not be allowed to be unset.
+	#
+	def on_variable_unset(var)
+	end
+
 	attr_reader   :framework
 	attr_accessor :active_module