Support specifying multiple payload classes (closes #141).

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

@@ -63,18 +63,28 @@ def self.included(exploit)
         #
         module ClassMethods
           #
-          # Gets or sets the payload base class that is compatible with the
+          # Gets or sets the payload base class(es) that is compatible with the
           # exploit.
           #
-          # @param [Class<Ronin::Payloads::Payload>, nil] new_payload_class
-          #   The optional new payload base class to set.
+          # @param [Array<Class<Ronin::Payloads::Payload>>] new_payload_classes
+          #   The optional new payload base class(es) to set.
           #
-          # @return [Class<Ronin::Payloads::Payload>]
-          #   The exploit's compatible payload base class.
+          # @return [Class<Ronin::Payloads::Payload>, Array<Class<Ronin::Payloads::Payload>>]
+          #   The exploit's compatible payload base class(es).
           #
-          def payload_class(new_payload_class=nil)
-            if new_payload_class
-              @payload_class = new_payload_class
+          # @example
+          #   payload_class ShellcodePayload
+          #
+          # @example setting multiple payload classes:
+          #   payload_class JavaScriptPayload, CommandPayload
+          #
+          def payload_class(*new_payload_classes)
+            unless new_payload_classes.empty?
+              @payload_class = if new_payload_classes.length == 1
+                                 new_payload_classes.first
+                               else
+                                 new_payload_classes
+                               end
             else
               @payload_class ||= if superclass.kind_of?(ClassMethods)
                                    superclass.payload_class
@@ -116,8 +126,15 @@ def initialize(payload: nil, **kwargs)
         #
         def payload=(new_payload)
           if new_payload.kind_of?(Payloads::Payload)
-            unless new_payload.kind_of?(self.class.payload_class)
-              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

@@ -8,6 +8,9 @@ module TestHasPayload
     class TestPayload < Ronin::Payloads::Payload
     end
 
+    class TestPayload2 < Ronin::Payloads::Payload
+    end
+
     class InheritedPayload < TestPayload
     end
 
@@ -24,6 +27,12 @@ class WithPayloadClass < Ronin::Exploits::Exploit
       payload_class TestPayload
     end
 
+    class WithPayloadClasses < Ronin::Exploits::Exploit
+      include Ronin::Exploits::Mixins::HasPayload
+
+      payload_class TestPayload, TestPayload2
+    end
+
     class InheritesPayloadClass < WithPayloadClass
     end
 
@@ -44,10 +53,22 @@ class InheritesAndOverridesPayloadClass < WithPayloadClass
     end
 
     context "when the payload_class has been set in the Exploit class" do
-      let(:test_class) { TestHasPayload::WithPayloadClass }
+      context "with a single payload class" do
+        let(:test_class) { TestHasPayload::WithPayloadClass }
+
+        it "must set the payload_class to the given payload class" do
+          expect(subject.payload_class).to be(TestHasPayload::TestPayload)
+        end
+      end
 
-      it "must set the payload_class to the given payload class" do
-        expect(subject.payload_class).to be(TestHasPayload::TestPayload)
+      context "with multiple payload classes" do
+        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 eq(
+            [TestHasPayload::TestPayload, TestHasPayload::TestPayload2]
+          )
+        end
       end
     end
 
@@ -156,35 +177,71 @@ class InheritesAndOverridesPayloadClass < WithPayloadClass
       end
 
       context "but the Exploit has defined a payload_class" do
-        let(:test_class) { TestHasPayload::WithPayloadClass }
+        context "with a single payload class" do
+          let(:test_class) { TestHasPayload::WithPayloadClass }
 
-        context "and the given payload object is a kind of payload_class" do
-          let(:payload) { test_class.payload_class.new }
+          context "and the given payload object is a kind of payload_class" do
+            let(:payload) { test_class.payload_class.new }
 
-          before { subject.payload = payload }
+            before { subject.payload = payload }
 
-          it "must set #payload" do
-            expect(subject.payload).to be(payload)
+            it "must set #payload" do
+              expect(subject.payload).to be(payload)
+            end
           end
-        end
 
-        context "and the given payload object inherits from payload_class" do
-          let(:payload) { TestHasPayload::InheritedPayload.new }
+          context "and the given payload object inherits from payload_class" do
+            let(:payload) { TestHasPayload::InheritedPayload.new }
 
-          before { subject.payload = payload }
+            before { subject.payload = payload }
 
-          it "must set #payload" do
-            expect(subject.payload).to be(payload)
+            it "must set #payload" do
+              expect(subject.payload).to be(payload)
+            end
+          end
+
+          context "but the given payload is not a kind of payload_class" do
+            let(:payload) { TestHasPayload::TestOtherPayload.new }
+
+            it do
+              expect {
+                subject.payload = payload
+              }.to raise_error(Ronin::Exploits::IncompatiblePayload,"incompatible payload, must be a #{test_class.payload_class} payload: #{payload.inspect}")
+            end
           end
         end
 
-        context "but the given payload is not a kind of payload_class" do
-          let(:payload) { TestHasPayload::TestOtherPayload.new }
+        context "with multiple payload classes" do
+          let(:test_class) { TestHasPayload::WithPayloadClasses }
+
+          context "and the given payload object is a kind of payload_class" do
+            let(:payload) { test_class.payload_class.last.new }
+
+            before { subject.payload = payload }
+
+            it "must set #payload" do
+              expect(subject.payload).to be(payload)
+            end
+          end
+
+          context "and the given payload object inherits from payload_class" do
+            let(:payload) { TestHasPayload::InheritedPayload.new }
+
+            before { subject.payload = payload }
 
-          it do
-            expect {
-              subject.payload = payload
-            }.to raise_error(Ronin::Exploits::IncompatiblePayload,"incompatible payload, must be a #{test_class.payload_class} payload: #{payload.inspect}")
+            it "must set #payload" do
+              expect(subject.payload).to be(payload)
+            end
+          end
+
+          context "but the given payload is not a kind of payload_class" do
+            let(:payload) { TestHasPayload::TestOtherPayload.new }
+
+            it do
+              expect {
+                subject.payload = payload
+              }.to raise_error(Ronin::Exploits::IncompatiblePayload,"incompatible payload, must be a #{test_class.payload_class.join(', ')} payload: #{payload.inspect}")
+            end
           end
         end
       end