module Msf

# this module provides instance methods to be used in overloading to do single byte sending of data
module SmallSend
	def write(buf, opts = {})
		warn "smallsend write"
		
		tsent = 0; bidx  = 0

		if self._send_size == nil or self._send_size == 0
			self._send_size = 1
		end
	
		while (bidx < buf.length)
			str = buf[bidx, _send_size]
			sent = super(str, opts)
			bidx += sent if sent > 0
			tsent += sent
			
			if self.is_a?(SlowSend)
				sleep(self._send_delay)
			else
				sleep(0)
			end
		end

		return tsent
	end

	attr_accessor :_send_size
end

# this module provides instance mehtods to be used in overloading of Socket to insert delays inbetween each write
module SlowSend
	def write(buf, opts = {})
		warn 'slowsend write'
		if !self.is_a?(SmallSend)
			sleep(_send_delay)
		end
		response = super(buf, opts)
		return response
	end

	attr_accessor :_send_delay
end

###
#
# This module provides methods for establish a connection to a remote host and
# communicating with it.
#
###
module Exploit::Remote::Tcp

	#
	# Initializes an instance of an exploit module that exploits a
	# vulnerability in a TCP server.
	#
	def initialize(info = {})
		super

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

		register_advanced_options(
			[
				Opt::SSL,
				Opt::Proxies,
				Opt::CPORT,
				Opt::CHOST,
				OptInt.new('ConnectTimeout', [ true, 'Maximum number of seconds to establish a TCP connection', 10])
			], Msf::Exploit::Remote::Tcp
		)
		
		register_evasion_options(
			[
				OptInt.new('TCP::max_send_size', [false, 'Maxiumum tcp segment size.  (0 = disable)', 0]),
				OptInt.new('TCP::send_delay', [false, 'Delays inserted before every send.  (0 = disable)', 0])
			], Msf::Exploit::Remote::Tcp
		)
	end

	#
	# Establishes a TCP connection to the specified RHOST/RPORT
	#
	def connect(global = true, opts={})
		nsock = Rex::Socket::Tcp.create(
			'PeerHost'  =>  opts['RHOST'] || rhost,
			'PeerPort'  => (opts['RPORT'] || rport).to_i,
			'LocalHost' =>  opts['CHOST'] || chost || "0.0.0.0",
			'LocalPort' => (opts['CPORT'] || cport || 0).to_i,
			'SSL'       => ssl,
			'Proxies'   => proxies,
			'Timeout'   => (opts['ConnectTimeout'] || connect_timeout || 10).to_i,
			'Context'   =>
				{
					'Msf'        => framework,
					'MsfExploit' => self,
				})

		# enable evasions on this socket
		# XXX implement evasions!!!!
		# evasions(nsock)

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

		# Add this socket to the list of sockets created by this exploit
		sockets << nsock

		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)
			if (nsock == self.sock)
				self.sock = nil
			end

			# Remove this socket from the list of sockets so that it will not be
			# aborted.
			sockets.delete(nsock)
		end

		return rv
	end

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

		if (nsock == sock)
			self.sock = nil
		end
	
		# Remove this socket from the list of sockets created by this exploit
		sockets.delete(nsock)
	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

	#
	# Returns the local host for outgoing connections
	#
	def chost
		datastore['CHOST']
	end

	#
	# Returns the local port for outgoing connections
	#
	def cport
		datastore['CPORT']
	end

	#
	# Returns the boolean indicating SSL
	#
	def ssl
		datastore['SSL']
	end

	#
	# Returns the proxy configuration
	#	
	def proxies
		datastore['Proxies']
	end	

	#
	# Returns the TCP connection timeout
	#	
	def connect_timeout
		datastore['ConnectTimeout']
	end
	
protected

	attr_accessor :sock

end


###
#
# 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.", Rex::Socket.source_address ]),
				OptPort.new('SRVPORT',    [ true, "The local port to listen on.", 8080 ])
			], Msf::Exploit::Remote::TcpServer)
		
		register_evasion_options(
			[
				OptInt.new('TCP::max_send_size', [false, 'Maxiumum tcp segment size.  (0 = disable)', 0]),
				OptInt.new('TCP::send_delay', [false, 'Delays inserted before every send.  (0 = disable)', 0])
			], Msf::Exploit::Remote::Tcp
		)
	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

	#
	# Called when a client connects.
	#
	def on_client_connect(client)
	end

	#
	# Called when a client has data available for reading.
	#
	def on_client_data(client)
	end

	#
	# Called when a client has disconnected.
	#
	def on_client_close(client)
	end

	#
	# Starts the service.
	#
	def start_service(*args)
		self.service = Rex::Socket::TcpServer.create(
			'LocalHost' => srvhost,
			'LocalPort' => srvport,
			'Context'   =>
				{
					'Msf'        => framework,
					'MsfExploit' => self,
				})

		self.service.on_client_connect_proc = Proc.new { |client|
			on_client_connect(client)
		}
		self.service.on_client_data_proc = Proc.new { |client|
			on_client_data(client)
		}
		self.service.on_client_close_proc = Proc.new { |client|
			on_client_close(client)
		}

		self.service.start
	end

	#
	# Stops the service.
	#
	def stop_service
		if (service)
			self.service.deref if self.service.kind_of?(Rex::Service)
			self.service.close if self.service.kind_of?(Rex::Socket)
			self.service = nil
		end
	end

	# Enable evasions on a given client
	def evasions(socket)
		# XXX - oooogly
	    return if socket.instance_variables.member?('@tcp_evasion')

		if !socket.is_a?(SmallSend) and datastore['TCP::max_send_size'] > 0
			socket.extend(SmallSend)
			socket._send_size = datastore['TCP::max_send_size']
		end

		if !socket.is_a?(SlowSend) and datastore['TCP::send_delay'] > 0
			socket.extend(SlowSend)
			socket._send_delay = datastore['TCP::send_delay']
		end
		socket.instance_eval('@tcp_evasion = 1')
	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 # :nodoc:

end

end