require 'rex/ui'
require 'rex/post/meterpreter'
require 'rex/logging'

module Rex
module Post
module Meterpreter
module Ui

###
#
# This class provides a shell driven interface to the meterpreter client API.
#
###
class Console

	include Rex::Ui::Text::DispatcherShell

	# Dispatchers
	require 'rex/post/meterpreter/ui/console/interactive_channel'
	require 'rex/post/meterpreter/ui/console/command_dispatcher'
	require 'rex/post/meterpreter/ui/console/command_dispatcher/core'

	#
	# Initialize the meterpreter console.
	#
	def initialize(client)
		if (Rex::Compat.is_windows())
			super("meterpreter")
		else
			super("%undmeterpreter%clr")
		end

		# The meterpreter client context
		self.client = client

		# Queued commands array
		self.commands = []

		# Point the input/output handles elsewhere
		reset_ui

		enstack_dispatcher(Console::CommandDispatcher::Core)

		# Set up logging to whatever logsink 'core' is using
		if ! $dispatcher['meterpreter']
			$dispatcher['meterpreter'] = $dispatcher['core']
		end
	end

	#
	# Called when someone wants to interact with the meterpreter client.  It's
	# assumed that init_ui has been called prior.
	#
	def interact(&block)
		init_tab_complete

		# Run queued commands
		commands.delete_if { |ent|
			run_single(ent)
			true
		}

		# Run the interactive loop
		run { |line|
			# Run the command
			run_single(line)

			# If a block was supplied, call it, otherwise return false
			if (block)
				block.call
			else
				false
			end
		}
	end

	#
	# Interacts with the supplied channel.
	#
	def interact_with_channel(channel)
		channel.extend(InteractiveChannel) unless (channel.kind_of?(InteractiveChannel) == true)
		channel.on_command_proc = self.on_command_proc if self.on_command_proc
		channel.on_print_proc   = self.on_print_proc if self.on_print_proc

		channel.interact(input, output)
		channel.reset_ui
	end

	#
	# Queues a command to be run when the interactive loop is entered.
	#
	def queue_cmd(cmd)
		self.commands << cmd
	end

	#
	# Runs the specified command wrapper in something to catch meterpreter
	# exceptions.
	#
	def run_command(dispatcher, method, arguments)
		begin
			super
		rescue Timeout::Error
			log_error("Operation timed out.")
		rescue RequestError => info
			log_error(info.to_s)
		rescue ::Errno::EPIPE, ::OpenSSL::SSL::SSLError, ::IOError
			self.client.kill
		rescue  ::Exception => e
			log_error("Error running command #{method}: #{e.class} #{e}")
		end
	end

	#
	# Logs that an error occurred and persists the callstack.
	#
	def log_error(msg)
		print_error(msg)

		elog(msg, 'meterpreter')

		dlog("Call stack:\n#{$@.join("\n")}", 'meterpreter')
	end

	attr_reader :client # :nodoc:

protected

	attr_writer :client # :nodoc:
	attr_accessor :commands # :nodoc:

end

end
end
end
end