require 'rex/ui'

module Rex
module Ui
module Text

###
#
# The shell class provides a command-prompt style interface in a
# generic fashion.
#
###
module Shell

	###
	#
	# This module is meant to be mixed into an input medium class instance as a
	# means of extending it to display a prompt before each call to gets.
	#
	###
	module InputShell
		attr_accessor :prompt, :output

		def pgets()

			output.print(prompt)
			output.flush

			output.prompting
			buf = gets
			output.prompting(false)

			buf
		end
	end

	#
	# Initializes a shell that has a prompt and can be interacted with.
	#
	def initialize(prompt, prompt_char = '>', histfile = nil, framework = nil)
		# Set the stop flag to false
		self.stop_flag      = false
		self.disable_output = false
		self.stop_count	    = 0

		# Initialize the prompt
		self.init_prompt = prompt
		self.prompt_char = prompt_char

		self.histfile = histfile
		self.hist_last_saved = 0

		self.framework = framework
	end

	def init_tab_complete
		if (self.input and self.input.supports_readline)
			self.input = Input::Readline.new(lambda { |str| tab_complete(str) })
			if Readline::HISTORY.length == 0 and histfile and File.exists?(histfile)
				File.readlines(histfile).each { |e|
					Readline::HISTORY << e.chomp
				}
				self.hist_last_saved = Readline::HISTORY.length
			end
			self.input.output = self.output
			update_prompt(input.prompt)
		end
	end

	#
	# Initializes the user interface input/output classes.
	#
	def init_ui(in_input = nil, in_output = nil)
		# Initialize the input and output methods
		self.input  = in_input
		self.output = in_output

		if (self.input)
			# Extend the input medium as an input shell if the input medium
			# isn't intrinsicly a shell.
			if (self.input.intrinsic_shell? == false)
				self.input.extend(InputShell)
			end

			self.input.output = self.output
		end
		update_prompt('')
	end

	#
	# Resets the user interface handles.
	#
	def reset_ui
		init_ui
	end

	#
	# Sets the log source that should be used for logging input and output.
	#
	def set_log_source(log_source)
		self.log_source = log_source
	end

	#
	# Unsets the log source so that logging becomes disabled.
	#
	def unset_log_source
		set_log_source(nil)
	end

	#
	# Performs tab completion on the supplied string.
	#
	def tab_complete(str)
		return tab_complete_proc(str) if (tab_complete_proc)
	end

	#
	# Run the command processing loop.
	#
	def run(&block)

		begin

			while true
				# If the stop flag was set or we've hit EOF, break out
				break if (self.stop_flag or self.stop_count > 1)

				init_tab_complete

				if framework
					if input.prompt.include?("%T")
						t = Time.now
						if framework.datastore['PromptTimeFormat']
							t = t.strftime(framework.datastore['PromptTimeFormat'])
						end
						input.prompt.gsub!(/%T/, t.to_s)
					end

					if input.prompt.include?("%H")
						hostname = ENV['HOSTNAME']
						if hostname.nil?
							hostname = `hostname`.split('.')[0]
						end

						# check if hostname is still nil
						if hostname.nil?
							hostname = ENV['COMPUTERNAME']
						end

						if hostname.nil?
							hostname = 'unknown'
						end

						input.prompt.gsub!(/%H/, hostname.chomp)
					end

					if input.prompt.include?("%U")
						user = ENV['USER']
						if user.nil?
							user = `whoami`
						end

						# check if username is still nil
						if user.nil?
							user = ENV['USERNAME']
						end

						if user.nil?
							user = 'unknown'
						end

						input.prompt.gsub!(/%U/, user.chomp)
					end

					input.prompt.gsub!(/%S/, framework.sessions.length.to_s)
					input.prompt.gsub!(/%J/, framework.jobs.length.to_s)
					input.prompt.gsub!(/%L/, Rex::Socket.source_address("50.50.50.50"))
					input.prompt.gsub!(/%D/, ::Dir.getwd)
					self.init_prompt = input.prompt
				end

				line = input.pgets()
				log_output(input.prompt)

				# If a block was passed in, pass the line to it.  If it returns true,
				# break out of the shell loop.
				if (block)
					break if (line == nil or block.call(line))
				elsif(input.eof? or line == nil)
				# If you have sessions active, this will give you a shot to exit gravefully
				# If you really are ambitious, 2 eofs will kick this out
					self.stop_count += 1
					next if(self.stop_count > 1)
					run_single("quit")
				else
				# Otherwise, call what should be an overriden instance method to
				# process the line.
					ret = run_single(line)
					# don't bother saving lines that couldn't be found as a
					# command, create the file if it doesn't exist
					if ret and self.histfile
						File.open(self.histfile, "a+") { |f|
							f.puts(line)
						}
					end
					self.stop_count = 0
				end

			end
		# Prevent accidental console quits
		rescue ::Interrupt
			output.print("Interrupt: use the 'exit' command to quit\n")
			retry
		end
	end

	#
	# Stop processing user input.
	#
	def stop
		self.stop_flag = true
	end

	#
	# Checks to see if the shell has stopped.
	#
	def stopped?
		self.stop_flag
	end

	#
	# Change the input prompt.
	#
	# prompt - the actual prompt
	# new_prompt_char the char to append to the prompt
	# mode - append or not to append - false = append true = make a new prompt
	def update_prompt(prompt = nil, new_prompt_char = nil, mode = false)
		if (self.input)
			if prompt
				new_prompt = self.init_prompt + ' ' + prompt + prompt_char + ' '
			else
				new_prompt = self.prompt || ''
			end

			if mode
			  new_prompt = prompt + (new_prompt_char || prompt_char) + ' '
			end

			# Save the prompt before any substitutions
			self.prompt = new_prompt

			# Set the actual prompt to the saved prompt with any substitutions
			# or updates from our output driver, be they color or whatever
			self.input.prompt = self.output.update_prompt(new_prompt)
			self.prompt_char  = new_prompt_char if (new_prompt_char)
		end
	end

	#
	# Output shortcuts
	#

	#
	# Prints an error message to the output handle.
	#
	def print_error(msg='')
		return if (output.nil?)

		self.on_print_proc.call(msg) if self.on_print_proc
		# Errors are not subject to disabled output
		log_output(output.print_error(msg))
	end

	#
	# Prints a status message to the output handle.
	#
	def print_status(msg='')
		return if (disable_output == true)

		self.on_print_proc.call(msg) if self.on_print_proc
		log_output(output.print_status(msg))
	end

	#
	# Prints a good message to the output handle.
	#
	def print_good(msg='')
		return if (disable_output == true)

		self.on_print_proc.call(msg) if self.on_print_proc
		log_output(output.print_good(msg))
	end

	#
	# Prints a line of text to the output handle.
	#
	def print_line(msg='')
		return if (disable_output == true)

		self.on_print_proc.call(msg) if self.on_print_proc
		log_output(output.print_line(msg))
	end

	#
	# Prints a raw message to the output handle.
	#
	def print(msg='')
		return if (disable_output == true)
		self.on_print_proc.call(msg) if self.on_print_proc
		log_output(output.print(msg))
	end

	#
	# Whether or not output has been disabled.
	#
	attr_accessor :disable_output
	#
	# The input handle to read user input from.
	#
	attr_reader   :input
	#
	# The output handle to write output to.
	#
	attr_reader   :output

	attr_accessor :on_command_proc
	attr_accessor :on_print_proc
	attr_accessor :framework

