Land rapid7#9201, enhanced tab completion

Author: wvu

lib/msf/ui/console/command_dispatcher/core.rb

@@ -2198,16 +2198,7 @@ def option_values_dispatch(o, str, words)
         if rh and not rh.empty?
           res << Rex::Socket.source_address(rh)
         else
-          res << Rex::Socket.source_address
-          # getifaddrs was introduced in 2.1.2
-          if Socket.respond_to?(:getifaddrs)
-            ifaddrs = Socket.getifaddrs.find_all do |ifaddr|
-              ((ifaddr.flags & Socket::IFF_LOOPBACK) == 0) &&
-                ifaddr.addr &&
-                ifaddr.addr.ip?
-            end
-            res += ifaddrs.map { |ifaddr| ifaddr.addr.ip_address }
-          end
+          res += tab_complete_source_address
         end
       else
       end

lib/msf/ui/console/command_dispatcher/exploit.rb

@@ -164,6 +164,21 @@ def cmd_exploit(*args)
     end
   end
 
+  def cmd_exploit_tabs(str, words)
+    fmt = {
+      '-e' => [ framework.encoders.map { |refname, mod| refname } ],
+      '-f' => [ nil                                               ],
+      '-h' => [ nil                                               ],
+      '-j' => [ nil                                               ],
+      '-n' => [ framework.nops.map { |refname, mod| refname }     ],
+      '-o' => [ true                                              ],
+      '-p' => [ framework.payloads.map { |refname, mod| refname } ],
+      '-t' => [ true                                              ],
+      '-z' => [ nil                                               ]
+    }
+    tab_complete_generic(fmt, str, words)
+  end
+
   alias cmd_run cmd_exploit
 
   def cmd_exploit_help

lib/msf/ui/console/command_dispatcher/jobs.rb

@@ -341,6 +341,19 @@ def cmd_handler(*args)
 
             print_status "Payload handler running as background job #{job_id}."
           end
+
+          def cmd_handler_tabs(str, words)
+            fmt = {
+              '-h' => [ nil                                               ],
+              '-x' => [ nil                                               ],
+              '-p' => [ framework.payloads.map { |refname, mod| refname } ],
+              '-P' => [ true                                              ],
+              '-H' => [ :address                                          ],
+              '-e' => [ framework.encoders.map { |refname, mod| refname } ],
+              '-n' => [ true                                              ]
+            }
+            tab_complete_generic(fmt, str, words)
+          end
         end
       end
     end

lib/msf/ui/console/command_dispatcher/payload.rb

@@ -13,7 +13,7 @@ class Payload
           include Msf::Ui::Console::ModuleCommandDispatcher
 
           # Load supported formats
-          supported_formats = \
+          @@supported_formats = \
             Msf::Simple::Buffer.transform_formats + \
             Msf::Util::EXE.to_executable_fmt_formats
 
@@ -25,7 +25,7 @@ class Payload
             "-o" => [ true,  "A comma separated list of options in VAR=VAL format." ],
             "-s" => [ true,  "NOP sled length."                                     ],
             "-f" => [ true,  "The output file name (otherwise stdout)"              ],
-            "-t" => [ true,  "The output format: #{supported_formats.join(',')}"    ],
+            "-t" => [ true,  "The output format: #{@@supported_formats.join(',')}"    ],
             "-p" => [ true,  "The Platform for output."                             ],
             "-k" => [ false, "Keep the template executable functional"              ],
             "-x" => [ true,  "The executable template to use"                       ],
@@ -151,6 +151,24 @@ def cmd_generate(*args)
             end
             true
           end
+
+          def cmd_generate_tabs(str, words)
+            fmt = {
+              '-b' => [ true                                              ],
+              '-E' => [ nil                                               ],
+              '-e' => [ framework.encoders.map { |refname, mod| refname } ],
+              '-h' => [ nil                                               ],
+              '-o' => [ true                                              ],
+              '-s' => [ true                                              ],
+              '-f' => [ :file                                             ],
+              '-t' => [ @@supported_formats                               ],
+              '-p' => [ true                                              ],
+              '-k' => [ nil                                               ],
+              '-x' => [ :file                                             ],
+              '-i' => [ true                                              ]
+            }
+            tab_complete_generic(fmt, str, words)
+          end
         end
       end
     end

lib/rex/ui/text/dispatcher_shell.rb

@@ -249,6 +249,56 @@ def tab_complete_filenames(str, words)
       matches
     end
 
+    #
+    # Provide a generic tab completion function based on the specification
+    # pass as fmt. The fmt argument in a hash where values are an array
+    # defining how the command should be completed. The first element of the
+    # array can be one of:
+    #   nil      - This argument is a flag and takes no option.
+    #   true     - This argument takes an option with no suggestions.
+    #   :address - This option is a source address.
+    #   :bool    - This option is a boolean.
+    #   :file    - This option is a file path.
+    #   Array    - This option is an array of possible values.
+    #
+    def tab_complete_generic(fmt, str, words)
+      last_word = words[-1]
+      fmt = fmt.select { |key, value| last_word == key || !words.include?(key) }
+
+      val = fmt[last_word]
+      return fmt.keys if !val  # the last word does not look like a fmtspec
+      arg = val[0]
+      return fmt.keys if !arg  # the last word is a fmtspec that takes no argument
+
+      tabs = []
+      if arg.to_s.to_sym == :address
+        tabs = tab_complete_source_address
+      elsif arg.to_s.to_sym == :bool
+        tabs = ['true', 'false']
+      elsif arg.to_s.to_sym == :file
+        tabs = tab_complete_filenames(str, words)
+      elsif arg.kind_of?(Array)
+        tabs = arg.map {|a| a.to_s}
+      end
+      tabs
+    end
+
+    #
+    # Return a list of possible source addresses for tab completion.
+    #
+    def tab_complete_source_address
+      addresses = [Rex::Socket.source_address]
+      # getifaddrs was introduced in 2.1.2
+      if Socket.respond_to?(:getifaddrs)
+        ifaddrs = Socket.getifaddrs.find_all do |ifaddr|
+          ((ifaddr.flags & Socket::IFF_LOOPBACK) == 0) &&
+            ifaddr.addr &&
+            ifaddr.addr.ip?
+        end
+        addresses += ifaddrs.map { |ifaddr| ifaddr.addr.ip_address }
+      end
+      addresses
+    end
   end
 
   #