some readline stuff, tab completion, add spoon's hashed payload stuff

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

lib/msf/core/module_manager.rb

@@ -56,11 +56,24 @@ class ModuleSet < Hash
 				next if (mod_platform_hash[mod].include?(opts['platform']) == false)
 			end
 
+			# Custom filtering
+			next if (each_module_filter(opts, name, entry) == true)
+
 			block.call(name, mod)
 		}
 	end
 
+	#
+	# Custom each_module filtering if an advanced set supports doing extended
+	# filtering.  Returns true if the entry should be filtered.
+	#
+	def each_module_filter(opts, name, entry)
+		return false
+	end
+
+	#
 	# Dummy placeholder to relcalculate aliases and other fun things
+	#
 	def recalculate
 	end
 

lib/msf/core/payload.rb

@@ -27,6 +27,9 @@ class Payload < Msf::Module
 		Stage  = (1 << 2)
 	end
 
+	class <<self
+	end
+
 	def initialize(info = {})
 		super
 

lib/msf/core/payload_set.rb

@@ -34,8 +34,29 @@ class PayloadSet < ModuleSet
 			Payload::Type::Stager,
 			Payload::Type::Stage
 		].each { |type|
-			self.payload_type_modules[type] = []
+			self.payload_type_modules[type] = {}
 		}
+
+		# Initialize hashes for each of the stages and singles.  Stagers
+		# never exist independent.  The stages hash will have entries that
+		# point to another hash that point to the per-stager implementation
+		# payload class.  For instance:
+		#
+		# ['windows/shell']['reverse_tcp']
+		#
+		# Singles will simply point to the single payload class.
+		self.stages  = {}
+		self.singles = {}
+	end
+
+	#
+	# Performs custom filtering during each_module enumeration.  This allows us
+	# to filter out certain stagers as necessary.
+	#
+	# TODO: stager-based customf iltering
+	#
+	def each_module_filter(opts, name, mod)
+		return false
 	end
 
 	# Build the actual hash of alias names based on all the permutations
@@ -48,8 +69,8 @@ class PayloadSet < ModuleSet
 		self.clear
 
 		# Recalculate single payloads
