Land rapid7#2154, @wchen-r7's msfcli optimizations and refactoring

Author: jvazquez-r7

lib/msf/base/simple/framework/module_paths.rb

@@ -5,25 +5,25 @@ module ModulePaths
 				# Initialize the module paths
 				#
 				# @return [void]
-				def init_module_paths
+				def init_module_paths(opts={})
 					# Ensure the module cache is accurate
 					self.modules.refresh_cache_from_database
 
 					# Initialize the default module search paths
 					if (Msf::Config.module_directory)
-						self.modules.add_module_path(Msf::Config.module_directory)
+						self.modules.add_module_path(Msf::Config.module_directory, opts)
 					end
 
 					# Initialize the user module search path
 					if (Msf::Config.user_module_directory)
-						self.modules.add_module_path(Msf::Config.user_module_directory)
+						self.modules.add_module_path(Msf::Config.user_module_directory, opts)
 					end
 
 					# If additional module paths have been defined globally, then load them.
 					# They should be separated by semi-colons.
 					if self.datastore['MsfModulePaths']
 						self.datastore['MsfModulePaths'].split(";").each { |path|
-							self.modules.add_module_path(path)
+							self.modules.add_module_path(path, opts)
 						}
 					end
 				end

lib/msf/core/module_manager/loading.rb

@@ -53,39 +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 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
+  # 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
 
@@ -106,9 +106,10 @@ def loaders
   # @param [Hash] options
   # @option options [Boolean] :force Whether the force loading the modules even if they are unchanged and already
   #   loaded.
+  # @option options [Array] :modules An array of regex patterns to search for specific modules
   # @return [Hash{String => Integer}] Maps module type to number of modules loaded
   def load_modules(path, options={})
-    options.assert_valid_keys(:force)
+    options.assert_valid_keys(:force, :whitelist)
 
     count_by_type = {}
 
@@ -122,4 +123,4 @@ def load_modules(path, options={})
 
     count_by_type
   end
-end
+end

lib/msf/core/module_manager/module_paths.rb

@@ -12,8 +12,10 @@ module Msf::ModuleManager::ModulePaths
   # Adds a path to be searched for new modules.
   #
   # @param [String] path
+  # @param [Hash] opts
+  # @option opts [Array] whitelist An array of regex patterns to search for specific modules
   # @return (see Msf::Modules::Loader::Base#load_modules)
-  def add_module_path(path)
+  def add_module_path(path, opts={})
     nested_paths = []
 
     # remove trailing file separator
