moves logic for bad paths to common.rb

lib/msf/ui/console/command_dispatcher/common.rb

@@ -158,6 +158,44 @@ def index_from_list(list, index, &block)
     yield list[idx]
   end
 
+  # Trims starting `.`, `./` `/`, `+path_head+/`, & `/+path_head+/` from +path+. Also trims trailing `.+extension+`
+  # from +path+, and any possible combination of misspellings of +extension+.
+  #
+  # @param path [String] The path to be trimmed
+  # @param path_head [String] The top-level directory that should be removed from the path
+  # @param extensions [Array] File extensions to be trimmed from +path+. `.` is automatically included. Defaults to ['rb', 'py', 'go'].
+  # @return [String] Altered +path+. Will return unaltered +path+ if regex constructed with +path_head+ & +path+ is not detected
+  def trim_path(path, path_head, extensions: ['rb', 'py', 'go'])
+    #Builds capture groups for all supported file extensions
+    regex_extension = ''
+    extensions.each do |ext|
+      regex_extension << "([#{ext}])+|"
+    end
+    regex_extension.delete_suffix!('|')
+
+    regexp = %r{
+        (
+          ^\.?                      # Dot at beginning of path
+          /?                        # Slash at beginning of path
+          (#{path_head}/)?          # top level directory (slash prepending directory name is optional)
+        )
+
+        |                           # OR
+
+        (
+          \.(#{regex_extension})$   # any possible file extension at end of path
+        )
+
+        |                           # OR
+
+        (
+          \.$                       # trailing dot
+        )
+    }ix
+
+    path.gsub(regexp, '')
+  end
+
 end
 
 end

lib/msf/ui/console/command_dispatcher/core.rb

@@ -1579,18 +1579,15 @@ def cmd_set(*args)
 
     # Set PAYLOAD
     if name.upcase == 'PAYLOAD' && active_module && (active_module.exploit? || active_module.evasion?)
-      if value.start_with?('/', 'payload/')
-        # Trims starting `/`, `payload/`, `/payload/` from user input
-        value.sub!(%r{^/?(?:payload/)?}, '')
-      else
-        # Checking set PAYLOAD by index
-        index_from_list(payload_show_results, value) do |mod|
-          return false unless mod && mod.respond_to?(:first)
+      value = trim_path(value, 'payload')
 
-          # [name, class] from payload_show_results
-          value = mod.first
-        end
+      index_from_list(payload_show_results, value) do |mod|
+        return false unless mod && mod.respond_to?(:first)
+
+        # [name, class] from payload_show_results
+        value = mod.first
       end
+
     end
 
     # If the driver indicates that the value is not valid, bust out.

lib/msf/ui/console/command_dispatcher/modules.rb

@@ -675,12 +675,7 @@ def cmd_use(*args)
             mod_resolved = args[1] == true ? true : false
 
             # Ensure we have a reference name and not a path
-            if mod_name.start_with?('./', 'modules/')
-              mod_name.sub!(%r{^(?:\./)?modules/}, '')
-            end
-            if mod_name.end_with?('.rb')
-              mod_name.sub!(/\.rb$/, '')
-            end
+            mod_name = trim_path(mod_name, "modules")
 
             begin
               mod = framework.modules.create(mod_name)

spec/lib/msf/ui/console/command_dispatcher/common_spec.rb

@@ -0,0 +1,71 @@
+require 'spec_helper'
+require 'msf/ui/console/command_dispatcher/common'
+
+RSpec.describe ::Msf::Ui::Console::CommandDispatcher::Common do
+  let(:subject) do
+    dummy_class = Class.new
+    dummy_class.extend(described_class)
+  end
+
+  def with_path_extensions(original_paths)
+    extensions = ['rb', 'br', 'r', 'b']
+    extensions += ['py', 'yp', 'p', 'y']
+    extensions += ['go', 'og', 'g', 'o']
+    extensions += ['']
+
+    paths_with_extensions = extensions.flat_map do |extension|
+      original_paths.map { |path| "#{path}.#{extension}" }
+    end
+    original_paths + paths_with_extensions
+  end
+
+  describe 'Trimming a path of errors' do
+    context 'when a user is inputting a Module path' do
+      let(:valid_path) { 'windows/smb/ms08_067_netapi' }
+      let(:all_paths) do
+        [
+          './exploits/windows/smb/ms08_067_netapi',
+          '/exploits/windows/smb/ms08_067_netapi',
+          '.exploits/windows/smb/ms08_067_netapi',
+          'exploits/windows/smb/ms08_067_netapi',
+
+          './windows/smb/ms08_067_netapi',
+          '/windows/smb/ms08_067_netapi',
+          '.windows/smb/ms08_067_netapi',
+
+          valid_path
+        ]
+      end
+
+      it 'corrects the module paths to be valid' do
+        with_path_extensions(all_paths).each do |path|
+          expect(subject.trim_path(path, 'exploits')).to eql(valid_path)
+        end
+      end
+    end
+
+    context 'when a user is inputting a Payload path' do
+      let(:valid_path) { 'windows/x64/vncinject/reverse_winhttps' }
+      let(:all_paths) do
+        [
+          './payload/windows/x64/vncinject/reverse_winhttps',
+          '/payload/windows/x64/vncinject/reverse_winhttps',
+          '.payload/windows/x64/vncinject/reverse_winhttps',
+          'payload/windows/x64/vncinject/reverse_winhttps',
+
+          './windows/x64/vncinject/reverse_winhttps',
+          '/windows/x64/vncinject/reverse_winhttps',
+          '.windows/x64/vncinject/reverse_winhttps',
+
+          valid_path
+        ]
+      end
+
+      it "when a user is inputting a Modules path" do
+        with_path_extensions(all_paths).each do |path|
+          expect(subject.trim_path(path, 'payload')).to eql(valid_path)
+        end
+      end
+    end
+  end
+end