protected

	#
	# Parse a line into an array of arguments.
	#
	def parse_line(line)
		log_input(line)

		line.gsub!(/(\r|\n)/, '')

		begin
			return args = Rex::Parser::Arguments.from_s(line)
		rescue ::ArgumentError
			print_error("Parse error: #{$!}")
		end

		return []
	end

	#
	# Print the prompt, but do not log it.
	#
	def _print_prompt(prompt)
		output.print(prompt)
	end

	#
	# Writes the supplied input to the log source if one has been registered.
	#
	def log_input(buf)
		rlog(buf, log_source) if (log_source)
	end

	#
	# Writes the supplied output to the log source if one has been registered.
	#
	def log_output(buf)
		rlog(buf, log_source) if (log_source)
	end

	attr_writer   :input, :output # :nodoc:
	attr_accessor :stop_flag, :init_prompt # :nodoc:
	attr_accessor :prompt # :nodoc:
	attr_accessor :prompt_char, :tab_complete_proc # :nodoc:
	attr_accessor :histfile # :nodoc:
	attr_accessor :hist_last_saved # the number of history lines when last saved/loaded
	attr_accessor :log_source, :stop_count # :nodoc:

end

###
#
# Pseudo-shell interface that simply includes the Shell mixin.
#
###
class PseudoShell
	include Shell
end


end end end