Land rapid7#18875, Add conditional option validation depending on SESSION/RHOST connection

Author: cgranleese-r7

lib/msf/core/module/options.rb

@@ -73,14 +73,18 @@ def register_options(options, owner = self.class)
   # @param name [String] Name for the group
   # @param description [String] Description of the group
   # @param option_names [Array<String>] List of datastore option names
+  # @param required_options [Array<String>] List of required datastore option names
   # @param merge [Boolean] whether to merge or overwrite the groups option names
-  def register_option_group(name:, description:, option_names: [], merge: true)
+  def register_option_group(name:, description:, option_names: [], required_options: [], 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)
+      option_group = Msf::OptionGroup.new(name: name,
+                                          description: description,
+                                          option_names: option_names,
+                                          required_options: required_options)
       options.add_group(option_group)
     end
   end

lib/msf/core/option_group.rb

@@ -9,14 +9,18 @@ class OptionGroup
     attr_accessor :description
     # @return [Array<String>] List of datastore option names
     attr_accessor :option_names
+    # @return [Array<String>] List of options that if present must have a value set
+    attr_accessor :required_options
 
     # @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: [])
+    # @param required_options [Array<String>] List of options that if present must have a value set
+    def initialize(name:, description:, option_names: [], required_options: [])
       self.name = name
       self.description = description
       self.option_names = option_names
+      self.required_options = required_options
     end
 
     # @param option_name [String] Name of the datastore option to be added to the group
@@ -28,5 +32,19 @@ def add_option(option_name)
     def add_options(option_names)
       @option_names.concat(option_names)
     end
+
+    # Validates that any registered and required options are set
+    #
+    # @param options [Array<Msf::OptBase>] A modules registered options
+    # @param datastore [Msf::DataStore|Msf::DataStoreWithFallbacks] A modules datastore
+    def validate(options, datastore)
+      issues = {}
+      required_options.each do |option_name|
+        if options[option_name] && !datastore[option_name]
+          issues[option_name] = "#{option_name} must be specified"
+        end
+      end
+      raise Msf::OptionValidateError.new(issues.keys.to_a, reasons: issues) unless issues.empty?
+    end
   end
 end

lib/msf/core/optional_session.rb

@@ -7,5 +7,52 @@
 module Msf
   module OptionalSession
     include Msf::SessionCompatibility
+
+    # Validates options depending on whether we are using  SESSION or an RHOST for our connection
+    def validate
+      super
+      return unless optional_session_enabled?
+
+      # If the session is set use that by default regardless of rhost being (un)set
+      if session
+        validate_session
+      elsif rhost
+        validate_rhost
+      else
+        raise Msf::OptionValidateError.new(message: 'A SESSION or RHOST must be provided')
+      end
+    end
+
+    def session
+      return nil unless optional_session_enabled?
+
+      super
+    end
+
+    protected
+
+    # Used to validate options when RHOST has been set
+    def validate_rhost
+      validate_group('RHOST')
+    end
+
+    # Used to validate options when SESSION has been set
+    def validate_session
+      issues = {}
+      if session_types && !session_types.include?(session.type)
+        issues['SESSION'] = "Incompatible session type: #{session.type}. This module works with: #{session_types.join(', ')}."
+      end
+      raise Msf::OptionValidateError.new(issues.keys.to_a, reasons: issues) unless issues.empty?
+
+      validate_group('SESSION')
+    end
+
+    # Validates the options within an option group
+    #
+    # @param group_name [String] Name of the option group
+    def validate_group(group_name)
+      option_group = options.groups[group_name]
+      option_group.validate(options, datastore) if option_group
+    end
   end
 end

lib/msf/core/optional_session/mssql.rb

@@ -15,13 +15,14 @@ def initialize(info = {})
           )
         )
 
-        if framework.features.enabled?(Msf::FeatureManager::MSSQL_SESSION_TYPE)
+        if optional_session_enabled?
           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)
