metasploit-gs/lib/rex/parser/arguments.rb

# -*- coding: binary -*-
# frozen_string_literal: true
require 'shellwords'

module Rex
  module Parser
    ###
    #
    # This class parses arguments in a getopt style format, kind of.
    # Unfortunately, the default ruby getopt implementation will only
    # work on ARGV, so we can't use it.
    #
    ###
    class Arguments
      # e.g. '-x' or '-xyz'
      SHORT_FLAG = /^-[a-zA-Z]+$/.freeze
      private_constant :SHORT_FLAG
      # e.g. '--verbose', '--very-verbose' or '--very-very-verbose'
      LONG_FLAG = /^--([a-zA-Z]+)(-[a-zA-Z]+)*$/.freeze
      private_constant :LONG_FLAG

      #
      # Initializes the format list with an array of formats like:
      #
      # Arguments.new(
      #    '-b'                => [ false, "some text"                 ],
      #    ['-b']              => [ false, "some text"                 ],
      #    ['-x', '--execute'] => [ true, "mixing long and short args" ],
      #    ['-t', '--test']    => [ true, "testing custom <opt> value", "<arg_to_test>" ],
      #    ['--long-flag']     => [ false, "sample long flag" ]
      # )
      #
      def initialize(fmt)
        normalised_fmt = fmt.map { |key, metadata| [Array(key), metadata] }.to_h
        self.fmt = normalised_fmt
        self.longest = normalised_fmt.each_pair.map { |key, value| key.flatten.join(', ') + (value[0] ? ' ' + value[2].to_s : '') }.max_by(&:length)
      end

      #
      # Takes a string and converts it into an array of arguments.
      #
      def self.from_s(str)
        Shellwords.shellwords(str)
      end

      #
      # Parses the supplied arguments into a set of options.
      #
      def parse(args, &_block)
        skip_next = false

        args.each_with_index do |arg, idx|
          if skip_next
            skip_next = false
            next
          end

          param = nil
          if arg =~ SHORT_FLAG
            arg.split('')[1..-1].each do |letter|
              next unless include?("-#{letter}")

              if arg_required?("-#{letter}")
                param = args[idx + 1]
                skip_next = true
              end

              yield "-#{letter}", idx, param
            end
          elsif arg =~ LONG_FLAG && include?(arg)
            if arg_required?(arg)
              param = args[idx + 1]
              skip_next = true
            end

            # Try to yield the short hand version of our argument if possible
            # This will result in less areas of code that would need to be changed
            to_return = short_arg_from_long_arg(arg)
            if to_return.nil?
              yield arg, idx, param
            else
              yield to_return, idx, param
            end
          else
            # else treat the passed in flag as argument
            yield nil, idx, arg
          end
        end
      end

      #
      # Returns usage information for this parsing context.
      #
      def usage
        txt = ["\nOPTIONS:\n"]

        fmt.sort_by { |key, _metadata| key.to_s.downcase }.each do |key, val|
          # if the arg takes in a parameter, get parameter string
          opt = val[0] ? " #{val[2]}" : ''

          # Get all arguments for a command
          output = key.join(', ')
          output += opt

          # Left align the fmt options and <opt> string
          aligned_option = "    #{output.ljust(longest.length)}"
          txt << "#{aligned_option}  #{val[1]}"
        end

        txt << ""
        txt.join("\n")
      end

      def include?(search)
        fmt.keys.flatten.include?(search)
      end

      def arg_required?(opt)
        value = select_value_from_fmt_option(opt)
        return false if value.nil?

        value.first
      end

      def option_keys
        fmt.keys.flatten
      end

      # Return new Parser object featuring options from the base object and including the options hash that was passed in
      def merge(to_merge)
        return fmt unless to_merge.is_a?(Hash)

        Rex::Parser::Arguments.new(fmt.clone.merge(to_merge))
      end

      private

      attr_accessor :fmt     # :nodoc:
      attr_accessor :longest # :nodoc:

      def select_value_from_fmt_option(option)
        fmt_option = fmt.find { |key, value| value if key.include?(option) }
        return if fmt_option.nil?

        fmt_option[1]
      end

      # Returns the short-flag equivalent of any option passed in, if one exists
      # Returns nil if one does not exist
      def short_arg_from_long_arg(long_arg)
        fmt_option = fmt.find { |key, value| value if key.include?(long_arg) }.first
        # if fmt_option == [long_arg] that means that a short flag option for it does not exist
        return if fmt_option.nil? || fmt_option == [long_arg]

        fmt_option.each { |opt| return opt if opt =~ SHORT_FLAG }
      end
    end
  end
end