Switched to storing multiple payloads in payload_class as an Array …

…of classes.

lib/ronin/exploits/cli/commands/show.rb

@@ -163,7 +163,14 @@ def print_metadata(exploit)
 
             if defined?(Mixins::HasPayload) &&
                exploit.include?(Mixins::HasPayload)
-              fields['Payload Type'] = payload_type(exploit.payload_class)
+              fields['Payload Type'] = case (payload_class = exploit.payload_class)
+                                       when Array
+                                         payload_class.map { |klass|
+                                           payload_type(klass)
+                                         }.join(', ')
+                                       else
+                                         payload_type(payload_class)
+                                       end
             end
 
             fields['Summary'] = exploit.summary if exploit.summary

lib/ronin/exploits/mixins/has_payload.rb

@@ -58,57 +58,6 @@ def self.included(exploit)
           exploit.extend ClassMethods
         end
 
-        #
-        # Helper class for matching a payload against multiple accepted payload
-        # classes.
-        #
-        # @api private
-        #
-        # @since 1.3.0
-        #
-        class PayloadClasses
-
-          # The payload classes to match against.
-          #
-          # @return [Array<Class<Ronin::Payloads::Payload>>]
-          attr_reader :classes
-
-          #
-          # Initializes the payload classes.
-          #
-          # @param [Array<Class<Ronin::Payloads::Payload>>] classes
-          #   The payload classes.
-          #
-          def initialize(classes)
-            @classes = classes
-          end
-
-          #
-          # Compares the payload object to the payload classes.
-          #
-          # @param [Ronin::Payloads::Payload] payload
-          #   The payload object to match.
-          #
-          # @return [Boolean]
-          #   Indicates whether the payload object inherits from any of the
-          #   payload classes.
-          #
-          def ===(payload)
-            @classes.any? { |payload_class| payload_class === payload }
-          end
-
-          #
-          # Converts the list payload classes to a String.
-          #
-          # @return [String]
-          #   The comma separated list of payload class names.
-          #
-          def to_s
-            @classes.join(', ')
-          end
-
-        end
-
         #
         # Class methods.
         #
@@ -117,18 +66,18 @@ module ClassMethods
           # Gets or sets the payload base class that is compatible with the
           # exploit.
           #
-          # @param [Class<Ronin::Payloads::Payload>, nil] new_payload_class
+          # @param [Array<Class<Ronin::Payloads::Payload>>] new_payload_classes
           #   The optional new payload base class to set.
           #
-          # @return [Class<Ronin::Payloads::Payload>]
+          # @return [Class<Ronin::Payloads::Payload>, Array<Class<Ronin::Payloads::Payload>>]
           #   The exploit's compatible payload base class.
           #
           def payload_class(*new_payload_classes)
             unless new_payload_classes.empty?
               @payload_class = if new_payload_classes.length == 1
                                  new_payload_classes.first
                                else
-                                 PayloadClasses.new(new_payload_classes)
+                                 new_payload_classes
                                end
             else
               @payload_class ||= if superclass.kind_of?(ClassMethods)
@@ -168,8 +117,15 @@ def initialize(payload: nil, **kwargs)
         #
         def payload=(new_payload)
           if new_payload.kind_of?(Payloads::Payload)
-            unless self.class.payload_class === new_payload
-              raise(IncompatiblePayload,"incompatible payload, must be a #{self.class.payload_class} payload: #{new_payload.inspect}")
+            case (payload_class = self.class.payload_class)
+            when Array
+              unless payload_class.any? { |klass| new_payload.kind_of?(klass) }
+                raise(IncompatiblePayload,"incompatible payload, must be a #{payload_class.join(', ')} payload: #{new_payload.inspect}")
+              end
+            else
+              unless new_payload.kind_of?(payload_class)
+                raise(IncompatiblePayload,"incompatible payload, must be a #{payload_class} payload: #{new_payload.inspect}")
+              end
             end
           end
 

spec/mixins/has_payload_spec.rb

@@ -41,46 +41,6 @@ class InheritesAndOverridesPayloadClass < WithPayloadClass
     end
   end
 
-  describe described_class::PayloadClasses do
-    let(:classes) do
-      [TestHasPayload::TestPayload, TestHasPayload::TestPayload2]
-    end
-
-    subject { described_class.new(classes) }
-
-    describe "#initialize" do
-      it "must set #classes" do
-        expect(subject.classes).to eq(classes)
-      end
-    end
-
-    describe "#===" do
-      context "when the given payload object is kind of one of the payload classes" do
-        let(:payload) { TestHasPayload::TestPayload2.new }
-
-        it "must return true" do
-          expect(subject === payload).to be(true)
-        end
-      end
-
-      context "when the given payload object is not kind of any of the payload classes" do
-        let(:payload) { TestHasPayload::TestOtherPayload.new }
-
-        it "must return false" do
-          expect(subject === payload).to be(false)
-        end
-      end
-    end
-
-    describe "#to_s" do
-      it "must return a comma-separated String of the class names" do
-        expect(subject.to_s).to eq(
-          "#{TestHasPayload::TestPayload}, #{TestHasPayload::TestPayload2}"
-        )
-      end
-    end
-  end
-
   describe ".payload_class" do
     subject { test_class }
 
@@ -105,8 +65,7 @@ class InheritesAndOverridesPayloadClass < WithPayloadClass
         let(:test_class) { TestHasPayload::WithPayloadClasses }
 
         it "must set payload_class to a #{described_class}::PayloadClasses objects with the given payload classes" do
-          expect(subject.payload_class).to be_kind_of(described_class::PayloadClasses)
-          expect(subject.payload_class.classes).to eq(
+          expect(subject.payload_class).to eq(
             [TestHasPayload::TestPayload, TestHasPayload::TestPayload2]
           )
         end
@@ -256,7 +215,7 @@ class InheritesAndOverridesPayloadClass < WithPayloadClass
           let(:test_class) { TestHasPayload::WithPayloadClasses }
 
           context "and the given payload object is a kind of payload_class" do
-            let(:payload) { test_class.payload_class.classes.last.new }
+            let(:payload) { test_class.payload_class.last.new }
 
             before { subject.payload = payload }
 
@@ -281,7 +240,7 @@ class InheritesAndOverridesPayloadClass < WithPayloadClass
             it do
               expect {
                 subject.payload = payload
-              }.to raise_error(Ronin::Exploits::IncompatiblePayload,"incompatible payload, must be a #{test_class.payload_class} payload: #{payload.inspect}")
+              }.to raise_error(Ronin::Exploits::IncompatiblePayload,"incompatible payload, must be a #{test_class.payload_class.join(', ')} payload: #{payload.inspect}")
             end
           end
         end