Add convenience methods for adding option groups

rapid7 · Feb 20, 2024 · be98b6a · be98b6a
1 parent 88b177c
commit be98b6a
Show file tree

Hide file tree

Showing 9 changed files with 64 additions and 34 deletions.
diff --git a/lib/msf/base/serializer/readable_text.rb b/lib/msf/base/serializer/readable_text.rb
@@ -571,7 +571,7 @@ def self.dump_generic_module(mod, indent = '')
   def self.dump_options(mod, indent = '', missing = false, advanced: false, evasion: false)
     filtered_options = mod.options.values.select { |opt| opt.advanced? == advanced && opt.evasion? == evasion }
 
-    option_groups = mod.options.option_groups.map { |_name, group| group }
+    option_groups = mod.options.groups.map { |_name, group| group }.sort_by(&:name)
     options_by_group = option_groups.map do |group|
       [group, group.option_names.map { |name| mod.options[name] }.compact]
     end.to_h

diff --git a/lib/msf/core/module/options.rb b/lib/msf/core/module/options.rb
@@ -67,4 +67,19 @@ def register_options(options, owner = self.class)
     self.options.add_options(options, owner)
     import_defaults(false)
   end
+
+  def register_option_group(name:, description:, option_names: [], merge: true)
+    existing_group = options.groups[name]
+    if merge && existing_group
+      existing_group.description = description
+      existing_group.add_options(option_names)
+    else
+      option_group = Msf::OptionGroup.new(name: name, description: description, option_names: option_names)
+      options.add_group(option_group)
+    end
+  end
+
+  def deregister_option_group(name:)
+    options.remove_group(name)
+  end
 end
diff --git a/lib/msf/core/option_container.rb b/lib/msf/core/option_container.rb
@@ -49,7 +49,7 @@ class OptionContainer < Hash
     #
     def initialize(opts = {})
       self.sorted = []
-      self.option_groups = {}
+      self.groups = {}
 
       add_options(opts)
     end
@@ -314,8 +314,18 @@ def merge_sort(other_container)
       result.sort
     end
 
-    def add_option_group(option_group)
-      option_groups[option_group.name] = option_group
+    # Adds an option group to the container
+    #
+    # @param option_group [Msf::OptionGroup]
+    def add_group(option_group)
+      groups[option_group.name] = option_group
+    end
+
+    # Removes an option group from the container by name
+    #
+    # @param group_name [String]
+    def remove_group(group_name)
+      groups.delete(group_name)
     end
 
     #
@@ -324,13 +334,13 @@ def add_option_group(option_group)
     attr_reader :sorted
 
     # @return [Hash<String, Msf::OptionGroup>]
-    attr_reader :option_groups
+    attr_reader :groups
 
     protected
 
     attr_writer :sorted # :nodoc:
 
-    attr_writer :option_groups
+    attr_writer :groups
   end
 
 end
diff --git a/lib/msf/core/option_group.rb b/lib/msf/core/option_group.rb
@@ -3,7 +3,12 @@
 module Msf
   class OptionGroup
 
-    attr_accessor :name, :description, :option_names
+    # @return [String] Name for the group
+    attr_accessor :name
+    # @return [String] Description to be displayed to the user
+    attr_accessor :description
+    # @return [Array<String>] List of datastore option names
+    attr_accessor :option_names
 
     # @param name [String] Name for the group
     # @param description [String] Description to be displayed to the user

diff --git a/lib/msf/core/optional_session/mssql.rb b/lib/msf/core/optional_session/mssql.rb
@@ -16,12 +16,12 @@ def initialize(info = {})
         )
 
         if framework.features.enabled?(Msf::FeatureManager::MSSQL_SESSION_TYPE)
