Extend replicants via arbitrary Ruby code

MSP-11673

* Implements a #register_extensions method on Msf::Module
* Any registered Ruby modules will extend the cloned module returned by #replicant

lib/msf/core/module.rb

@@ -53,6 +53,10 @@ class Module
   include Msf::Module::UI
   include Msf::Module::UUID
 
+  # The key where a comma-separated list of Ruby module names will live in the
+  # datastore, consumed by #replicant to allow clean override of MSF module methods.
+  REPLICANT_EXTENSION_DS_KEY = 'ReplicantExtensions'
+
   # Make include public so we can runtime extend
   public_class_method :include
 
@@ -142,7 +146,6 @@ class Module
   # Creates a fresh copy of an instantiated module
   #
   def replicant
-
     obj = self.class.new
     self.instance_variables.each { |k|
       v = instance_variable_get(k)
@@ -154,9 +157,40 @@ class Module
     obj.user_input   = self.user_input
     obj.user_output  = self.user_output
     obj.module_store = self.module_store.clone
+
+    obj.perform_extensions
     obj
   end
 
+  # Extends self with the constant list in the datastore
+  # @return [void]
+  def perform_extensions
+    if datastore[REPLICANT_EXTENSION_DS_KEY].present?
+      if datastore[REPLICANT_EXTENSION_DS_KEY].respond_to?(:each)
+        datastore[REPLICANT_EXTENSION_DS_KEY].each do |const_name|
+          self.extend(const_name.constantize)
+        end
+      else
+        fail "Invalid settings in datastore at key #{REPLICANT_EXTENSION_DS_KEY}"
+      end
+    end
+  end
+
+  # @overload register_extensions(name)
+  #   @param[String] name of Ruby module
+  #
+  # @overload register_extensions(name_array)
+  #   @param[Array<String>] name_array array of Ruby module names
+  #
+  # @return [void]
+  def register_extensions(*rb_modules)
+    datastore[REPLICANT_EXTENSION_DS_KEY] = [] unless datastore[REPLICANT_EXTENSION_DS_KEY].present?
+    rb_modules.each do |rb_mod|
+      name = rb_mod.name
+      datastore[REPLICANT_EXTENSION_DS_KEY] << name unless datastore[REPLICANT_EXTENSION_DS_KEY].include? name
+    end
+  end
+
   #
   # Returns the unduplicated class associated with this module.
   #

spec/lib/msf/core/module_spec.rb

@@ -45,4 +45,50 @@ describe Msf::Module do
     it { is_expected.to respond_to :cached? }
     it { is_expected.to respond_to :is_usable }
   end
+
+  describe "cloning modules into replicants" do
+    module MsfExtensionTestFoo; def my_test1; true; end; end;
+    module MsfExtensionTestBar; def my_test2; true; end; end;
+
+    describe "#perform_extensions" do
+      describe "when there are extensions registered" do
+        before(:each) do
+          msf_module.register_extensions(MsfExtensionTestFoo, MsfExtensionTestBar)
+        end
+
+        it 'should extend the module replicant with the constants referenced in the datastore' do
+          expect(msf_module.replicant).to respond_to(:my_test1)
+          expect(msf_module.replicant).to respond_to(:my_test2)
+        end
+      end
+
+      describe "when the datastore key has invalid data" do
+        before(:each) do
+          msf_module.datastore[Msf::Module::REPLICANT_EXTENSION_DS_KEY] = "invalid"
+        end
+
+        it 'should raise an exception' do
+          expect{msf_module.replicant}.to raise_error(RuntimeError)
+        end
+      end
+    end
+
+    describe "#register_extensions" do
+      describe "with single module" do
+        it 'should place the named module in the datastore' do
+          msf_module.register_extensions(MsfExtensionTestFoo)
+          expect(msf_module.replicant.datastore[Msf::Module::REPLICANT_EXTENSION_DS_KEY]).to eql(["MsfExtensionTestFoo"])
+        end
+      end
+
+      describe "with multiple modules" do
+        it 'should place the named modules in the datastore' do
+          msf_module.register_extensions(MsfExtensionTestFoo, MsfExtensionTestBar)
+          expect(msf_module.replicant.datastore[Msf::Module::REPLICANT_EXTENSION_DS_KEY]).to eql(["MsfExtensionTestFoo", "MsfExtensionTestBar"])
+        end
+      end
+
+    end
+  end
+
 end