module Msf

###
#
# This module provides methods for establish a connection to a remote host and
# communicating with it.
#
###
module Exploit::Remote::Tcp
	
	def initialize(info = {})
		super

		register_options(
			[
				Opt::RHOST,
				Opt::RPORT,
				Opt::SSL
			], Msf::Exploit::Remote::Tcp)
	end

	#
	# Establishes a TCP connection to the specified RHOST/RPORT
	#
	def connect(global = true)
		nsock = Rex::Socket::Tcp.create(
			'PeerHost'  => datastore['RHOST'],
			'PeerPort'  => datastore['RPORT'].to_i,
			'LocalHost' => datastore['CHOST'] || "0.0.0.0",
			'LocalPort' => datastore['CPORT'] ? datastore['CPORT'].to_i : 0,
			'SSL'       => datastore['SSL'])

		# Set this socket to the global socket as necessary
		self.sock = nsock if (global)

		return nsock
	end

	def handler(nsock = self.sock)
		# If the handler claims the socket, then we don't want it to get closed
		# during cleanup
		if (((rv = super) == Handler::Claimed) and
		    (nsock == self.sock))
			self.sock = nil
		end

		return rv
	end

	#
	# Closes the TCP connection
	#
	def disconnect(nsock = self.sock)
		if (nsock)
			nsock.shutdown
			nsock.close
		end

		if (nsock == sock)
			self.sock = nil
		end
	end

	#
	# Performs cleanup, disconnects the socket if necessary
	#
	def cleanup
		super

		disconnect
	end

	##
	#
	# Wrappers for getters
	#
	##

	#
	# Returns the target host
	#
	def rhost
		datastore['RHOST']
	end

	#
	# Returns the remote port
	#
	def rport
		datastore['RPORT']
	end

	#
	# Returns the local host
	#
	def lhost
		datastore['LHOST']
	end

	#
	# Returns the local port
	#
	def lport
		datastore['LPORT']
	end

protected

	attr_accessor :sock

end


###
#
# TcpServer
# ---------
#
# This mixin provides a generic interface for running a TCP server of some
# sort that is designed to exploit clients.  Exploits that include this mixin
# automatically take a passive stance.
#
###
module Exploit::Remote::TcpServer

	def initialize(info = {})
		super(update_info(info,
			'Stance' => Msf::Exploit::Stance::Passive))

		register_options(
			[
				OptAddress.new('SRVHOST', [ true, "The local host to listen on."       ]),
				OptPort.new('SRVPORT',    [ true, "The local port to listen on.", 8080 ])
			], Msf::Exploit::Remote::TcpServer)
	end

	#
	# This mixin overrides the exploit method so that it can initiate the
	# service that corresponds with what the client has requested.
	#
	def exploit
		start_service

		print_status("Server started.");
	end

	#
	# Stops the service, if one was created.
	#
	def cleanup
		stop_service

		print_status("Server stopped.");
	end

	#
	# Starts the service.
	#
	def start_service(*args)
	end

	#
	# Stops the service.
	#
	def stop_service
		if (service)
			Rex::ServiceManager.stop_service(self.service)

			self.service.deref
			self.service = nil
		end
	end

	#
	# Returns the local host that is being listened on.
	#
	def srvhost
		datastore['SRVHOST']
	end

	#
	# Returns the local port that is being listened on.
	#
	def srvport
		datastore['SRVPORT']
	end

protected
	
	attr_accessor :service

end

end