# frozen_string_literal: true
#
# ronin-exploits - A Ruby library for ronin-rb that provides exploitation and
# payload crafting functionality.
#
# Copyright (c) 2007-2023 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/cli/exploit_command'
require 'ronin/exploits/cli/ruby_shell'
require 'ronin/exploits/mixins/has_payload'
require 'ronin/exploits/mixins/has_targets'
require 'ronin/exploits/mixins/loot'

require 'ronin/payloads/cli/encoder_methods'
require 'ronin/payloads/cli/payload_methods'
require 'ronin/payloads/mixins/post_ex'

require 'ronin/core/cli/options/param'
require 'ronin/core/cli/options/values/arches'
require 'ronin/core/cli/options/values/oses'
require 'ronin/core/cli/logging'

require 'command_kit/printing/indent'

module Ronin
  module Exploits
    class CLI
      module Commands
        #
        # Runs an exploit.
        #
        # ## Usage
        #
        #     ronin-exploits run [options] {NAME | -f FILE}
        #
        # ## Options
        #
        #     -f, --file FILE                  The exploit file to load
        #     -p, --param NAME=VALUE           Sets a param
        #     -D, --dry-run                    Builds the exploit but does not launch it
        #         --payload-file FILE          Load the payload from the given Ruby file
        #         --read-payload FILE          Reads the payload string from the file
        #         --payload-string STRING      Uses the raw payload string instead
        #     -P, --payload NAME               The payload to load and use
        #         --payload-param NAME=VALUE   Sets a param in the payload
        #         --encoder-file FILE          Load the payload encoder from the Ruby file
        #     -E, --encoder NAME               Loads the payload encoder by name
        #         --encoder-param ENCODER.NAME=VALUE
        #                                      Sets a param of the ENCODER
        #     -t, --target INDEX               Selects the target by index
        #     -A x86|x86-64|amd64|ia64|ppc|ppc64|arm|armbe|arm64|arm64be|mips|mipsle|mips64|mips64le,
        #         --target-arch                Selects the target with the matching arch
        #     -O linux|macos|windows|freebsd|openbsd|netbsd,
        #         --target-os                  Selects the target with the matching OS
        #         --target-os-version VERSION  Selects the target with the matching OS version
        #     -S, --target-software NAME       Selects the target with the matching software name
        #     -V, --target-version VERSION     Selects the target with the matching software version
        #     -L, --save-loot DIR              Saves any found loot to the DIR
        #     -d, --debug                      Enables debugging messages
        #         --irb                        Open an interactive Ruby shell inside the exploit
        #     -h, --help                       Print help information
        #
        # ## Arguments
        #
        #     [NAME]                           The exploit name to load
        #
        class Run < ExploitCommand

          include Payloads::CLI::EncoderMethods
          include Payloads::CLI::PayloadMethods
          include Core::CLI::Options::Param
          include Core::CLI::Logging
          include CommandKit::Printing::Indent

          # Exploit options
          option :dry_run, short: '-D',
                           desc: 'Builds the exploit but does not launch it'

          # Payload options
          option :payload_file, value: {
                                  type: String,
                                  usage: 'FILE'
                                },
                                desc: 'Load the payload from the given Ruby file'
          option :read_payload, value: {
                                  type:  String,
                                  usage: 'FILE'
                                },
                                desc: 'Reads the payload string from the file'

          option :payload_string, value: {
                                    type:  String,
                                    usage: 'STRING'
                                  },
                                  desc: 'Uses the raw payload string instead'

          option :payload, short: '-P',
                           value: {
                             type: String,
                             usage: 'NAME'
                           },
                           desc: 'The payload to load and use'

          option :payload_param, value: {
                                   type: /\A[^=\s]+=.+\z/,
                                   usage: 'NAME=VALUE'
                                 },
                                 desc: 'Sets a param on the payload' do |param|
                                   name, value = param.split('=',2)

                                   @payload_params[name.to_sym] = value
                                 end

          # Encoder options
          option :encoder_file, value: {
                                  type: String,
                                  usage: 'FILE'
                                },
                                desc: 'Load the payload encoder from the Ruby file' do |file|
                                  @encoders_to_load << [:file, file]
                                end

          option :encoder, short: '-E',
                           value: {
                             type: String,
                             usage: 'NAME'
                           },
                           desc: 'Loads the payload encoder by name' do |name|
                             @encoders_to_load << [:name, name]
                           end

          option :encoder_param, value: {
                                   type: /\A[^\.\=\s]+\.[^=\s]+=.+\z/,
                                   usage: 'ENCODER.NAME=VALUE'
                                 },
                                 desc: 'Sets a param on the ENCODER' do |str|
                                   prefix, value = str.split('=',2)
                                   encoder, name = prefix.split('.',2)

                                   @encoder_params[encoder][name.to_sym] = value
                                 end

          # Target options
          option :target, short: '-t',
                          value: {
                            type: Integer,
                            usage: 'INDEX'
                          },
                          desc: 'Selects the target by index'

          option :target_arch, short: '-A',
                               value: {
                                 type: Core::CLI::Options::Values::ARCHES
                               },
                               desc: 'Selects the target with the matching arch' do |arch|
                                 @target_kwargs[:arch] = arch
                               end

          option :target_os, short: '-O',
                             value: {
                               type: Core::CLI::Options::Values::OSES
                             },
                             desc: 'Selects the target with the matching OS' do |os|
                               @target_kwargs[:os] = os
                             end

          option :target_os_version, value: {
                                       type: String,
                                       usage: 'VERSION'
                                     },
                                     desc: 'Selects the target with the matching OS version' do |version|
                                       @target_kwargs[:os_version] = version
                                     end

          option :target_software, short: '-S',
                                   value: {
                                     type: String,
                                     usage: 'NAME'
                                   },
                                   desc: 'Selects the target with the matching software name' do |software|
                                     @target_kwargs[:software] = software
                                   end

          option :target_version, short: '-V',
                                  value: {
                                    type: String,
                                    usage: 'VERSION'
                                  },
                                  desc: 'Selects the target with the matching software version' do |version|
                                    @target_kwargs[:version] = version
                                  end

          option :save_loot, short: '-L',
                             value: {
                               type:  String,
                               usage: 'DIR'
                             },
                             desc: 'Saves any found loot to the DIR'

          option :debug, short: '-d',
                         desc: 'Enables debugging messages' do
                           Support::CLI::Printing.debug = true
                         end

          option :irb, desc: 'Open an interactive Ruby shell inside the exploit'

          description 'Runs an exploit'

          man_page 'ronin-exploits-run.1'

          # Thte encoder names and paths to load.
          #
          # @return [Array<(Symbol, String)>]
          attr_reader :encoders_to_load

          # The encoder params.
          #
          # @return [Hash{String => Hash{String => String}}]
          attr_reader :encoder_params

          # The payload params.
          #
          # @return [Hash{String => String}]
          attr_reader :payload_params

          # The keyword arguments to select a target with.
          #
          # @return [Hash{Symbol => Object}]
          attr_reader :target_kwargs

          #
          # Initializes the `ronin-exploits run` command.
          #
          # @param [Hash{Symbol => Object}] kwargs
          #   Additional keyword arguments.
          #
          def initialize(**kwargs)
            super(**kwargs)

            @encoders_to_load  = []
            @encoder_params    = Hash.new { |hash,key| hash[key] = {} }
            @payload_params    = {}
            @target_kwargs     = {}
          end

          #
          # Runs the `ronin-exploits run` command.
          #
          # @param [String, nil] name
          #   The optional exploit name to load.
          #
          def run(name=nil)
            super(name)

            load_encoders
            load_payload
            initialize_encoders
            initialize_payload
            validate_payload
            initialize_exploit
            validate_exploit
            run_exploit

            if options[:irb]
              start_shell
            else
              post_exploitation
            end

            perform_cleanup
          end

          #
          # Loads the payload encoder classes specified by `--encoder` or
          # `--encoder-file`.
          #
          def load_encoders
            @encoder_classes = @encoders_to_load.map do |(type,value)|
              case type
              when :name then load_encoder(value)
              when :file then load_encoder_from(value)
              else
                raise(NotImplementedError,"invalid encoder type: #{type.inspect}")
              end
            end
          end

          #
          # Initializes the payload encoders specified by `--encoder` or
          # `--encoder-file`.
          #
          def initialize_encoders
            @encoders = @encoder_classes.map do |encoder_class|
              encoder_class.new(params: @encoder_params[encoder_class.id])
            end
          end

          #
          # Loads the payload class specified by `--payload` or
          # `--payload-file`.
          #
          def load_payload
            @payload_class = if options[:payload]
                               super(options[:payload])
                             elsif options[:payload_file]
                               load_payload_from(options[:payload_file])
                             end
          end

          #
          # Initializes the payload specified by `--payload`, `--payload-file`,
          # `--read-payload`, or `--payload-string`.
          #
          def initialize_payload
            @payload = if @payload_class
                         super(@payload_class, params:   @payload_params,
                                               encoders: @encoders)
                       elsif options[:read_payload]
                         File.binread(options[:read_payload])
                       elsif options[:payload_string]
                         options[:payload_string]
                       end
          end

          #
          # Validates the payload.
          #
          def validate_payload
            super(@payload) if @payload
          end

          #
          # Initializes the exploit.
          #
          def initialize_exploit
            kwargs = {params: @params}

            if @exploit_class.include?(Mixins::HasPayload)
              kwargs[:payload] = @payload
            end

            if @exploit_class.include?(Mixins::HasTargets)
              kwargs[:target] = if options[:target]
                                  options[:target]
                                elsif !@target_kwargs.empty?
                                  @target_kwargs
                                end
            end

            super(**kwargs)
          end

          #
          # Runs the exploit.
          #
          def run_exploit
            log_info "Running exploit #{@exploit.class_id} ..."

            begin
              @exploit.exploit(dry_run: options[:dry_run])
            rescue ExploitError => error
              print_error "failed to run exploit #{@exploit.class_id}: #{error.message}"
              exit(1)
            rescue => error
              print_exception(error)
              print_error "an unhandled exception occurred while running the exploit #{@exploit.class_id}"
              exit(-1)
            end
          end

          #
          # Starts an interactive ruby shell within the exploit object.
          #
          def start_shell
            log_info "Exploit #{@exploit.class_id} launched!"
            log_info "Starting interactive Ruby shell ..."

            RubyShell.start(name: @exploit_class.name, context: @exploit)
          end

          #
          # Performs the post-exploitation stage.
          #
          def post_exploitation
            if @exploit_class.include?(Mixins::HasPayload) &&
               @exploit.payload.kind_of?(Ronin::Payloads::Payload) &&
               @exploit.payload.kind_of?(Ronin::Payloads::Mixins::PostEx)
              unless @exploit.payload.session
                print_error "payload (#{@exploit.payload.class_id}) did not create a post-exploitation session"

                perform_cleanup
                exit(1)
              end

              @exploit.payload.session.system.interact
            elsif @exploit_class.include?(Mixins::Loot)
              print_loot
              save_loot if options[:save_loot]
            end
          end

          #
          # Prints any loot collected by the exploit.
          #
          def print_loot
            unless @exploit.loot.empty?
              log_info "Exploit found the following loot:"

              indent do
                @exploit.loot.each do |file|
                  puts
                  puts "#{file.path}:"
                  puts

                  indent do
                    file.to_s.each_line do |line|
                      puts line
                    end
                  end
                  puts
                end
              end
            else
              log_error "Exploit did not find any loot :("
            end
          end

          #
          # Saves the collected loot to the `--save-loot` directory.
          #
          def save_loot
            @exploit.loot.save(options.fetch(:save_loot))
          end

          #
          # Performs the cleanup stage of the exploit.
          #
          def perform_cleanup
            @exploit.perform_cleanup
          rescue ExploitError => error
            print_error "failed to cleanup exploit #{@exploit.class_id}: #{error.message}"
            exit(1)
          rescue => error
            print_exception(error)
            print_error "an unhandled exception occurred while cleaning up the exploit #{@exploit.class_id}"
            exit(-1)
          end

        end
      end
    end
  end
end