-		singles.each { |p|
-			mod, name, handler = p
+		_singles.each_pair { |name, p|
+			mod, handler = p
 
 			# Build the payload dupe using the determined handler
 			# and module
@@ -58,21 +79,17 @@ class PayloadSet < ModuleSet
 			# Sets the modules derived name
 			p.refname = name
 
-			# Associate this class with the single payload's name
-			self[name] = p
-
-			manager.add_module(p, name)
-
-			dlog("Built single payload #{name}.", 'core', LEV_1)
+			# Add it to the set
+			add_single(p, name)
 		}
 
 		# Recalculate stagers and stages
-		stagers.each { |p|
-			stager_mod, stager_name, handler, stager_platform, stager_arch = p
+		_stagers.each_pair { |stager_name, p|
+			stager_mod, handler, stager_platform, stager_arch = p
 
 			# Walk the array of stages
-			stages.each { |p|
-				stage_mod, stage_name, junk, stage_platform, stage_arch = p
+			_stages.each_pair { |stage_name, p|
+				stage_mod, junk, stage_platform, stage_arch = p
 
 				# No intersection between architectures on the payloads?
 				if ((stager_arch) and
@@ -102,35 +119,17 @@ class PayloadSet < ModuleSet
 				combined  = stage_name
 
 				# If a valid handler exists for this stager, then combine it
-				combined += '/stg/' + handler.handler_type if (handler)
+				combined += '/' + handler.handler_type
 
 				# Sets the modules derived name
 				p.refname = combined
 
-				self[combined] = p
-
-				manager.add_module(p, combined)
-			
-				dlog("Built staged payload #{combined}.", 'core', LEV_1)
+				# Add the stage
+				add_stage(p, combined, stage_name, handler.handler_type)
 			}
 		}
 	end
 
-	# Return the array of single payloads
-	def singles
-		return payload_type_modules[Payload::Type::Single] || []
-	end
-
-	# Return the array of stager payloads
-	def stagers
-		return payload_type_modules[Payload::Type::Stager] || []
-	end
-
-	# Return the array of stage payloads
-	def stages
-		return payload_type_modules[Payload::Type::Stage] || []
-	end
-
 	# Called when a new payload module class is loaded up.  For the payload
 	# set we simply create an instance of the class and do some magic to figure
 	# out if it's a single, stager, or stage.  Depending on which it is, we 
@@ -150,28 +149,84 @@ class PayloadSet < ModuleSet
 		pinfo = 
 			[
 				pmodule,
-				instance.alias || name,
 				instance.handler,
 				instance.platform,
 				instance.arch
 			]
 
+		# Use the module's preferred alias if it has one
+		name = instance.alias if (instance.alias)
+
 		# Store the module and alias name for this payload.  We
 		# also convey other information about the module, such as
 		# the platforms and architectures it supports
-		payload_type_modules[instance.payload_type] << pinfo
+		payload_type_modules[instance.payload_type][name] = pinfo
 
 		# If the payload happens to be a single, but has no defined
 		# connection, then it can also be staged.  Insert it into
 		# the staged list.
 		if ((instance.payload_type == Payload::Type::Single) and
 		    (instance.handler == nil))
-			payload_type_modules[Payload::Type::Stage] << pinfo
+			payload_type_modules[Payload::Type::Stage][name] = pinfo
 		end
 	end
 
+	#
+	# Adds a single payload to the set and adds it to the singles hash
+	#
+	def add_single(p, name)
+		# Associate this class with the single payload's name
+		self[name] = p
+
+		# Add the singles hash
+		singles[name] = p
+
+		# Add it to the global module set
+		manager.add_module(p, name)
+
+		dlog("Built single payload #{name}.", 'core', LEV_1)
+	end
+
+	#
+	# Adds a stage payload to the set and adds it to the stages hash
+	# using the supplied handler type.
+	#
+	def add_stage(p, full_name, stage_name, handler_type)
+		# Associate this stage's full name with the payload class in the set
+		self[full_name] = p
+
+		# Add the full name association in the global module set
+		manager.add_module(p, full_name)
+
+		# Create the hash entry for this stage and then create
+		# the associated entry for the handler type
+		stages[stage_name] = {} if (!stages[stage_name])
+
+		# Add it to this stage's stager hash
+		stages[stage_name][handler_type] = p
+			
+		dlog("Built staged payload #{full_name}.", 'core', LEV_1)
+	end
+
+	attr_reader :stages, :singles
+
 protected
 
+	# Return the hash of single payloads
+	def _singles
+		return payload_type_modules[Payload::Type::Single] || {}
+	end
+
+	# Return the hash of stager payloads
+	def _stagers
+		return payload_type_modules[Payload::Type::Stager] || {}
+	end
+
+	# Return the hash of stage payloads
+	def _stages
+		return payload_type_modules[Payload::Type::Stage] || {}
+	end
+
 	# Builds a duplicate, extended version of the Payload base
 	# class using the supplied modules.
 	def build_payload(*modules)
@@ -187,8 +242,8 @@ protected
 		return klass
 	end
 
-
 	attr_accessor :manager, :payload_type_modules
+	attr_writer   :stages, :singles
 
 end
 

lib/msf/ui/console/command_dispatcher.rb

@@ -6,6 +6,7 @@ module CommandDispatcher
 
 	def initialize(in_driver)
 		self.driver = in_driver
+		self.tab_complete_items = []
 	end
 
 	def print_error(msg = '')
@@ -40,6 +41,11 @@ module CommandDispatcher
 		return driver.datastore['_ActiveModule']
 	end
 
+	#
+	# No tab completion items by default
+	#
+	attr_accessor :tab_complete_items
+
 protected
 
 	attr_accessor :driver

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

@@ -129,6 +129,8 @@ class Core
 		}
 
 		print_line("Added #{args.length} search paths.")
+
+		recalculate_tab_complete
 	end
 
 	#
@@ -314,6 +316,17 @@ class Core
 
 protected
 
+	#
+	# Recalculates the tab completion list
+	#
+	def recalculate_tab_complete
+		self.tab_complete_items = []
+
+		framework.modules.each_module { |refname, mod|
+			self.tab_complete_items << refname
+		}
+	end
+
 	#
 	# Module list enumeration
 	#

lib/msf/ui/console/driver.rb

@@ -3,8 +3,8 @@ require 'msf/base'
 require 'msf/ui'
 require 'msf/ui/console/shell'
 require 'msf/ui/console/command_dispatcher'
-
 require 'msf/ui/console/table'
+require 'find'
 
 module Msf
 module Ui
@@ -23,7 +23,7 @@ class Driver < Msf::Ui::Driver
 
 	include Msf::Ui::Console::Shell
 
-	def initialize(prompt = "%umsf", prompt_char = ">%c%b")
+	def initialize(prompt = "%umsf", prompt_char = ">%c")
 		# Initialize attributes
 		self.framework        = Msf::Framework.new
 		self.dispatcher_stack = []
@@ -36,6 +36,29 @@ class Driver < Msf::Ui::Driver
 		super
 	end
 
+	#
+	# Performs tab completion on shell input if supported
+	#
+	def tab_complete(str)
+		items = []
+
+		# Next, try to match internal command or value completion
+		# Enumerate each entry in the dispatcher stack
+		dispatcher_stack.each { |dispatcher|
+			# If it supports commands, query them all
+			if (dispatcher.respond_to?('commands'))
+				items.concat(dispatcher.commands.to_a.map { |x| x[0] })
+			end
+
+			# If the dispatcher has custom tab completion items, use them
+			items.concat(dispatcher.tab_complete_items || [])
+		}
+
+		items.find_all { |e| 
+			e =~ /^#{str}/
+		}
+	end
+
 	# Run a single command line
 	def run_single(line)
 		arguments = parse_line(line)

lib/msf/ui/console/input_methods.rb

@@ -15,7 +15,7 @@ module Console
 #
 ###
 class InputMethod
-	def initialize
+	def initialize(tab_complete_proc = nil)
 		self.eof = false
 	end
 
@@ -47,6 +47,12 @@ begin
 	class ReadlineInputMethod < InputMethod
 		include Readline
 
+		def initialize(tab_complete_proc = nil)
+			if (tab_complete_proc)
+				Readline.completion_proc = tab_complete_proc
+			end
+		end
+
 		def gets
 			if ((line = readline(prompt, true)))
 				HISTORY.pop if (line.empty?)

lib/msf/ui/console/shell.rb

@@ -23,7 +23,7 @@ module Shell
 		self.output = StdioOutputMethod.new
 
 		begin
-			self.input = ReadlineInputMethod.new
+			self.input = ReadlineInputMethod.new(lambda { |str| tab_complete(str) })
 		rescue
 		end
 
@@ -39,6 +39,13 @@ module Shell
 		super()
 	end
 
+	#
+	# Performs tab completion on the supplied string
+	#
+	def tab_complete(str)
+		return tab_complete_proc(str) if (tab_complete_proc)
+	end
+
 	#
 	# Run the command processing loop
 	#
@@ -149,7 +156,7 @@ protected
 
 
 	attr_accessor :input, :output, :stop_flag, :init_prompt
-	attr_accessor :prompt_char
+	attr_accessor :prompt_char, :tab_complete_proc
 
 end