From b49d6a86ab1e9dba4c3d4583cc4cd04ba31bfba1 Mon Sep 17 00:00:00 2001 From: Postmodern Date: Thu, 8 Aug 2024 22:51:32 -0700 Subject: [PATCH] Support specifying multiple payload classes (closes #141). --- lib/ronin/exploits/cli/commands/show.rb | 9 ++- lib/ronin/exploits/mixins/has_payload.rb | 37 ++++++--- spec/mixins/has_payload_spec.rb | 99 +++++++++++++++++++----- 3 files changed, 113 insertions(+), 32 deletions(-) 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 633c2f3c..2aca496f 100644 --- a/lib/ronin/exploits/mixins/has_payload.rb +++ b/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, nil] new_payload_class - # The optional new payload base class to set. + # @param [Array>] new_payload_classes + # The optional new payload base class(es) to set. # - # @return [Class] - # The exploit's compatible payload base class. + # @return [Class, Array>] + # 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 diff --git a/spec/mixins/has_payload_spec.rb b/spec/mixins/has_payload_spec.rb index cb052160..8b43a388 100644 --- a/spec/mixins/has_payload_spec.rb +++ b/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