Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place.

1
# -*- coding: binary -*-
2
# frozen_string_literal: true
3
require 'shellwords'
4

5
module Rex
6
  module Parser
7
    ###
8
    #
9
    # This class parses arguments in a getopt style format, kind of.
10
    # Unfortunately, the default ruby getopt implementation will only
11
    # work on ARGV, so we can't use it.
12
    #
13
    ###
14
    class Arguments
15
      # e.g. '-x' or '-xyz'
16
      SHORT_FLAG = /^-[a-zA-Z]+$/.freeze
17
      private_constant :SHORT_FLAG
18
      # e.g. '--verbose', '--very-verbose' or '--very-very-verbose'
19
      LONG_FLAG = /^--([a-zA-Z]+)(-[a-zA-Z]+)*$/.freeze
20
      private_constant :LONG_FLAG
21

22
      #
23
      # Initializes the format list with an array of formats like:
24
      #
25
      # Arguments.new(
26
      #    '-z'                => [ has_argument, "some text", "<argument_description>" ],
27
      #    '-b'                => [ false, "some text"                 ],
28
      #    ['-b']              => [ false, "some text"                 ],
29
      #    ['-x', '--execute'] => [ true, "mixing long and short args" ],
30
      #    ['-t', '--test']    => [ true, "testing custom <opt> value", "<arg_to_test>" ],
31
      #    ['--long-flag']     => [ false, "sample long flag" ]
32
      # )
33
      #
34
      def initialize(fmt)
35
        normalised_fmt = fmt.map { |key, metadata| [Array(key), metadata] }.to_h
36
        self.fmt = normalised_fmt
37
        self.longest = normalised_fmt.each_pair.map { |key, value| key.flatten.join(', ') + (value[0] ? ' ' + value[2].to_s : '') }.max_by(&:length)
38
      end
39

40
      #
41
      # Takes a string and converts it into an array of arguments.
42
      #
43
      def self.from_s(str)
44
        Shellwords.shellwords(str)
45
      end
46

47
      #
48
      # Parses the supplied arguments into a set of options.
49
      #
50
      def parse(args, &_block)
51
        skip_next = 0
52

53
        args.each_with_index do |arg, idx|
54
          if skip_next > 0
55
            skip_next -= 1
56
            next
57
          end
58

59
          param = nil
60
          if arg =~ SHORT_FLAG
61
            # parsing needs to take into account a couple requirements
62
            #  1. longest `short` flag found in 'arg' should be extracted first
63
            #    * consider passing the short flag to a tokenizer that returns a list of tokens in order with any invalid tokens
64
            #  2. any short flag arguments that need an option will consume the next option from the list
65
            short_args_from_token(arg).each do |letter|
66
              next unless include?("-#{letter}")
67

68
              if arg_required?("-#{letter}")
69
                skip_next += 1
70
                param = args[idx + skip_next]
71
              end
72

73
              yield "-#{letter}", idx, param
74
            end
75
          elsif arg =~ LONG_FLAG && include?(arg)
76
            if arg_required?(arg)
77
              skip_next = 1
78
              param = args[idx + skip_next]
79
            end
80

81
            # Try to yield the short hand version of our argument if possible
82
            # This will result in less areas of code that would need to be changed
83
            to_return = short_arg_from_long_arg(arg)
84
            if to_return.nil?
85
              yield arg, idx, param
86
            else
87
              yield to_return, idx, param
88
            end
89
          else
90
            # else treat the passed in flag as argument
91
            yield nil, idx, arg
92
          end
93
        end
94
      end
95

96
      #
97
      # Returns usage information for this parsing context.
98
      #
99
      def usage
100
        txt = ["\nOPTIONS:\n"]
101

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

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

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

115
        txt << ""
116
        txt.join("\n")
117
      end
118

119
      def include?(search)
120
        fmt.keys.flatten.include?(search)
121
      end
122

123
      def arg_required?(opt)
124
        value = select_value_from_fmt_option(opt)
125
        return false if value.nil?
126

127
        value.first
128
      end
129

130
      def option_keys
131
        fmt.keys.flatten
132
      end
133

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

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

141
      private
142

143
      attr_accessor :fmt     # :nodoc:
144
      attr_accessor :longest # :nodoc:
145

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

150
        fmt_option[1]
151
      end
152

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

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

163
      # Parsing takes into account longest `short` flag found in 'arg' should as extracted first
164
      #
165
      # Returns Array of short arguments found in `arg`
166
      def short_args_from_token(arg)
167
        compare_arg = arg.dup[1..-1]
168
        short_args = []
169
        found_args = {}
170
        fmt.keys.each do |keys|
171
          if keys.first =~ SHORT_FLAG
172
            short_args << keys.first[1..-1]
173
          end
174
        end
175
        short_args.sort_by! { |value| value.downcase }.reverse!
176
        short_args.each do |short_arg|
177
          break if compare_arg.empty?
178
          if compare_arg.include? short_arg
179
            found_args[arg.index(short_arg)] = short_arg
180
            compare_arg.gsub!(short_arg, '')
181
          end
182
        end
183
        found_args.sort_by { |key, _value| key }.to_h.values
184
      end
185
    end
186
  end
187
end
188

189