@@ -51,7 +53,7 @@ def add_module_path(path)
     # Load all of the modules from the nested paths
     count_by_type = {}
     nested_paths.each { |path|
-      path_count_by_type = load_modules(path, :force => false)
+      path_count_by_type = load_modules(path, opts.merge({:force => false}))
 
       # merge hashes
       path_count_by_type.each do |type, path_count|
@@ -74,4 +76,4 @@ def remove_module_path(path)
   protected
 
   attr_accessor :module_paths # :nodoc:
-end
+end

lib/msf/core/modules/loader/archive.rb

@@ -27,10 +27,12 @@ def loadable?(path)
   # Yields the module_reference_name for each module file in the Fastlib archive at path.
   #
   # @param path [String] The path to the Fastlib archive file.
+  # @param opts [Hash] Additional options
   # @yield (see Msf::Modules::Loader::Base#each_module_reference_name)
   # @yieldparam (see Msf::Modules::Loader::Base#each_module_reference_name)
   # @return (see Msf::Modules::Loader::Base#each_module_reference_name)
-  def each_module_reference_name(path)
+  def each_module_reference_name(path, opts={})
+    whitelist = opts[:whitelist] || []
     entries = ::FastLib.list(path)
 
     entries.each do |entry|
@@ -45,11 +47,22 @@ def each_module_reference_name(path)
         next
       end
 
-      if module_path?(entry)
-        # The module_reference_name doesn't have a file extension
-        module_reference_name = module_reference_name_from_path(entry)
+      if whitelist.empty?
 
-        yield path, type, module_reference_name
+        if module_path?(entry)
+          # The module_reference_name doesn't have a file extension
+          module_reference_name = module_reference_name_from_path(entry)
+
+          yield path, type, module_reference_name
+        end
+      else
+        whitelist.each do |pattern|
+          if entry =~ pattern
+            yield path, type, module_reference_name
+          else
+            next
+          end
+        end
       end
     end
   end

lib/msf/core/modules/loader/base.rb

@@ -248,16 +248,25 @@ def load_module(parent_path, type, module_reference_name, options={})
   # @param [Hash] options
   # @option options [Boolean] force (false) Whether to force loading of
   #   the module even if the module has not changed.
+  # @option options [Array] whitelist An array of regex patterns to search for specific modules
   # @return [Hash{String => Integer}] Maps module type to number of
   #   modules loaded
   def load_modules(path, options={})
-    options.assert_valid_keys(:force)
+    options.assert_valid_keys(:force, :whitelist)
 
     force = options[:force]
     count_by_type = {}
     recalculate_by_type = {}
 
-    each_module_reference_name(path) do |parent_path, type, module_reference_name|
+    # This is used to avoid loading the same thing twice
+    loaded_items = []
+
+    each_module_reference_name(path, options) do |parent_path, type, module_reference_name|
+      # In msfcli mode, if a module is already loaded, avoid loading it again
+      next if loaded_items.include?(module_reference_name) and options[:whitelist]
+
+      # Keep track of loaded modules in msfcli mode
+      loaded_items << module_reference_name if options[:whitelist]
       load_module(
           parent_path,
           type,
@@ -649,4 +658,3 @@ def usable?(metasploit_class)
     usable
   end
 end
-

lib/msf/core/modules/loader/directory.rb

@@ -18,12 +18,14 @@ def loadable?(path)
   # Yields the module_reference_name for each module file found under the directory path.
   #
   # @param [String] path The path to the directory.
+  # @param [Array] modules An array of regex patterns to search for specific modules
   # @yield (see Msf::Modules::Loader::Base#each_module_reference_name)
   # @yieldparam [String] path The path to the directory.
   # @yieldparam [String] type The type correlated with the directory under path.
   # @yieldparam module_reference_name (see Msf::Modules::Loader::Base#each_module_reference_name)
   # @return (see Msf::Modules::Loader::Base#each_module_reference_name)
-  def each_module_reference_name(path)
+  def each_module_reference_name(path, opts={})
+    whitelist = opts[:whitelist] || []
     ::Dir.foreach(path) do |entry|
       if entry.downcase == '.svn'
         next
@@ -49,7 +51,24 @@ def each_module_reference_name(path)
           # The module_reference_name doesn't have a file extension
           module_reference_name = module_reference_name_from_path(relative_entry_descendant_path)
 
-          yield path, type, module_reference_name
+          # If the modules argument is set, this means we only want to load specific ones instead
+          # of loading everything to memory - see msfcli.
+          if whitelist.empty?
+            # Load every module we see, which is the default behavior.
+            yield path, type, module_reference_name
+          else
+              whitelist.each do |pattern|
+              # We have to use entry_descendant_path to see if this is the module we want, because
+              # this is easier to identify the module type just by looking at the file path.
+              # For example, if module_reference_name is used (or a parsed relative path), you can't
+              # really tell if php/generic is a NOP module, a payload, or an encoder.
+              if entry_descendant_path =~ pattern
+                yield path, type, module_reference_name
+              else
+                next
+              end
+            end
+          end
         end
       end
     end
@@ -76,18 +95,18 @@ def read_module_content(parent_path, type, module_reference_name)
     module_content = ''
 
     begin
-	    # force to read in binary mode so Pro modules won't be truncated on Windows
-	    File.open(full_path, 'rb') do |f|
-		    # Pass the size of the file as it leads to faster reads due to fewer buffer resizes. Greatest effect on Windows.
-		    # @see http://www.ruby-forum.com/topic/209005
-		    # @see https://github.com/ruby/ruby/blob/ruby_1_8_7/io.c#L1205
-		    # @see https://github.com/ruby/ruby/blob/ruby_1_9_3/io.c#L2038
-		    module_content = f.read(f.stat.size)
-	    end
+      # force to read in binary mode so Pro modules won't be truncated on Windows
+      File.open(full_path, 'rb') do |f|
+        # Pass the size of the file as it leads to faster reads due to fewer buffer resizes. Greatest effect on Windows.
+        # @see http://www.ruby-forum.com/topic/209005
+        # @see https://github.com/ruby/ruby/blob/ruby_1_8_7/io.c#L1205
+        # @see https://github.com/ruby/ruby/blob/ruby_1_9_3/io.c#L2038
+        module_content = f.read(f.stat.size)
+      end
     rescue Errno::ENOENT => error
-	    load_error(full_path, error)
+      load_error(full_path, error)
     end
 
     module_content
   end
-end
+end