diff --git a/lib/ronin/exploits/cli/commands/show.rb b/lib/ronin/exploits/cli/commands/show.rb index 5f7de490..3bff1279 100644 --- a/lib/ronin/exploits/cli/commands/show.rb +++ b/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 diff --git a/lib/ronin/exploits/mixins/has_payload.rb b/lib/ronin/exploits/mixins/has_payload.rb index 6318286f..83959842 100644 --- a/lib/ronin/exploits/mixins/has_payload.rb +++ b/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>] - attr_reader :classes - - # - # Initializes the payload classes. - # - # @param [Array>] 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,10 +66,10 @@ module ClassMethods # Gets or sets the payload base class that is compatible with the # exploit. # - # @param [Class, nil] new_payload_class + # @param [Array>] new_payload_classes # The optional new payload base class to set. # - # @return [Class] + # @return [Class, Array>] # The exploit's compatible payload base class. # def payload_class(*new_payload_classes) @@ -128,7 +77,7 @@ def payload_class(*new_payload_classes) @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 diff --git a/spec/mixins/has_payload_spec.rb b/spec/mixins/has_payload_spec.rb index bdb1e507..8b43a388 100644 --- a/spec/mixins/has_payload_spec.rb +++ b/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