Add support for custom modules in home directory

Author: rastating

lib/wpxf/cli/module_cache.rb

@@ -45,6 +45,7 @@ def rebuild_cache
         print_warning 'Refreshing the module cache...'
 
         Wpxf::Models::Module.truncate
+        Wpxf.load_custom_modules
 
         create_module_models 'exploit'
         create_module_models 'auxiliary'

lib/wpxf/modules.rb

@@ -1,19 +1,36 @@
 # frozen_string_literal: true
 
 module Wpxf
-  def self.build_module_list(namespace, folder_name = '')
+  def self.modules_path
+    File.join(Wpxf.app_path, 'lib', 'wpxf', 'modules')
+  end
+
+  def self.custom_modules_path
+    File.join(Wpxf.home_directory, 'modules')
+  end
+
+  def self.payloads_path
+    File.join(Wpxf.app_path, 'lib', 'wpxf', 'payloads')
+  end
+
+  def self.build_module_list(namespace, source_folders = [])
     modules = namespace.constants.select do |c|
       namespace.const_get(c).is_a? Class
     end
 
-    modules_directory = File.join(Wpxf.app_path, 'lib', 'wpxf', folder_name)
-
     modules.map do |m|
       klass = namespace.const_get(m)
       filename = klass.new.method(:initialize).source_location[0]
+
+      # Remove any source folders from the path and store the
+      # relative path that will be used in the CLI.
+      source_folders.each do |source_folder|
+        filename = filename.sub(source_folder, '')
+      end
+
       {
         class: klass,
-        name: filename.sub(modules_directory, '').sub(/^\//, '').sub(/\.rb$/, '')
+        name: filename.sub(/^\//, '').sub(/\.rb$/, '')
       }
     end
   end
@@ -24,15 +41,20 @@ def self.load_module(path)
     Object.const_get(mod.class_name).new
   end
 
+  def self.load_custom_modules
+    custom_modules_path = File.join(Wpxf.home_directory, 'modules', '**', '*.rb')
+    Dir.glob(custom_modules_path).each { |p| load p }
+  end
+
   module Auxiliary
     def self.module_list
-      Wpxf.build_module_list(Wpxf::Auxiliary, 'modules')
+      Wpxf.build_module_list(Wpxf::Auxiliary, [Wpxf.modules_path, Wpxf.custom_modules_path])
     end
   end
 
   module Exploit
     def self.module_list
-      Wpxf.build_module_list(Wpxf::Exploit, 'modules')
+      Wpxf.build_module_list(Wpxf::Exploit, [Wpxf.modules_path, Wpxf.custom_modules_path])
     end
   end
 
@@ -46,7 +68,7 @@ def self.payload_count
     end
 
     def self.payload_list
-      @@payloads ||= Wpxf.build_module_list(Wpxf::Payloads, 'payloads')
+      @@payloads ||= Wpxf.build_module_list(Wpxf::Payloads, [Wpxf.payloads_path])
     end
 
     def self.load_payload(name)
@@ -60,3 +82,5 @@ def self.load_payload(name)
 require_rel 'modules/auxiliary'
 require_rel 'modules/exploit'
 require_rel 'payloads'
+
+Wpxf.load_custom_modules

spec/lib/wpxf/cli/module_cache_spec.rb

@@ -140,6 +140,12 @@ def initialize
       expect(Wpxf::Models::Module.first(name: 'test')).to be_nil
     end
 
+    it 'should reload the custom modules' do
+      allow(Wpxf).to receive(:load_custom_modules).and_call_original
+      subject.rebuild_cache
+      expect(Wpxf).to have_received(:load_custom_modules).exactly(1).times
+    end
+
     it 'should create a {Models::Module} for each class in the Wpxf::Exploit namespace' do
       modules = Wpxf::Exploit.constants.select do |c|
         Wpxf::Exploit.const_get(c).is_a? Class

spec/lib/wpxf/modules_spec.rb

@@ -6,14 +6,53 @@
 describe Wpxf do
   describe '.build_module_list' do
     it 'builds an array of hashes containing the modules for the specified namespace' do
-      result = Wpxf.build_module_list(Wpxf::Exploit, 'modules')
+      result = Wpxf.build_module_list(Wpxf::Exploit, [Wpxf.modules_path])
       expect(result[0]).to include(:class, :name)
 
       mod = result.find { |m| m[:name] == 'exploit/shell/admin_shell_upload' }
       expect(mod).to_not be_nil
       expect(mod[:class]).to be Wpxf::Exploit::AdminShellUpload
     end
   end
+
+  describe '.modules_path' do
+    it 'should return the path to the standard modules directory' do
+      path = Wpxf.modules_path
+      expect(path).to eql(File.join(Wpxf.app_path, 'lib', 'wpxf', 'modules'))
+    end
+  end
+
+  describe '.payloads_path' do
+    it 'should return the path to the payloads directory' do
+      path = Wpxf.payloads_path
+      expect(path).to eql(File.join(Wpxf.app_path, 'lib', 'wpxf', 'payloads'))
+    end
+  end
+
+  describe '.custom_modules_path' do
+    it 'should return the path to the custom modules directory' do
+      path = Wpxf.custom_modules_path
+      expect(path).to eql(File.join(Wpxf.home_directory, 'modules'))
+    end
+  end
+
+  describe '.load_custom_modules' do
+    it 'should load all modules found in the home directory of the current user' do
+      allow(Wpxf).to receive(:load)
+      allow(Dir).to receive(:glob)
+        .with(File.join(Wpxf.home_directory, 'modules', '**', '*.rb'))
+        .and_return %w[/path1.rb /path2.rb]
+
+      Wpxf.load_custom_modules
+      expect(Wpxf).to have_received(:load)
+        .with('/path1.rb')
+        .exactly(1).times
+
+      expect(Wpxf).to have_received(:load)
+        .with('/path2.rb')
+        .exactly(1).times
+    end
+  end
 end
 
 describe 'Wpxf::Auxiliary' do