Add ability to use dots in module and payload file names

Author: rastating

lib/cli/auto_complete.rb

@@ -28,7 +28,7 @@ def refresh_autocomplete_options
 
         if mod.exploit_module?
           opts_hash['payload'] = {}
-          Wpxf::Payloads.payload_list.each { |p| opts_hash['payload'][p] = {} }
+          Wpxf::Payloads.payload_list.each { |p| opts_hash['payload'][p[:name]] = {} }
         end
       end
 
@@ -39,8 +39,8 @@ def refresh_autocomplete_options
     def build_cmd_list
       cmds = {}
       permitted_commands.each { |c| cmds[c] = {} }
-      Wpxf::Auxiliary.module_list.each { |m| cmds['use'][m] = {} }
-      Wpxf::Exploit.module_list.each { |m| cmds['use'][m] = {} }
+      Wpxf::Auxiliary.module_list.each { |m| cmds['use'][m[:name]] = {} }
+      Wpxf::Exploit.module_list.each { |m| cmds['use'][m[:name]] = {} }
       cmds['show'] = {
         'options' => {},
         'advanced' => {},

lib/cli/context.rb

@@ -1,9 +1,6 @@
 module Cli
   # A context which modules will be used in.
   class Context
-    def initialize
-    end
-
     def class_name(path_name)
       return path_name if path_name !~ /_/ && path_name =~ /[A-Z]+.*/
       path_name.split('_').map(&:capitalize).join
@@ -14,22 +11,7 @@ def verbose?
     end
 
     def load_module(path)
-      match = path.match(/(auxiliary|exploit)\/(.+)/i)
-      raise 'Invalid module path' unless match
-
-      type = match.captures[0]
-      name = class_name(match.captures[1])
-
-      begin
-        if type.eql? 'auxiliary'
-          @module = Wpxf::Auxiliary.const_get(name).new
-        elsif type.eql? 'exploit'
-          @module = Wpxf::Exploit.const_get(name).new
-        end
-      rescue NameError
-        raise 'Invalid module name'
-      end
-
+      @module = Wpxf.load_module(path)
       @module_path = path
       @module
     end
@@ -45,19 +27,7 @@ def reload
     end
 
     def load_payload(name)
-      clsid = class_name(name)
-
-      if Wpxf::Payloads.const_defined?(clsid)
-        payload_class = Wpxf::Payloads.const_get(clsid)
-        if payload_class.is_a?(Class)
-          self.module.payload = payload_class.new
-        else
-          fail "\"#{name}\" is not a valid payload"
-        end
-      else
-        fail "\"#{name}\" is not a valid payload"
-      end
-
+      self.module.payload = Wpxf::Payloads.load_payload(name)
       self.module.payload.check(self.module)
       self.module.payload
     end

modules/modules.rb

@@ -1,29 +1,40 @@
 module Wpxf
-  def self.underscore(module_name)
-    module_name.gsub(/::/, '/').
-                gsub(/([A-Z]+)([A-Z][a-z])/,'\1_\2').
-                gsub(/([a-z\d])([A-Z])/,'\1_\2').
-                tr("-", "_").
-                downcase
+  def self.build_module_list(namespace, id_prefix = '')
+    modules = namespace.constants.select do |c|
+      namespace.const_get(c).is_a? Class
+    end
+
+    modules.map do |m|
+      klass = namespace.const_get(m)
+      filename = klass.new.method(:initialize).source_location[0]
+      {
+        class: klass,
+        name: "#{id_prefix}#{File.basename(filename, '.rb')}"
+      }
+    end
+  end
+
+  def self.load_module(name)
+    match = name.match(/(auxiliary|exploit)\/(.+)/i)
+    raise 'Invalid module path' unless match
+
+    type = match.captures[0]
+    list = type == 'auxiliary' ? Wpxf::Auxiliary.module_list : Wpxf::Exploit.module_list
+
+    mod = list.find { |p| p[:name] == name }
+    raise "\"#{name}\" is not a valid module" if mod.nil?
+    mod[:class].new
   end
 
   module Auxiliary
     def self.module_list
-      modules = Wpxf::Auxiliary.constants.select do |c|
-        Wpxf::Auxiliary.const_get(c).is_a? Class
-      end
-
-      modules.map { |m| "auxiliary/#{Wpxf.underscore(m.to_s)}" }
+      @@modules ||= Wpxf.build_module_list(Wpxf::Auxiliary, 'auxiliary/')
     end
   end
 
   module Exploit
     def self.module_list
-      modules = Wpxf::Exploit.constants.select do |c|
-        Wpxf::Exploit.const_get(c).is_a? Class
-      end
-
-      modules.map { |m| "exploit/#{Wpxf.underscore(m.to_s)}" }
+      @@modules ||= Wpxf.build_module_list(Wpxf::Exploit, 'exploit/')
     end
   end
 
@@ -37,11 +48,13 @@ def self.payload_count
     end
 
     def self.payload_list
-      payloads = Wpxf::Payloads.constants.select do |c|
-        Wpxf::Payloads.const_get(c).is_a? Class
-      end
+      @@payloads ||= Wpxf.build_module_list(Wpxf::Payloads)
+    end
 
-      payloads.map { |p| Wpxf.underscore(p.to_s) }
+    def self.load_payload(name)
+      payload = payload_list.find { |p| p[:name] == name }
+      raise "\"#{name}\" is not a valid payload" if payload.nil?
+      payload[:class].new
     end
   end
 end

spec/modules_spec.rb

@@ -0,0 +1,45 @@
+require_relative 'spec_helper'
+require 'modules'
+
+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, 'exploit/')
+      expect(result[0]).to include(:class, :name)
+
+      mod = result.find { |m| m[:name] == 'exploit/admin_shell_upload' }
+      expect(mod).to_not be_nil
+      expect(mod[:class]).to be Wpxf::Exploit::AdminShellUpload
+    end
+  end
+end
+
+describe 'Wpxf::Auxiliary' do
+  describe '.module_list' do
+    it 'builds an array of hashes containing the auxiliary modules' do
+      list = Wpxf::Auxiliary.module_list
+      expect(list).to_not be_nil
+      expect(list[0]).to include(:class, :name)
+    end
+  end
+end
+
+describe 'Wpxf::Exploit' do
+  describe '.module_list' do
+    it 'builds an array of hashes containing the exploit modules' do
+      list = Wpxf::Exploit.module_list
+      expect(list).to_not be_nil
+      expect(list[0]).to include(:class, :name)
+    end
+  end
+end
+
+describe 'Wpxf::Payloads' do
+  describe '.payload_list' do
+    it 'builds an array of hashes containing the module payloads' do
+      list = Wpxf::Payloads.payload_list
+      expect(list).to_not be_nil
+      expect(list[0]).to include(:class, :name)
+    end
+  end
+end