+                                option_names: RHOST_GROUP_OPTIONS,
+                                required_options: RHOST_GROUP_OPTIONS)
           register_options(
             [
               Msf::OptInt.new('SESSION', [ false, 'The session to run this module on' ]),
@@ -36,10 +37,8 @@ def initialize(info = {})
         end
       end
 
-      def session
-        return nil unless framework.features.enabled?(Msf::FeatureManager::MSSQL_SESSION_TYPE)
-
-        super
+      def optional_session_enabled?
+        framework.features.enabled?(Msf::FeatureManager::MSSQL_SESSION_TYPE)
       end
     end
   end

lib/msf/core/optional_session/mysql.rb

@@ -15,13 +15,14 @@ def initialize(info = {})
           )
         )
 
-        if framework.features.enabled?(Msf::FeatureManager::MYSQL_SESSION_TYPE)
+        if optional_session_enabled?
           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)
+                                option_names: RHOST_GROUP_OPTIONS,
+                                required_options: RHOST_GROUP_OPTIONS)
           register_options(
             [
               Msf::OptInt.new('SESSION', [ false, 'The session to run this module on' ]),
@@ -34,10 +35,8 @@ def initialize(info = {})
         end
       end
 
-      def session
-        return nil unless framework.features.enabled?(Msf::FeatureManager::MYSQL_SESSION_TYPE)
-
-        super
+      def optional_session_enabled?
+        framework.features.enabled?(Msf::FeatureManager::MYSQL_SESSION_TYPE)
       end
     end
   end

lib/msf/core/optional_session/postgresql.rb

@@ -15,13 +15,14 @@ def initialize(info = {})
           )
         )
 
-        if framework.features.enabled?(Msf::FeatureManager::POSTGRESQL_SESSION_TYPE)
+        if optional_session_enabled?
           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)
+                                option_names: RHOST_GROUP_OPTIONS,
+                                required_options: RHOST_GROUP_OPTIONS)
           register_options(
             [
               Msf::OptInt.new('SESSION', [ false, 'The session to run this module on' ]),
@@ -36,10 +37,8 @@ def initialize(info = {})
         end
       end
 
-      def session
-        return nil unless framework.features.enabled?(Msf::FeatureManager::POSTGRESQL_SESSION_TYPE)
-
-        super
+      def optional_session_enabled?
+        framework.features.enabled?(Msf::FeatureManager::POSTGRESQL_SESSION_TYPE)
       end
     end
   end

lib/msf/core/optional_session/smb.rb

@@ -15,13 +15,14 @@ def initialize(info = {})
           )
         )
 
-        if framework.features.enabled?(Msf::FeatureManager::SMB_SESSION_TYPE)
+        if optional_session_enabled?
           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)
