This patch allows the framework to be created with only specific modu…

…le types enabled, speeding up msfpayload and msfencode.

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

lib/msf/base/simple/framework.rb

@@ -66,8 +66,7 @@ def on_module_created(instance)
 	#   instance is created.
 	#
 	def self.create(opts = {})
-		framework = Msf::Framework.new
-
+		framework = Msf::Framework.new(opts)
 		return simplify(framework, opts)
 	end
 
@@ -84,7 +83,7 @@ def self.simplify(framework, opts)
 
 		# Initialize the simplified framework
 		framework.init_simplified()
-
+		
 		# Call the creation procedure if one was supplied
 		if (opts['OnCreateProc'])
 			opts['OnCreateProc'].call(framework)
@@ -94,6 +93,7 @@ def self.simplify(framework, opts)
 		Msf::Config.init
 		Msf::Logging.init
 
+
 		# Load the configuration
 		framework.load_config
 
@@ -178,4 +178,4 @@ def save_config
 end
 
 end
-end
+end

lib/msf/core/event_dispatcher.rb

@@ -221,4 +221,4 @@ def remove_event_subscriber(array, subscriber) # :nodoc:
 
 end
 
-end
+end

lib/msf/core/framework.rb

@@ -39,9 +39,13 @@ module Offspring
 	#
 	# Creates an instance of the framework context.
 	#
-	def initialize()
+	def initialize(opts={})
+
+		# Allow specific module types to be loaded
+		types = opts[:module_types] || MODULE_TYPES
+
 		self.events    = EventDispatcher.new(self)
-		self.modules   = ModuleManager.new(self)
+		self.modules   = ModuleManager.new(self,types)
 		self.sessions  = SessionManager.new(self)
 		self.datastore = DataStore.new
 		self.jobs      = Rex::JobContainer.new

lib/msf/core/module_manager.rb

@@ -325,17 +325,20 @@ class ModuleManager < ModuleSet
 
 	#
 	# Initializes an instance of the overall module manager using the supplied
-	# framework instance.
+	# framework instance. The types parameter can be used to only load specific
+	# module types on initialization
 	#
-	def initialize(framework)
+	def initialize(framework,types=MODULE_TYPES)
 		self.module_paths         = []
 		self.module_history       = {}
 		self.module_history_mtime = {}
 		self.module_sets          = {}
 		self.module_failed        = {}
+		self.enabled_types        = {}
 		self.framework            = framework
 
-		MODULE_TYPES.each { |type|
+		types.each { |type|
+			self.enabled_types[type] = true
 			case type
 				when MODULE_PAYLOAD
 					instance = PayloadSet.new(self)
@@ -444,7 +447,9 @@ def set_module_cache_file(file_path)
 			@modcache.add_group(type, false)
 
 			@modcache[type].each_key { |name|
-
+				next if not @modcache[type]
+				next if not module_sets[type]
+
 				fullname = type + '/' + name
 
 				# Make sure the files associated with this module exist.  If it
@@ -508,6 +513,7 @@ def save_module_cache
 				@modcache['ModuleTypeCounts'].clear
 
 				MODULE_TYPES.each { |type|
+					next if not @modcache['ModuleTypeCounts'][type]
 					@modcache['ModuleTypeCounts'][type] = module_sets[type].length.to_s
 				}
 			end
@@ -797,38 +803,47 @@ def module_names(set)
 	# Load all of the modules from the supplied module path (independent of
 	# module type).
 	#
-	def load_modules(path, demand = false)
+	def load_modules(bpath, demand = false)
 		loaded = {}
 		recalc = {}
 		counts = {}
 		delay  = {}
 		ks     = true
+
+		dbase  = Dir.new(bpath)
+		dbase.entries.each do |ent|
+			path  = File.join(bpath, ent)
+			mtype = ent.gsub(/s$/, '')
+			next if not File.directory?(path)
+			next if not MODULE_TYPES.include?(mtype)
+			next if not enabled_types[mtype]
 
-		# Try to load modules from all the files in the supplied path
-		Rex::Find.find(path) { |file|
-
-			# Skip non-ruby files
-			next if (file !~ /\.rb$/i)
-
-			# Skip unit test files
-			next if (file =~ /rb\.(ut|ts)\.rb$/)
+			# Try to load modules from all the files in the supplied path
+			Rex::Find.find(path) do |file|
 
-			# Skip files with a leading period
-			next if (file =~ /^\./)
+				# Skip non-ruby files
+				next if (file !~ /\.rb$/i)
 
-			begin
-				load_module_from_file(path, file, loaded, recalc, counts, demand)
-			rescue NameError
-
-				# As of Jan-06-2007 this code isn't hit with the official module tree
-
-				# If we get a name error, it's possible that this module depends
-				# on another one that we haven't loaded yet.  Let's postpone
-				# the load operation for now so that we can resolve all 
-				# dependencies.  This is pretty much a hack.
-				delay[file] = $!
+				# Skip unit test files
+				next if (file =~ /rb\.(ut|ts)\.rb$/)
+
+				# Skip files with a leading period
+				next if (file =~ /^\./)
+
+				begin
+					load_module_from_file(bpath, file, loaded, recalc, counts, demand)
+				rescue NameError
+
+					# As of Jan-06-2007 this code isn't hit with the official module tree
+
+					# If we get a name error, it's possible that this module depends
+					# on another one that we haven't loaded yet.  Let's postpone
+					# the load operation for now so that we can resolve all 
+					# dependencies.  This is pretty much a hack.
+					delay[file] = $!
+				end
 			end
-		}
+		end
 
 		# Keep processing all delayed module loads until we've gotten
 		# all the dependencies resolved or until we just suck.
@@ -1049,6 +1064,7 @@ def auto_subscribe_module(mod)
 	attr_accessor :module_paths # :nodoc:
 	attr_accessor :module_history, :module_history_mtime # :nodoc:
 	attr_accessor :module_failed # :nodoc:
+	attr_accessor :enabled_types # :nodoc:
 
 end
 

msfencode

@@ -76,7 +76,9 @@ def usage
 end
 
 # Initialize the simplified framework instance.
-$framework = Msf::Simple::Framework.create
+$framework = Msf::Simple::Framework.create(
+	:module_types => [ Msf::MODULE_ENCODER ]
+)
 
 # Defaults
 cmd      = "encode"

msfpayload

@@ -29,7 +29,10 @@ def dump_payloads
 end
 
 # Initialize the simplified framework instance.
-$framework = Msf::Simple::Framework.create
+$framework = Msf::Simple::Framework.create(
+	:module_types => [ Msf::MODULE_PAYLOAD ]
+)
+
 
 if (ARGV.length <= 1)
 	puts "\n" + "    Usage: #{$0} <payload> [var=val] <S[ummary]|C|P[erl]|[Rub]y|R[aw]|J[avascript]|e[X]ecutable|[V]BA>\n"