#
# ronin-exploits - A Ruby library for ronin-rb that provides exploitation and
# payload crafting functionality.
#
# Copyright (c) 2007-2022 Hal Brodigan (postmodern.mod3 at gmail.com)
#
# ronin-exploits is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# ronin-exploits is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with ronin-exploits.  If not, see <https://www.gnu.org/licenses/>.
#

require 'ronin/exploits/exceptions'
require 'ronin/payloads/payload'

module Ronin
  module Exploits
    module Mixins
      #
      # Adds the ability to use a payload in an exploit.
      #
      # ## Examples
      #
      #     module Ronin
      #       module Exploits
      #         class MyExploit < Exploit
      #     
      #           include Mixins::HasPayload
      #     
      #           payload_class Ronin::Payloads::JavaScriptPayload
      #     
      #         end
      #       end
      #     end
      #
      # @api public
      #
      module HasPayload
        #
        # Adds {ClassMethods} to the exploit.
        #
        # @param [Class] exploit
        #   The exploit class including {HasPayload}.
        #
        # @api private
        #
        def self.included(exploit)
          exploit.extend ClassMethods
        end

        #
        # Class methods.
        #
        module ClassMethods
          #
          # Gets or sets the payload base class that is compatible with the
          # exploit.
          #
          # @param [Class<Ronin::Payloads::Payload>, nil] new_payload_class
          #   The optional new payload base class to set.
          #
          # @return [Class<Ronin::Payloads::Payload>]
          #   The exploit's compatible payload base class.
          #
          def payload_class(new_payload_class=nil)
            if new_payload_class
              @payload_class = new_payload_class
            else
              @payload_class ||= if superclass.kind_of?(ClassMethods)
                                   superclass.payload_class
                                 else
                                   Ronin::Payloads::Payload
                                 end
            end
          end
        end

        # The payload the exploit can use.
        #
        # @return [Ronin::Payloads::Payload, String, nil]
        attr_reader :payload

        #
        # Initializes the exploit and sets the {#payload}.
        #
        # @param [Ronin::Payloads::Payload, String, nil] payload
        #   The payload to use.
        #
        def initialize(payload: nil, **kwargs)
          super(**kwargs)

          self.payload = payload
        end

        #
        # Sets the payload to use with the exploit.
        #
        # @param [Ronin::Payloads::Payload, String, nil] new_payload
        #   The new payload to use with the exploit.
        #
        # @return [Ronin::Payloads::Payload, String, nil]
        #   The new payload of the exploit.
        #
        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}")
            end
          end

          @payload = new_payload
        end

        #
        # Validates {#payload} and the exploit.
        #
        # @raise [MissingPayload]
        #   {#payload} was never set.
        #
        # @raise [Ronin::Core::Params::RequiredParam]
        #   One of the required params in the exploit or {#payload} is not
        #   set.
        #
        # @api semipublic
        #
        def perform_validate
          unless @payload
            raise(MissingPayload,"exploit requires a payload")
          end

          if @payload.kind_of?(Ronin::Core::Params::Mixin)
            @payload.validate_params
          end

          super
        end

        #
        # Calls the payload's `perform_build` method first before the exploit
        # is built.
        #
        def perform_build
          if @payload.kind_of?(Ronin::Payloads::Payload)
            @payload.perform_build
          end

          super
        end

        #
        # Overrides the payload's `perform_prelaunch` method, then calls the
        # exploit's {Exploit#perform_launch perform_launch} method, and finally
        # calls the payload's `perform_postlaunch` method.
        #
        # @note
        #   If any exception is raised by the exploit's `launch` method, then
        #   the payload's `perform_cleanup` method is called and the exception
        #   is re-raised.
        #
        def perform_launch
          if @payload.kind_of?(Ronin::Payloads::Payload)
            @payload.perform_prelaunch
          end

          begin
            super()

            if @payload.kind_of?(Ronin::Payloads::Payload)
              @payload.perform_postlaunch
            end
          rescue => error
            if @payload.kind_of?(Ronin::Payloads::Payload)
              @payload.perform_cleanup
            end

            raise(error)
          end
        end

        #
        # Calls the payload's `perform_cleanup` method first after the exploit
        # is cleaned up.
        #
        def perform_cleanup
          super

          if @payload.kind_of?(Ronin::Payloads::Payload)
            @payload.perform_cleanup
          end
        end
      end
    end
  end
end