dispatcher_shell.rb in librex-0.0.7

- old
+ new

@@ -1,6 +1,7 @@
 require 'rex/ui'
+require 'pp'
 
 module Rex
 module Ui
 module Text
 
@@ -78,23 +79,85 @@
 		def update_prompt(prompt=nil)
 			shell.update_prompt(prompt)
 		end
 
 		#
+		# Displays the help banner.  With no arguments, this is just a list of
+		# all commands grouped by dispatcher.  Otherwise, tries to use a method
+		# named cmd_#{+cmd+}_help for the first dispatcher that has a command
+		# named +cmd+.
+		#
+		def cmd_help(cmd=nil, *ignored)
+			if cmd
+				help_found = false
+				cmd_found = false
+				shell.dispatcher_stack.each do |dispatcher|
+					next unless dispatcher.respond_to?(:commands)
+					next if (dispatcher.commands.nil?)
+					next if (dispatcher.commands.length == 0)
+
+					if dispatcher.respond_to?("cmd_#{cmd}")
+						cmd_found = true
+						break unless dispatcher.respond_to? "cmd_#{cmd}_help"
+						dispatcher.send("cmd_#{cmd}_help")
+						help_found = true
+						break
+					end
+				end
+				print_error("No help for #{cmd}, try -h") if cmd_found and not help_found
+				print_error("No such command") if not cmd_found
+			else
+				print(shell.help_to_s)
+			end
+		end
+
+		#
+		# Tab completion for the help command
+		#
+		# By default just returns a list of all commands in all dispatchers.
+		#
+		def cmd_help_tabs(str, words)
+			return [] if words.length > 1
+
+			tabs = []
+			shell.dispatcher_stack.each { |dispatcher|
+				tabs += dispatcher.commands.keys
+			}
+			return tabs
+		end
+
+		alias cmd_? cmd_help
+
+
+		#
 		# No tab completion items by default
 		#
 		attr_accessor :shell, :tab_complete_items
 
+		#
+		# Provide a generic tab completion for file names.
+		#
+		# If the only completion is a directory, this descends into that directory
+		# and continues completions with filenames contained within.
+		#
+		def tab_complete_filenames(str, words)
+			matches = ::Readline::FILENAME_COMPLETION_PROC.call(str)
+			if matches and matches.length == 1 and File.directory?(matches[0])
+				dir = matches[0]
+				dir += File::SEPARATOR if dir[-1,1] != File::SEPARATOR
+				matches = ::Readline::FILENAME_COMPLETION_PROC.call(dir) 
+			end
+			matches
+		end
+
 	end
 
 	#
 	# DispatcherShell derives from shell.
 	#
 	include Shell
 
-	attr_accessor :on_command_proc
-
 	#
 	# Initialize the dispatcher shell.
 	#
 	def initialize(prompt, prompt_char = '>', histfile = nil)
 		super
@@ -146,24 +209,26 @@
 		# Enumerate each entry in the dispatcher stack
 		dispatcher_stack.each { |dispatcher|
 		
 			# If no command is set and it supports commands, add them all
 			if (tab_words.empty? and dispatcher.respond_to?('commands'))
-				items.concat(dispatcher.commands.to_a.map { |x| x[0] })
+				items.concat(dispatcher.commands.keys)
 			end
 
 			# If the dispatcher exports a tab completion function, use it
 			if(dispatcher.respond_to?('tab_complete_helper'))
 				res = dispatcher.tab_complete_helper(str, tab_words)
+			else
+				res = tab_complete_helper(dispatcher, str, tab_words)
+			end
 
-				if (res.nil?)
-					# A nil response indicates no optional arguments
-					return [''] if items.empty?
-				else
-					# Otherwise we add the completion items to the list
-					items.concat(res)
-				end
+			if (res.nil?)
+				# A nil response indicates no optional arguments
+				return [''] if items.empty?
+			else
+				# Otherwise we add the completion items to the list
+				items.concat(res)
 			end
 		}
 
 		# Verify that our search string is a valid regex
 		begin
@@ -180,9 +245,29 @@
 			e =~ /^#{str}/
 		# Prepend the rest of the command (or it gets replaced!)
 		}.map { |e| 
 			tab_words.dup.push(e).join(' ')
 		}
+	end
+
+	#
+	# Provide command-specific tab completion
+	#
+	def tab_complete_helper(dispatcher, str, words)
+		items = []
+
+		tabs_meth = "cmd_#{words[0]}_tabs"
+		# Is the user trying to tab complete one of our commands?
+		if (dispatcher.commands.include?(words[0]) and dispatcher.respond_to?(tabs_meth))
+			res = dispatcher.send(tabs_meth, str, words)
+			return [] if res.nil?
+			items.concat(res)
+		else
+			# Avoid the default completion list for known commands
+			return []
+		end
+
+		return items
 	end
 
 	#
 	# Run a single command line.
 	#