+                                option_names: RHOST_GROUP_OPTIONS,
+                                required_options: RHOST_GROUP_OPTIONS)
           register_options(
             [
               Msf::OptInt.new('SESSION', [ false, 'The session to run this module on' ]),
@@ -34,10 +35,8 @@ def initialize(info = {})
         end
       end
 
-      def session
-        return nil unless framework.features.enabled?(Msf::FeatureManager::SMB_SESSION_TYPE)
-
-        super
+      def optional_session_enabled?
+        framework.features.enabled?(Msf::FeatureManager::SMB_SESSION_TYPE)
       end
     end
   end

lib/msf/ui/formatter/option_validate_error.rb

@@ -13,7 +13,11 @@ def self.print_error(mod, error)
           raise ArgumentError, "invalid error type #{error.class}, expected ::Msf::OptionValidateError" unless error.is_a?(::Msf::OptionValidateError)
 
           if error.reasons.empty?
-            mod.print_error("#{error.class} The following options failed to validate: #{error.options.join(', ')}")
+            if error.message
+              mod.print_error("#{error.class} #{error.message}")
+            else
+              mod.print_error("#{error.class} The following options failed to validate: #{error.options.join(', ')}")
+            end
           else
             mod.print_error("#{error.class} The following options failed to validate:")
             error.options.sort.each do |option_name|

spec/lib/msf/core/option_group_spec.rb

@@ -47,4 +47,92 @@
       end
     end
   end
+
+  describe '#validate' do
+    let(:required_option_name) { 'required_name' }
+    let(:option_names) { ['not_required_name', required_option_name] }
+    let(:required_names) { [required_option_name] }
+    let(:options) { instance_double(Msf::OptionContainer) }
+    let(:datastore) { instance_double(Msf::DataStoreWithFallbacks) }
+
+    context 'when there are no required options' do
+      subject { described_class.new(name: 'name', description: 'description', option_names: option_names) }
+
+      context 'when no values are set for the options' do
+
+        before(:each) do
+          allow(options).to receive(:[]).and_return(instance_double(Msf::OptBase))
+          allow(datastore).to receive(:[]).and_return(nil)
+        end
+
+        it 'validates the options in the group' do
+          expect { subject.validate(options, datastore) }.not_to raise_error(Msf::OptionValidateError)
+        end
+      end
+
+      context 'when values are set for the options' do
+
+        before(:each) do
+          allow(options).to receive(:[]).and_return(instance_double(Msf::OptBase))
+          allow(datastore).to receive(:[]).and_return('OptionValue')
+        end
+
+        it 'validates the options in the group' do
+          expect { subject.validate(options, datastore) }.not_to raise_error(Msf::OptionValidateError)
+        end
+      end
+
+      context 'when the options have not been registered' do
+
+        before(:each) do
+          allow(options).to receive(:[]).and_return(nil)
+        end
+
+        it 'does not attempt to validate the options' do
+          expect { subject.validate(options, datastore) }.not_to raise_error(Msf::OptionValidateError)
+        end
+      end
+    end
+
+    context 'when there is a required option' do
+      subject { described_class.new(name: 'name', description: 'description', option_names: option_names, required_options: required_names) }
+      let(:error_message) { "The following options failed to validate: #{required_option_name}." }
+
+      context 'when no values are set for the options' do
+
+        before(:each) do
+          allow(options).to receive(:[]).and_return(instance_double(Msf::OptBase))
+          allow(datastore).to receive(:[]).and_return(nil)
+        end
+
+        it 'raises an error only for the required option' do
+          expect { subject.validate(options, datastore) }.to raise_error(Msf::OptionValidateError).with_message(error_message)
+        end
+      end
+
+      context 'when values are set for the options' do
+
+        before(:each) do
+          allow(options).to receive(:[]).and_return(instance_double(Msf::OptBase))
+          allow(datastore).to receive(:[]).and_return('OptionValue')
+        end
+
+        it 'validates the options in the group' do
+          expect { subject.validate(options, datastore) }.not_to raise_error(Msf::OptionValidateError)
+        end
+      end
+
+      context 'when the options have not been registered' do
+
+        before(:each) do
+          allow(options).to receive(:[]).and_return(nil)
+        end
+
+        it 'does not attempt to validate the options' do
+          expect { subject.validate(options, datastore) }.not_to raise_error(Msf::OptionValidateError)
+        end
+      end
+    end
+
+  end
 end

spec/lib/msf/core/optional_session/mssql_spec.rb

@@ -0,0 +1,19 @@
+# -*- coding:binary -*-
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Msf::OptionalSession::MSSQL do
+  subject(:mod) do
+    mod = ::Msf::Module.new
+    mod.extend described_class
+    mod
+  end
+
+  before(:each) do
+    allow(Msf::FeatureManager.instance).to receive(:enabled?).and_call_original
+    allow(Msf::FeatureManager.instance).to receive(:enabled?).with(Msf::FeatureManager::MSSQL_SESSION_TYPE).and_return(true)
+  end
+
+  it_behaves_like Msf::OptionalSession
+end