Adds reloading for stagers/stages/adapters

Author: msutovsky-r7

lib/msf/core/module_manager/reloading.rb

@@ -1,4 +1,5 @@
 # -*- coding: binary -*-
+
 # Concerns reloading modules
 
 module Msf::ModuleManager::Reloading
@@ -14,29 +15,41 @@ def reload_module(mod)
       metasploit_class = mod
     end
 
-    if aliased_as = self.inv_aliases[metasploit_class.fullname]
+    if aliased_as = inv_aliases[metasploit_class.fullname]
       aliased_as.each do |a|
-        self.aliases.delete a
+        aliases.delete a
       end
-      self.inv_aliases.delete metasploit_class.fullname
+      inv_aliases.delete metasploit_class.fullname
     end
 
     if mod.payload?
       return reload_payload_module(mod)
     end
-    
+
+    if aliased_as = inv_aliases[metasploit_class.fullname]
+      aliased_as.each do |a|
+        aliases.delete a
+      end
+      inv_aliases.delete metasploit_class.fullname
+    end
+
     namespace_module = metasploit_class.module_parent
-    
+
     # Check if the namespace module has a loader
     unless namespace_module.respond_to?(:loader)
-      elog("Module does not have loader")
+      elog('Module does not have loader')
       return mod
     end
 
     loader = namespace_module.loader
     loader.reload_module(mod)
   end
-  # 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. 
+
+  def manual_reload(parent_path, type, ref_name)
+    loaders.each { |loader| loader.load_module(parent_path, type, ref_name, { force: true }) }
+  end
+
+  # 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)
@@ -46,44 +59,39 @@ def reload_payload_module(mod)
     else
       metasploit_class = mod
       original_instance = nil
-    end 
-    if (module_set = self.module_set_by_type.fetch(metasploit_class.type, nil))
+    end
+    if (module_set = module_set_by_type.fetch(metasploit_class.type, nil))
       module_set.delete(metasploit_class.refname)
     end
-
-    module_info = self.module_info_by_path[metasploit_class.file_path]
+    module_info = module_info_by_path[metasploit_class.file_path]
     unless module_info && (parent_path = module_info[:parent_path])
-      elog("Failed to get parent_path from module object")
+      elog('Failed to get parent_path from module object')
       return mod
     end
 
-    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
+    # reload adapters if any
+    manual_reload(parent_path, module_info[:type], File.join('adapters', mod.adapter_refname)) if mod.adapter_refname
+
+    # reload stagers if any
+    manual_reload(parent_path, module_info[:type], File.join('stagers', mod.stager_refname)) if mod.stager_refname
 
-    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}) }
+    # reload stages if any
+    manual_reload(parent_path, module_info[:type], File.join('stages', mod.stage_refname)) if mod.stage_refname
 
+    # reload single if any
+    manual_reload(parent_path, module_info[:type], File.join('singles', module_info[:reference_name])) if original_instance.payload_type == Msf::Payload::Type::Single
 
     # Get reloaded module
     new_instance = framework.modules.create(metasploit_class.fullname)
 
     if new_instance.blank?
-      elog("Failed create new instance")
+      elog('Failed create new instance')
       return mod
     end
 
-    # Restore the datastore 
+    # 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
@@ -95,12 +103,12 @@ def reload_payload_module(mod)
   #
   # @return (see Msf::ModuleManager::Loading#load_modules)
   def reload_modules
-    self.enablement_by_type.each_key do |type|
+    enablement_by_type.each_key do |type|
       module_set_by_type[type].clear
       init_module_set(type)
     end
-    self.aliases.clear
-    self.inv_aliases.clear
+    aliases.clear
+    inv_aliases.clear
 
     # default the count to zero the first time a type is accessed
     count_by_type = Hash.new(0)