-          options.add_option_group(Msf::OptionGroup.new(name: 'SESSION',
-                                                        description: 'Used when connecting via an existing SESSION',
-                                                        option_names: ['SESSION']))
-          options.add_option_group(Msf::OptionGroup.new(name: 'RHOST',
-                                                        description: 'Used when making a new connection via RHOSTS',
-                                                        option_names: RHOST_GROUP_OPTIONS))
+          register_option_group(name: 'SESSION',
+                                description: 'Used when connecting via an existing SESSION',
+                                option_names: ['SESSION'])
+          register_option_group(name: 'RHOST',
+                                description: 'Used when making a new connection via RHOSTS',
+                                option_names: RHOST_GROUP_OPTIONS)
           register_options(
             [
               Msf::OptInt.new('SESSION', [ false, 'The session to run this module on' ]),

diff --git a/lib/msf/core/optional_session/mysql.rb b/lib/msf/core/optional_session/mysql.rb
@@ -16,12 +16,12 @@ def initialize(info = {})
         )
 
         if framework.features.enabled?(Msf::FeatureManager::MYSQL_SESSION_TYPE)
-          options.add_option_group(Msf::OptionGroup.new(name: 'SESSION',
-                                                        description: 'Used when connecting via an existing SESSION',
-                                                        option_names: ['SESSION']))
-          options.add_option_group(Msf::OptionGroup.new(name: 'RHOST',
-                                                        description: 'Used when making a new connection via RHOSTS',
-                                                        option_names: RHOST_GROUP_OPTIONS))
+          register_option_group(name: 'SESSION',
+                                description: 'Used when connecting via an existing SESSION',
+                                option_names: ['SESSION'])
+          register_option_group(name: 'RHOST',
+                                description: 'Used when making a new connection via RHOSTS',
+                                option_names: RHOST_GROUP_OPTIONS)
           register_options(
             [
               Msf::OptInt.new('SESSION', [ false, 'The session to run this module on' ]),

diff --git a/lib/msf/core/optional_session/postgresql.rb b/lib/msf/core/optional_session/postgresql.rb
@@ -16,12 +16,12 @@ def initialize(info = {})
         )
 
         if framework.features.enabled?(Msf::FeatureManager::POSTGRESQL_SESSION_TYPE)
-          options.add_option_group(Msf::OptionGroup.new(name: 'SESSION',
-                                                        description: 'Used when connecting via an existing SESSION',
-                                                        option_names: ['SESSION']))
-          options.add_option_group(Msf::OptionGroup.new(name: 'RHOST',
-                                                        description: 'Used when making a new connection via RHOSTS',
-                                                        option_names: RHOST_GROUP_OPTIONS))
+          register_option_group(name: 'SESSION',
+                                description: 'Used when connecting via an existing SESSION',
+                                option_names: ['SESSION'])
+          register_option_group(name: 'RHOST',
+                                description: 'Used when making a new connection via RHOSTS',
+                                option_names: RHOST_GROUP_OPTIONS)
           register_options(
             [
               Msf::OptInt.new('SESSION', [ false, 'The session to run this module on' ]),

diff --git a/lib/msf/core/optional_session/smb.rb b/lib/msf/core/optional_session/smb.rb
@@ -16,12 +16,12 @@ def initialize(info = {})
         )
 
         if framework.features.enabled?(Msf::FeatureManager::SMB_SESSION_TYPE)
-          options.add_option_group(Msf::OptionGroup.new(name: 'SESSION',
-                                                        description: 'Used when connecting via an existing SESSION',
-                                                        option_names: ['SESSION']))
-          options.add_option_group(Msf::OptionGroup.new(name: 'RHOST',
-                                                        description: 'Used when making a new connection via RHOSTS',
-                                                        option_names: RHOST_GROUP_OPTIONS))
+          register_option_group(name: 'SESSION',
+                                description: 'Used when connecting via an existing SESSION',
+                                option_names: ['SESSION'])
+          register_option_group(name: 'RHOST',
+                                description: 'Used when making a new connection via RHOSTS',
+                                option_names: RHOST_GROUP_OPTIONS)
           register_options(
             [
               Msf::OptInt.new('SESSION', [ false, 'The session to run this module on' ]),

diff --git a/spec/lib/msf/base/serializer/readable_text_spec.rb b/spec/lib/msf/base/serializer/readable_text_spec.rb
@@ -190,7 +190,7 @@ def initialize
       let(:group) { Msf::OptionGroup.new(name: group_name, description: group_description, option_names: option_names) }
       let(:aux_mod_with_grouped_options) do
         mod = aux_mod_with_set_options.replicant
-        mod.options.add_option_group(group)
+        mod.options.add_group(group)
         mod
       end
 
@@ -231,8 +231,8 @@ def initialize
 
       let(:aux_mod_with_grouped_options) do
         mod = aux_mod_with_set_options.replicant
-        mod.options.add_option_group(group_1)
-        mod.options.add_option_group(group_2)
+        mod.options.add_group(group_1)
+        mod.options.add_group(group_2)
         mod
       end