Merge pull request #2 from todb-r7/bug/modulemanager-each

Bug/modulemanager each

Author: jlee-r7

lib/msf/base/simple/exploit.rb

@@ -69,8 +69,7 @@ def self.exploit_simple(oexploit, opts, &block)
 
 			# Make sure parameters are valid.
 			if (opts['Payload'] == nil)
-				raise MissingPayloadError,
-					"You must specify a payload.", caller
+				raise MissingPayloadError.new, caller
 			end
 
 			# Verify the options
@@ -81,7 +80,7 @@ def self.exploit_simple(oexploit, opts, &block)
 
 			# Initialize the driver instance
 			driver.exploit    = exploit
-			driver.payload    = exploit.framework.modules.create(opts['Payload'])
+			driver.payload    = exploit.framework.payloads.create(opts['Payload'])
 
 			# Set the force wait for session flag if the caller requested force
 			# blocking.  This is so that passive exploits can be blocked on from

lib/msf/core/db_manager.rb

@@ -375,7 +375,6 @@ def update_all_module_details
 		refresh.each  {|md| md.destroy }
 		refresh = nil
 
-		stime = Time.now.to_f
 		[
 			[ 'exploit',   framework.exploits  ],
 			[ 'auxiliary', framework.auxiliary ],

lib/msf/core/module_manager.rb

@@ -12,11 +12,15 @@
 require 'msf/core/module_set'
 
 module Msf
-  # Upper management decided to throw in some middle management # because the modules were getting out of hand.  This
-  # bad boy takes care of the work of managing the interaction with modules in terms of loading and instantiation.
+  # Upper management decided to throw in some middle management
+  # because the modules were getting out of hand.  This bad boy takes
+  # care of the work of managing the interaction with modules in terms
+  # of loading and instantiation.
   #
   # @todo add unload support
-  class ModuleManager < ModuleSet
+  class ModuleManager
+    include Msf::Framework::Offspring
+
     require 'msf/core/payload_set'
 
     # require here so that Msf::ModuleManager is already defined
@@ -41,36 +45,28 @@ class ModuleManager < ModuleSet
     # Maps module type directory to its module type.
     TYPE_BY_DIRECTORY = Msf::Modules::Loader::Base::DIRECTORY_BY_TYPE.invert
 
-    # Overrides the module set method for adding a module so that some extra steps can be taken to subscribe the module
-    # and notify the event dispatcher.
-    #
-    # @param (see Msf::ModuleSet#add_module)
-    # @return (see Msf::ModuleSet#add_module)
-    def add_module(mod, name, file_paths)
-      # Call {Msf::ModuleSet#add_module} with same arguments
-      dup = super
-
-      # Automatically subscribe a wrapper around this module to the necessary
-      # event providers based on whatever events it wishes to receive.  We
-      # only do this if we are the module manager instance, as individual
-      # module sets need not subscribe.
-      auto_subscribe_module(dup)
-
-      # Notify the framework that a module was loaded
-      framework.events.on_module_load(name, dup)
-
-      dup
+    def [](key)
+      names = key.split("/")
+      type = names.shift
+
+      module_set = module_set_by_type[type]
+
+      module_reference_name = names.join("/")
+      module_set[module_reference_name]
     end
 
     # Creates a module instance using the supplied reference name.
     #
-    # @param [String] name a module reference name.  It may optionally be prefixed with a "<type>/", in which case the
-    #   module will be created from the {Msf::ModuleSet} for the given <type>.
+    # @param name [String] A module reference name.  It may optionally
+    #   be prefixed with a "<type>/", in which case the module will be
+    #   created from the {Msf::ModuleSet} for the given <type>.
+    #   Otherwise, we step through all sets until we find one that
+    #   matches.
     # @return (see Msf::ModuleSet#create)
     def create(name)
       # Check to see if it has a module type prefix.  If it does,
       # try to load it from the specific module set for that type.
-      names = name.split(File::SEPARATOR)
+      names = name.split("/")
       potential_type_or_directory = names.first
 
       # if first name is a type
@@ -81,15 +77,24 @@ def create(name)
         type = TYPE_BY_DIRECTORY[potential_type_or_directory]
       end
 
+      module_instance = nil
       if type
         module_set = module_set_by_type[type]
 
-        module_reference_name = names[1 .. -1].join(File::SEPARATOR)
-        module_set.create(module_reference_name)
-      # Otherwise, just try to load it by name.
+        # First element in names is the type, so skip it
+        module_reference_name = names[1 .. -1].join("/")
+        module_instance = module_set.create(module_reference_name)
       else
-        super
+        # Then we don't have a type, so we have to step through each set
+        # to see if we can create this module.
+        module_set_by_type.each do |_, set|
+          module_reference_name = names.join("/")
+          module_instance = set.create(module_reference_name)
+          break if module_instance
+        end
       end
+
+      module_instance
     end
 
 
@@ -128,18 +133,18 @@ def initialize(framework, types=Msf::MODULE_TYPES)
       types.each { |type|
         init_module_set(type)
       }
-
-      super(nil)
     end
 
     protected
 
-    # This method automatically subscribes a module to whatever event providers it wishes to monitor.  This can be used
-    # to allow modules to automatically # execute or perform other tasks when certain events occur.  For instance, when
-    # a new host is detected, other aux modules may wish to run such that they can collect more information about the
-    # host that was detected.
+    # This method automatically subscribes a module to whatever event
+    # providers it wishes to monitor.  This can be used to allow modules
+    # to automatically execute or perform other tasks when certain
+    # events occur.  For instance, when a new host is detected, other
+    # aux modules may wish to run such that they can collect more
+    # information about the host that was detected.
     #
-    # @param [Class] mod a Msf::Module subclass
+    # @param mod [Class] A subclass of Msf::Module
     # @return [void]
     def auto_subscribe_module(mod)
       # If auto-subscribe has been disabled

lib/msf/core/module_manager/cache.rb

@@ -24,8 +24,8 @@ def cache_empty?
   def load_cached_module(type, reference_name)
     loaded = false
 
-    module_info = self.module_info_by_path.values.find { |module_info|
-      module_info[:type] == type and module_info[:reference_name] == reference_name
+    module_info = self.module_info_by_path.values.find { |inner_info|
+      inner_info[:type] == type and inner_info[:reference_name] == reference_name
     }
 
     if module_info
@@ -116,8 +116,9 @@ def module_info_by_path_from_database!
 
         typed_module_set = module_set(type)
 
-        # Don't want to trigger as {Msf::ModuleSet#create} so check for key instead of using ||= which would call
-        # {Msf::ModuleSet#[]} which would potentially call {Msf::ModuleSet#create}.
+        # Don't want to trigger as {Msf::ModuleSet#create} so check for
+        # key instead of using ||= which would call {Msf::ModuleSet#[]}
+        # which would potentially call {Msf::ModuleSet#create}.
         unless typed_module_set.has_key? reference_name
           typed_module_set[reference_name] = Msf::SymbolicModule
         end
@@ -126,4 +127,4 @@ def module_info_by_path_from_database!
 
     self.module_info_by_path
   end
-end
+end

lib/msf/core/module_manager/loading.rb

@@ -57,21 +57,16 @@ def file_changed?(path)
   # categorized accordingly.
   #
   def on_module_load(mod, type, name, modinfo)
-    # Payload modules require custom loading as the individual files
-    # may not directly contain a logical payload that a user would
-    # reference, such as would be the case with a payload stager or
-    # stage.  As such, when payload modules are loaded they are handed
-    # off to a special payload set.  The payload set, in turn, will
-    # automatically create all the permutations after all the payload
-    # modules have been loaded.
-
-    if (type != Msf::MODULE_PAYLOAD)
-      # Add the module class to the list of modules and add it to the
-      # type separated set of module classes
-      add_module(mod, name, modinfo)
-    end
+    dup = module_set_by_type[type].add_module(mod, name, modinfo)
+
+    # Automatically subscribe a wrapper around this module to the necessary
+    # event providers based on whatever events it wishes to receive.
+    auto_subscribe_module(dup)
+
+    # Notify the framework that a module was loaded
+    framework.events.on_module_load(name, dup)
 
-    module_set_by_type[type].add_module(mod, name, modinfo)
+    dup
   end
 
   protected

lib/msf/core/module_manager/module_sets.rb

@@ -39,7 +39,7 @@ def init_module_set(type)
     self.enablement_by_type[type] = true
     case type
       when Msf::MODULE_PAYLOAD
-        instance = Msf::PayloadSet.new(self)
+        instance = Msf::PayloadSet.new
       else
         instance = Msf::ModuleSet.new(type)
     end
@@ -100,4 +100,4 @@ def type_enabled?(type)
 
   attr_accessor :enablement_by_type # :nodoc:
   attr_accessor :module_set_by_type # :nodoc:
-end
+end

lib/msf/core/module_set.rb

@@ -33,24 +33,17 @@ def [](name)
 
   # Create an instance of the supplied module by its name
   #
-  # @param [String] name the module reference name.
-  # @return [Msf::Module] instance of the named module.
+  # @param name [String] The module reference name.
+  # @return [Msf::Module,nil] Instance of the named module or nil if it
+  #   could not be created.
   def create(name)
     klass = fetch(name, nil)
     instance = nil
 
     # If there is no module associated with this class, then try to demand
     # load it.
     if klass.nil? or klass == Msf::SymbolicModule
-      # If we are the root module set, then we need to try each module
-      # type's demand loading until we find one that works for us.
-      if module_type.nil?
-        Msf::MODULE_TYPES.each { |type|
-          framework.modules.load_cached_module(type, name)
-        }
-      else
-        framework.modules.load_cached_module(module_type, name)
-      end
+      framework.modules.load_cached_module(module_type, name)
 
       recalculate
 
@@ -168,17 +161,6 @@ def initialize(type = nil)
   def on_module_reload(mod)
   end
 
-  # @!attribute [rw] postpone_recalc
-  #   Whether or not recalculations should be postponed.  This is used
-  #   from the context of the {#each_module_list} handler in order to
-  #   prevent the demand loader from calling recalc for each module if
-  #   it's possible that more than one module may be loaded.  This field
-  #   is not initialized until used.
-  #
-  #   @return [true] if {#recalculate} should not be called immediately
-  #   @return [false] if {#recalculate} should be called immediately
-  attr_accessor :postpone_recalculate
-
   # Dummy placeholder to recalculate aliases and other fun things.
   #
   # @return [void]
@@ -194,8 +176,6 @@ def valid?(name)
     (self[name]) ? true : false
   end
 
-  protected
-
   # Adds a module with a the supplied name.
   #
   # @param [Class] mod The module class: a subclass of Msf::Module.
@@ -226,25 +206,24 @@ def add_module(mod, name, modinfo = nil)
     mod
   end
 
+  protected
+
   # Load all modules that are marked as being symbolic.
   #
   # @return [void]
   def demand_load_modules
+    found_symbolics = false
     # Pre-scan the module list for any symbolic modules
     self.each_pair { |name, mod|
       if (mod == Msf::SymbolicModule)
-        self.postpone_recalculate = true
-
+        found_symbolics = true
         mod = create(name)
-
         next if (mod.nil?)
       end
     }
 
     # If we found any symbolic modules, then recalculate.
-    if (self.postpone_recalculate)
-      self.postpone_recalculate = false
-
+    if (found_symbolics)
       recalculate
     end
   end