Payload compatible cache_in_memory

[#47720609]

Msf::PayloadSet#add_module does NOT return an annotated module class as
Msf::ModuleSet#add_module does because a payload module is defined as a
ruby Module instead of a ruby Class.   Since add_module doesn't always
return an annotated_class, the logic in
Msf::ModuleManager#on_module_load needed to change to NOT use
annotated_class and create #add_module as return [void].  Thus, it is
necessary to pass in all the metasploit module metadata to
Msf::ModuleManager#cache_in_memory instead of assuming they can be
derived from the (payload) Module or (other) Class.

Author: limhoff-r7

lib/msf/core/module_manager/cache.rb

@@ -15,20 +15,39 @@ def cache_empty?
     module_info_by_path.empty?
   end
 
+	# @note path, reference_name, and type must be passed as options because when +class_or_module+ is a payload Module,
+	#   those attributes will either not be set or not exist on the module.
+	#
 	# Updates the in-memory cache so that {#file_changed?} will report +false+ if
 	# the module is loaded again.
 	#
-	# @param klass [Class<Msf::Module>] a module class
+	# @param class_or_module [Class<Msf::Module>, ::Module] either a module Class
+	#   or a payload Module.
+	# @param options [Hash{Symbol => String}]
+	# @option options [String] :path the path to the file from which
+	#   +class_or_module+ was loaded.
+	# @option options [String] :reference_name the reference name for
+	#   +class_or_module+.
+	# @option options [String] :type the module type
 	# @return [void]
-	def cache_in_memory(klass)
-		modification_time = File.mtime(klass.file_path)
-		parent_path = klass.parent.parent_path
+	# @raise [KeyError] unless +:path+ is given.
+	# @raise [KeyError] unless +:reference_name+ is given.
+	# @raise [KeyError] unless +:type+ is given.
+	def cache_in_memory(class_or_module, options={})
+		options.assert_valid_keys(:path, :reference_name, :type)
+
+		path = options.fetch(:path)
+		modification_time = File.mtime(path)
+
+		parent_path = class_or_module.parent.parent_path
+		reference_name = options.fetch(:reference_name)
+		type = options.fetch(:type)
 
-		module_info_by_path[klass.file_path] = {
+		module_info_by_path[path] = {
 				:modification_time => modification_time,
 				:parent_path => parent_path,
-				:reference_name => klass.refname,
-				:type => klass.type
+				:reference_name => reference_name,
+				:type => type
 		}
 	end
 

lib/msf/core/module_manager/loading.rb

@@ -53,32 +53,39 @@ def file_changed?(path)
 
   attr_accessor :module_load_error_by_path
 
-  # Called when a module is initially loaded such that it can be categorized
-  # accordingly.
-  #
-  # @param klass (see Msf::ModuleSet#add_module)
-  # @param type [String] The module type.
-  # @param reference_name (see Msf::ModuleSet#add_module)
-  # @param info [Hash{String => Array}] additional information about the module
-  # @option info [Array<String>] 'files' List of paths to the ruby source files
-  #   where +klass+ is defined.
-  # @option info [Array<String>] 'paths' List of module reference names.
-  # @option info [String] 'type' The module type, should match positional
-  #   +type+ argument.
-  # @return [void]
-  def on_module_load(klass, type, reference_name, info={})
-    module_set = module_set_by_type[type]
-    annotated_klass = module_set.add_module(klass, reference_name, info)
-
-    cache_in_memory(annotated_klass)
-
-    # Automatically subscribe a wrapper around this module to the necessary
-    # event providers based on whatever events it wishes to receive.
-    auto_subscribe_module(annotated_klass)
-
-    # Notify the framework that a module was loaded
-    framework.events.on_module_load(reference_name, annotated_klass)
-  end
+	# Called when a module is initially loaded such that it can be categorized
+	# accordingly.
+	#
+	# @param class_or_module [Class<Msf::Module>, ::Module] either a module Class
+	#   or a payload Module.
+	# @param type [String] The module type.
+	# @param reference_name The module reference name.
+	# @param info [Hash{String => Array}] additional information about the module
+	# @option info [Array<String>] 'files' List of paths to the ruby source files
+	#   where +class_or_module+ is defined.
+	# @option info [Array<String>] 'paths' List of module reference names.
+	# @option info [String] 'type' The module type, should match positional
+	#   +type+ argument.
+	# @return [void]
+	def on_module_load(class_or_module, type, reference_name, info={})
+		module_set = module_set_by_type[type]
+		module_set.add_module(class_or_module, reference_name, info)
+
+		path = info['files'].first
+		cache_in_memory(
+				class_or_module,
+				:path => path,
+				:reference_name => reference_name,
+				:type => type
+		)
+
+		# Automatically subscribe a wrapper around this module to the necessary
+		# event providers based on whatever events it wishes to receive.
+		auto_subscribe_module(class_or_module)
+
+		# Notify the framework that a module was loaded
+		framework.events.on_module_load(reference_name, class_or_module)
+	end
 
   protected
 

lib/msf/core/payload_set.rb

@@ -169,29 +169,38 @@ def recalculate
 		flush_blob_cache
 	end
 
-	#
 	# This method is 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 add it to the appropriate list.
 	#
-	def add_module(pmodule, name, modinfo = nil)
-
-		if (md = name.match(/^(singles|stagers|stages)#{File::SEPARATOR}(.*)$/))
+	# @param payload_module [::Module] The module name.
+	# @param reference_name [String] The module reference name.
+	# @param modinfo [Hash{String => Array}] additional information about the
+	#   module.
+	# @option modinfo [Array<String>] 'files' List of paths to the ruby source
+	#   files where +class_or_module+ is defined.
+	# @option modinfo [Array<String>] 'paths' List of module reference names.
+	# @option modinfo [String] 'type' The module type, should match positional
+	#   +type+ argument.
+	# @return [void]
+	def add_module(payload_module, reference_name, modinfo={})
+
+		if (md = reference_name.match(/^(singles|stagers|stages)#{File::SEPARATOR}(.*)$/))
 			ptype = md[1]
-			name  = md[2]
+			reference_name  = md[2]
 		end
 
 		# Duplicate the Payload base class and extend it with the module
 		# class that is passed in.  This allows us to inspect the actual
 		# module to see what type it is, and to grab other information for
 		# our own evil purposes.
-		instance = build_payload(pmodule).new
+		instance = build_payload(payload_module).new
 
 		# Create an array of information about this payload module
 		pinfo =
 			[
-				pmodule,
+				payload_module,
 				instance.handler_klass,
 				instance.platform,
 				instance.arch,
@@ -200,26 +209,12 @@ def add_module(pmodule, name, modinfo = nil)
 			]
 
 		# Use the module's preferred alias if it has one
-		name = instance.alias if (instance.alias)
+		reference_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][name] = pinfo
-
-		#
-		# Disable sending singles over stagers for now
-		#
-=begin
-		# 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_klass == Msf::Handler::None) or
-		     (instance.handler_klass == nil)))
-			payload_type_modules[Payload::Type::Stage][name] = pinfo
-		end
-=end
+		payload_type_modules[instance.payload_type][reference_name] = pinfo
 	end
 
 	#

spec/lib/msf/core/modules/loader/directory_spec.rb

@@ -62,8 +62,34 @@
 						subject.load_module(parent_path, type, module_reference_name)
 					end
 
-					it 'should not load the module' do
-						subject.load_module(parent_path, type, module_reference_name).should be_false
+					# Payloads are defined as ruby Modules so they can behave differently
+					context 'with payload' do
+						let(:reference_name) do
+							'stages/windows/x64/vncinject'
+						end
+
+						let(:type) do
+							'payload'
+						end
+
+						it 'should not load the module' do
+							subject.load_module(parent_path, type, module_reference_name).should be_false
+						end
+					end
+
+					# Non-payloads are defined as ruby Classes
+					context 'without payload' do
+						let(:reference_name) do
+							'windows/smb/ms08_067_netapi'
+						end
+
+						let(:type) do
+							'exploit'
+						end
+
+						it 'should not load the module' do
+							subject.load_module(parent_path, type, module_reference_name).should be_false
+						end
 					end
 				end
 			end

spec/support/shared/examples/msf/module_manager/cache.rb

@@ -60,21 +60,20 @@
 
 	context '#cache_in_memory' do
 		def cache_in_memory
-			module_manager.cache_in_memory(klass)
+			module_manager.cache_in_memory(
+					class_or_module,
+					:path => path,
+					:reference_name => reference_name,
+					:type => type
+			)
 		end
 
 		def module_info_by_path
 			module_manager.send(:module_info_by_path)
 		end
 
-		let(:klass) do
-			mock(
-					'Class<Msf::Module>',
-					:file_path => path,
-					:parent => namespace_module,
-					:refname => reference_name,
-					:type => type
-			)
+		let(:class_or_module) do
+			mock('Class<Msf::Module> or Module', :parent => namespace_module)
 		end
 
 		let(:namespace_module) do
@@ -96,29 +95,29 @@ def module_info_by_path
 				cache_in_memory
 			end
 
-			it 'should have entry for file_path' do
-				module_info_by_path[klass.file_path].should be_a Hash
+			it 'should have entry for path' do
+				module_info_by_path[path].should be_a Hash
 			end
 
 			context 'value' do
 				subject(:value) do
-					module_info_by_path[klass.file_path]
+					module_info_by_path[path]
 				end
 
-				it 'should have modification time of file_path for :modification_time' do
+				it 'should have modification time of :path option for :modification_time' do
 					value[:modification_time].should == pathname_modification_time
 				end
 
 				it 'should have parent path from namespace module for :parent_path' do
 					value[:parent_path].should == namespace_module.parent_path
 				end
 
-				it 'should have refname for :reference_name' do
-					value[:reference_name].should == klass.refname
+				it 'should use :reference_name option' do
+					value[:reference_name].should == reference_name
 				end
 
-				it 'should have type for :type' do
-					value[:type].should == klass.type
+				it 'should use :type option' do
+					value[:type].should == type
 				end
 			end
 		end

spec/support/shared/examples/msf/module_manager/loading.rb

@@ -78,26 +78,10 @@
 	end
 
 	context '#on_module_load' do
-		def module_info_by_path
-			module_manager.send(:module_info_by_path)
-		end
-
-	  def on_module_load
+		def on_module_load
 			module_manager.on_module_load(klass, type, reference_name, options)
 		end
 
-		let(:annotated_class) do
-			# stub out include? so calls in auto_subscribe_module work.
-			mock(
-					'Annotated Class',
-					:file_path => path,
-					:include? => false,
-					:parent => namespace_module,
-					:refname => reference_name,
-					:type => type
-			)
-		end
-
 		let(:klass) do
 			Class.new(Msf::Auxiliary)
 		end
@@ -149,7 +133,7 @@ def on_module_load
 					klass,
 					reference_name,
 					options
-			).and_return(annotated_class)
+			)
 
 			on_module_load
 		end
@@ -160,25 +144,19 @@ def on_module_load
 			on_module_load
 		end
 
-		context 'annotated class' do
-			before(:each) do
-				module_set.stub(:add_module).and_return(annotated_class)
-			end
-
-			it 'should pass annotated class to Msf::ModuleManager#auto_subscribe_module' do
-				module_manager.should_receive(:auto_subscribe_module).with(annotated_class)
+		it 'should pass class to #auto_subscribe_module' do
+			module_manager.should_receive(:auto_subscribe_module).with(klass)
 
-				on_module_load
-			end
+			on_module_load
+		end
 
-			it 'should fire on_module_load event' do
-				framework.events.should_receive(:on_module_load).with(
-						reference_name,
-						annotated_class
-				)
+		it 'should fire on_module_load event with class' do
+			framework.events.should_receive(:on_module_load).with(
+					reference_name,
+					klass
+			)
 
-				on_module_load
-			end
+			on_module_load
 		end
 	end
 end