Merge pull request #1 from msutovsky-r7/collab/fix/payload-reload

Collab/fix/payload reload

Author: xaitax

lib/msf/core/module_manager/reloading.rb

@@ -1,6 +1,6 @@
 # -*- coding: binary -*-
-
 # Concerns reloading modules
+
 module Msf::ModuleManager::Reloading
   # Reloads the module specified in mod.  This can either be an instance of a module or a module class.
   #
@@ -10,110 +10,103 @@ def reload_module(mod)
     # if it's an instance, then get its class
     if mod.is_a? Msf::Module
       metasploit_class = mod.class
-      original_instance = mod
     else
       metasploit_class = mod
-      original_instance = nil
     end
 
-    # Handle aliases cleanup
-    if (aliased_as = inv_aliases[metasploit_class.fullname])
+    if aliased_as = self.inv_aliases[metasploit_class.fullname]
       aliased_as.each do |a|
-        aliases.delete a
+        self.aliases.delete a
       end
-      inv_aliases.delete metasploit_class.fullname
+      self.inv_aliases.delete metasploit_class.fullname
     end
 
-    # Special handling for payload modules
-    if mod&.payload?
-      return reload_payload_module(metasploit_class, original_instance)
+    if mod.payload?
+      return reload_payload_module(mod)
     end
-
-    # Standard module reloading for non-payloads
+    
     namespace_module = metasploit_class.module_parent
-
+    
     # Check if the namespace module has a loader
     unless namespace_module.respond_to?(:loader)
-      raise "Module #{metasploit_class.fullname} namespace does not have a loader"
+      elog("Module does not have loader")
+      return mod
     end
 
     loader = namespace_module.loader
     loader.reload_module(mod)
   end
-
-  private
-
-  # Reloads a payload module. This must be done by reloading the entire parent
-  # directory to ensure the framework's complex payload "stitching" process
-  # (combining stages, stagers, and mixins) is correctly executed. This is slower
-  # but guarantees a fully-functional reloaded module.
-  def reload_payload_module(metasploit_class, original_instance = nil)
-    # Step 1: Get all necessary identifiers from the original module class.
-    fullname = metasploit_class.fullname
-    refname = metasploit_class.refname
-    type = metasploit_class.type
-    file_path = metasploit_class.file_path
-
-    # Get a reference to the original datastore now, but do not copy it yet.
-    original_datastore = original_instance&.datastore
-
-    # Step 2: Manually purge the old module from the framework's caches.
-    if (module_set = module_set_by_type.fetch(type, nil))
-      module_set.delete(refname)
-    end
-    if (aliases_for_fullname = inv_aliases[fullname])
-      aliases_for_fullname.each { |a| aliases.delete(a) }
-      inv_aliases.delete(fullname)
+  # Reload payload module, separately from other categories. This is due to complexity of payload module and due to the fact they don't follow class structure as rest of the modules. 
+  # @param [Msf::Module, Class] mod either an instance of a module or a module class
+  # @return (see Msf::Modules::Loader::Base#reload_module)
+  def reload_payload_module(mod)
+    if mod.is_a? Msf::Module
+      metasploit_class = mod.class
+      original_instance = mod
+    else
+      metasploit_class = mod
+      original_instance = nil
+    end 
+    if (module_set = self.module_set_by_type.fetch(metasploit_class.type, nil))
+      module_set.delete(metasploit_class.refname)
     end
 
-    # Step 3: Get the module's parent directory path.
-    module_info = module_info_by_path[file_path]
+    module_info = self.module_info_by_path[metasploit_class.file_path]
     unless module_info && (parent_path = module_info[:parent_path])
-      raise Msf::LoadError, "Could not find cached module information for path: #{file_path}"
+      elog("Failed to get parent_path from module object")
+      return mod
     end
 
-    # Step 4: Use the core framework loader to reload the entire parent directory.
-    # This is the only way to reliably trigger the payload stitching logic.
-    load_modules(parent_path, force: true)
+    case original_instance&.payload_type
+      when Msf::Payload::Type::Single
+        prepend_path = 'singles'
+      when Msf::Payload::Type::Stager
+        prepend_path = 'stagers'
+      when Msf::Payload::Type::Stage
+        prepend_path = 'stages'
+      when Msf::Payload::Type::Adapter
+        prepend_path = 'adapters'
+    end
+
+    full_reference_name = File.join(prepend_path, module_info[:reference_name])
+    self.loaders.each { |loader| loader.load_module(parent_path,module_info[:type], full_reference_name, {:force => true}) }
+
 
-    # Step 5: Now that the framework has completed its full reload process,
-    # use the public API to get a new instance of our reloaded module.
-    new_instance = framework.modules.create(fullname)
+    # Get reloaded module
+    new_instance = framework.modules.create(metasploit_class.fullname)
 
     if new_instance.blank?
-      raise Msf::LoadError, "Failed to create a new instance of #{fullname} after reloading. The module file may be broken."
+      elog("Failed create new instance")
+      return mod
     end
 
-    # Step 6: Restore the datastore to the new, fully-functional instance.
-    # Now we perform the copy, which is a method on the new datastore's merge function.
-    if original_datastore
-      new_instance.datastore.merge!(original_datastore)
-    end
+    # Restore the datastore 
+    new_instance.datastore.merge!(original_instance.datastore)
+
 
     # Return the new instance, which the framework will make the active module.
     return new_instance
   rescue StandardError => e
-    elog("Failed to reload payload #{fullname}", error: e)
-    raise Msf::LoadError, e.message
+    elog("Failed to reload payload #{fullname}: #{e.message}")
+    return mod
   end
 
-  public
-
   # Reloads modules from all module paths
   #
   # @return (see Msf::ModuleManager::Loading#load_modules)
   def reload_modules
-    enablement_by_type.each_key do |type|
+    self.enablement_by_type.each_key do |type|
       module_set_by_type[type].clear
       init_module_set(type)
     end
-
-    aliases.clear
-    inv_aliases.clear
+    self.aliases.clear
+    self.inv_aliases.clear
 
     # default the count to zero the first time a type is accessed
     count_by_type = Hash.new(0)
 
+    framework.init_module_paths unless framework.module_paths_inited
+
     module_paths.each do |path|
       path_count_by_type = load_modules(path, force: true)
 
@@ -124,6 +117,7 @@ def reload_modules
     end
 
     refresh_cache_from_module_files
+
     count_by_type
   end
 end