Fix issue where newly registered options aren't included in the group

lib/msf/base/serializer/readable_text.rb

@@ -571,7 +571,13 @@ 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 }
 
-    options_grouped_by_conditions = filtered_options.group_by(&:conditions)
+    option_groups = mod.options.option_groups.map { |_name, group| group }
+    options_by_group = option_groups.map do |group|
+      [group, group.option_names.map { |name| mod.options[name] }.compact]
+    end.to_h
+    grouped_option_names = option_groups.flat_map(&:option_names)
+    remaining_options = filtered_options.reject { |option| grouped_option_names.include?(option.name) }
+    options_grouped_by_conditions = remaining_options.group_by(&:conditions)
 
     option_tables = []
 
@@ -587,6 +593,11 @@ def self.dump_options(mod, indent = '', missing = false, advanced: false, evasio
       end
     end
 
+    options_by_group.each do |group, options|
+      tbl = options_table(missing, mod, options, indent)
+      option_tables << "#{indent}#{group.description}:\n\n#{tbl}"
+    end
+
     result = option_tables.join("\n\n")
     result
   end

lib/msf/core/opt.rb

@@ -44,8 +44,8 @@ def self.RHOSTS(default= nil, required=true, desc="The target host(s), see https
       Msf::OptRhosts.new('RHOSTS', [ required, desc, default ], aliases: [ 'RHOST' ])
     end
 
-    def self.RHOST(default=nil, required=true, desc="The target host(s), see https://docs.metasploit.com/docs/using-metasploit/basics/using-metasploit.html", **kwargs)
-      Msf::OptRhosts.new('RHOSTS', [ required, desc, default ], aliases: [ 'RHOST' ], **kwargs)
+    def self.RHOST(default=nil, required=true, desc="The target host(s), see https://docs.metasploit.com/docs/using-metasploit/basics/using-metasploit.html")
+      Msf::OptRhosts.new('RHOSTS', [ required, desc, default ], aliases: [ 'RHOST' ])
     end
 
     # @return [OptPort]

lib/msf/core/opt_base.rb

@@ -26,15 +26,14 @@ class OptBase
     #
     def initialize(in_name, attrs = [],
                    required: false, desc: nil, default: nil, conditions: [], enums: [], regex: nil, aliases: [], max_length: nil,
-                   fallbacks: [], group: nil)
+                   fallbacks: [])
       self.name     = in_name
       self.advanced = false
       self.evasion  = false
       self.aliases  = aliases
       self.max_length = max_length
       self.conditions = conditions
       self.fallbacks = fallbacks
-      self.group = group
 
       if attrs.is_a?(String) || attrs.length == 0
         self.required = required
@@ -231,9 +230,6 @@ def invalid_value_length?(value)
     #
     attr_accessor :max_length
 
-    # @return [Msf::OptionGroup, nil] The option group that this option belongs to if any
-    attr_accessor :group
-
     protected
 
     attr_writer   :required, :desc, :default # :nodoc:

lib/msf/core/option_container.rb

@@ -49,6 +49,7 @@ class OptionContainer < Hash
     #
     def initialize(opts = {})
       self.sorted = []
+      self.option_groups = {}
 
       add_options(opts)
     end
@@ -313,14 +314,23 @@ def merge_sort(other_container)
       result.sort
     end
 
+    def add_option_group(option_group)
+      option_groups[option_group.name] = option_group
+    end
+
     #
     # The sorted array of options.
     #
     attr_reader :sorted
 
+    # @return [Hash<String, Msf::OptionGroup>]
+    attr_reader :option_groups
+
     protected
 
     attr_writer :sorted # :nodoc:
+
+    attr_writer :option_groups
   end
 
 end

lib/msf/core/option_group.rb

@@ -3,11 +3,25 @@
 module Msf
   class OptionGroup
 
-    attr_accessor :name, :description
+    attr_accessor :name, :description, :option_names
 
-    def initialize(name:, description:)
+    # @param name [String] Name for the group
+    # @param description [String] Description to be displayed to the user
+    # @param option_names [Array<String>] List of datastore option names
+    def initialize(name:, description:, option_names: [])
       self.name = name
       self.description = description
+      self.option_names = option_names
+    end
+
+    # @param option_name [String] Name of the datastore option to be added to the group
+    def add_option(option_name)
+      @option_names << option_name
+    end
+
+    # @param option_names [Array<String>] List of datastore option names to be added to the group
+    def add_options(option_names)
+      @option_names.concat(option_names)
     end
   end
 end

lib/msf/core/optional_session.rb

@@ -7,9 +7,5 @@
 module Msf
   module OptionalSession
     include Msf::SessionCompatibility
-
-    def session_enabled?(session_feature_name)
-      framework.features.enabled?(session_feature_name)
-    end
   end
 end

lib/msf/core/optional_session/mssql.rb

@@ -5,6 +5,8 @@ module OptionalSession
     module MSSQL
       include Msf::OptionalSession
 
+      RHOST_GROUP_OPTIONS = %w[RHOSTS RPORT DATABASE USERNAME PASSWORD THREADS]
+
       def initialize(info = {})
         super(
           update_info(
@@ -14,6 +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_options(
             [
               Msf::OptInt.new('SESSION', [ false, 'The session to run this module on' ]),
@@ -23,6 +31,7 @@ def initialize(info = {})
               Msf::Opt::RPORT(1433, false)
             ]
           )
+
           add_info('New in Metasploit 6.4 - This module can target a %grnSESSION%clr or an %grnRHOST%clr')
         end
       end

lib/msf/core/optional_session/mysql.rb

@@ -5,6 +5,8 @@ module OptionalSession
     module MySQL
       include Msf::OptionalSession
 
+      RHOST_GROUP_OPTIONS = %w[RHOSTS RPORT DATABASE USERNAME PASSWORD THREADS]
+
       def initialize(info = {})
         super(
           update_info(
@@ -14,15 +16,21 @@ 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_options(
             [
               Msf::OptInt.new('SESSION', [ false, 'The session to run this module on' ]),
               Msf::Opt::RHOST(nil, false),
               Msf::Opt::RPORT(3306, false)
             ]
           )
-          add_info('New in Metasploit 6.4 - This module can target a %grnSESSION%clr or an %grnRHOST%clr')
 
+          add_info('New in Metasploit 6.4 - This module can target a %grnSESSION%clr or an %grnRHOST%clr')
         end
       end
 

lib/msf/core/optional_session/postgresql.rb

@@ -5,14 +5,23 @@ module OptionalSession
     module PostgreSQL
       include Msf::OptionalSession
 
+      RHOST_GROUP_OPTIONS = %w[RHOSTS RPORT DATABASE USERNAME PASSWORD THREADS]
+
       def initialize(info = {})
         super(
           update_info(
             info,
             'SessionTypes' => %w[postgresql]
           )
         )
+
         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_options(
             [
               Msf::OptInt.new('SESSION', [ false, 'The session to run this module on' ]),
@@ -22,6 +31,7 @@ def initialize(info = {})
               Msf::Opt::RPORT(5432, false)
             ]
           )
+
           add_info('New in Metasploit 6.4 - This module can target a %grnSESSION%clr or an %grnRHOST%clr')
         end
       end

lib/msf/core/optional_session/smb.rb

@@ -5,7 +5,6 @@ module OptionalSession
     module SMB
       include Msf::OptionalSession
 
-      FEATURE_NAME = Msf::FeatureManager::SMB_SESSION_TYPE
       RHOST_GROUP_OPTIONS = %w[RHOSTS RPORT SMBDomain SMBUser SMBPass THREADS]
 
       def initialize(info = {})
@@ -16,36 +15,30 @@ def initialize(info = {})
           )
         )
 
-        if session_enabled?
-          session_group = Msf::OptionGroup.new(name: 'SESSION', description: 'Used when connecting via an existing SESSION')
-          rhost_group = Msf::OptionGroup.new(name: 'RHOST', description: 'Used when making a new connection via RHOSTS')
+        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_options(
             [
-              Msf::OptInt.new('SESSION', [ false, 'The session to run this module on' ], group: session_group),
+              Msf::OptInt.new('SESSION', [ false, 'The session to run this module on' ]),
               Msf::Opt::RHOST(nil, false),
               Msf::Opt::RPORT(445, false),
             ]
           )
 
-          RHOST_GROUP_OPTIONS.each do |option_name|
-            if options[option_name]
-              options[option_name].group = rhost_group
-            end
-          end
-
           add_info('New in Metasploit 6.4 - This module can target a %grnSESSION%clr or an %grnRHOST%clr')
         end
       end
 
       def session
-        return nil unless session_enabled?
+        return nil unless framework.features.enabled?(Msf::FeatureManager::SMB_SESSION_TYPE)
 
         super
       end
-
-      def session_enabled?
-        super(FEATURE_NAME)
-      end
     end
   end
 end

spec/lib/msf/base/serializer/readable_text_spec.rb

@@ -186,14 +186,14 @@ def initialize
     context 'when some options are grouped' do
       let(:group_name) { 'group_name' }
       let(:group_description) { 'Used for example reasons' }
-      let(:group) { Msf::OptionGroup.new(name: group_name, description: group_description) }
+      let(:option_names) { %w[RHOSTS SMBUser SMBDomain] }
+      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['RHOSTS'].group = group
-        mod.options['SMBUser'].group = group
-        mod.options['SMBDomain'].group = group
+        mod.options.add_option_group(group)
         mod
       end
+
       it 'should return the grouped options separate to the rest of the options' do
         expect(described_class.dump_options(aux_mod_with_grouped_options, indent_string, false)).to match_table <<~TABLE
           Name                     Current Setting     Required  Description
@@ -221,18 +221,21 @@ def initialize
     context 'when there are multiple options groups' do
       let(:group_name_1) { 'group_name_1' }
       let(:group_description_1) { 'Used for example reasons_1' }
+      let(:option_names_1) {['RHOSTS']}
       let(:group_name_2) { 'group_name_2' }
       let(:group_description_2) { 'Used for example reasons_2' }
-      let(:group_1) { Msf::OptionGroup.new(name: group_name_1, description: group_description_1) }
-      let(:group_2) { Msf::OptionGroup.new(name: group_name_2, description: group_description_2) }
+      let(:option_names_2) { %w[SMBUser SMBDomain] }
+
+      let(:group_1) { Msf::OptionGroup.new(name: group_name_1, description: group_description_1, option_names: option_names_1) }
+      let(:group_2) { Msf::OptionGroup.new(name: group_name_2, description: group_description_2, option_names: option_names_2) }
 
       let(:aux_mod_with_grouped_options) do
         mod = aux_mod_with_set_options.replicant
-        mod.options['RHOSTS'].group = group_1
-        mod.options['SMBUser'].group = group_2
-        mod.options['SMBDomain'].group = group_2
+        mod.options.add_option_group(group_1)
+        mod.options.add_option_group(group_2)
         mod
       end
+
       it 'should return the grouped options separate to the rest of the options' do
         expect(described_class.dump_options(aux_mod_with_grouped_options, indent_string, false)).to match_table <<~TABLE
           Name                     Current Setting     Required  Description