require 'rex/sync/thread_safe'

module Rex
module IO

###
#
# This mixin is an abstract representation of a streaming connection.  Streams
# extend classes that must implement the following methods:
#
#   syswrite(buffer)
#   sysread(length)
#   shutdown(how)
#   close
#   peerinfo
#   localinfo
#
###
module Stream

	##
	#
	# Abstract methods
	#
	##

	#
	# This method writes the supplied buffer to the stream.  This method
	# intelligent reduces the size of supplied buffers so that ruby doesn't get
	# into a potential global thread blocking state when used on blocking
	# sockets.  That is, this method will send the supplied buffer in chunks
	# of, at most, 32768 bytes.
	#
	def write(buf, opts = {})
		tsent = 0
		bidx  = 0
		
		begin
			while (bidx < buf.length)
				sent  = fd.syswrite(buf[bidx, 32768])
				bidx  += sent if sent > 0
				tsent += sent
			end
		rescue IOError
			return nil if (fd.abortive_close == true)

			raise $!
		end

		tsent
	end

	#
	# This method reads data of the supplied length from the stream.
	#
	def read(length = nil, opts = {})
		# XXX handle length being nil
		begin
			fd.sysread(length)
		rescue IOError
			return nil if (fd.abortive_close == true)
			raise $!
		end
	end

	#
	# Polls the stream to see if there is any read data available.  Returns
	# true if data is available for reading, otherwise false is returned.
	#
	def has_read_data?(timeout = nil)
		begin
			if ((rv = Rex::ThreadSafe.select([ fd ], nil, nil, timeout)) and
			    (rv[0]) and
			    (rv[0][0] == fd))
				true
			else
				false
			end
		rescue StreamClosedError, IOError
			# If the thing that lead to the closure was an abortive close, then
			# don't raise the stream closed error.
			return false if (fd.abortive_close == true)

			raise $!
		end
	end

	#
	# This method returns the selectable file descriptor, or self by default.
	#
	def fd
		self
	end

	##
	#
	# Common methods
	#
	##

	#
	# This method writes the supplied buffer to the stream by calling the write
	# routine.
	#
	def <<(buf)
		return write(buf.to_s)
	end

	#
	# This method calls get_once() to read pending data from the socket
	#
	def >>
		get_once
	end
	
	#
	# This method writes to the stream, optionally timing out after a period of
	# time.
	#
	def timed_write(buf, wait = def_write_timeout, opts = {})
		if (wait and wait > 0)
			timeout(wait) {
				return write(buf, opts)
			}
		else
			return write(buf, opts)
		end
	end

	#
	# This method reads from the stream, optionally timing out after a period
	# of time.
	#
	def timed_read(length = nil, wait = def_read_timeout, opts = {})
		if (wait and wait > 0)
			timeout(wait) {
				return read(length, opts)
			}
		else
			return read(length, opts)
		end
	end

	#
	# This method writes the full contents of the supplied buffer, optionally
	# with a timeout.
	#
	def put(buf, opts = {})
		return 0 if (buf == nil or buf.length == 0)

		send_len = buf.length
		send_idx = 0
		wait     = opts['Timeout'] || 0

		# Keep writing until our send length drops to zero
		while (send_idx < send_len)
			curr_len  = timed_write(buf[send_idx, buf.length-send_idx], wait, opts)

			# If the write operation failed due to an IOError, then we fail.
			return buf.length - send_len if (curr_len == nil)

			send_len -= curr_len
			send_idx += curr_len
		end

		return buf.length - send_len
	end


	# 
	# This method emulates the behavior of Pex::Socket::Recv in MSF2
	#
	def get_once(length = -1, timeout = def_read_timeout)

		if (has_read_data?(timeout) == false)
			return nil
		end
		
		bsize = (length == -1) ? def_block_size : length

		begin
			return read(bsize)
		rescue Exception
		end
		
		return ''
	end

	#
	# This method reads as much data as it can from the wire given a maximum
	# timeout.
	#
	def get(timeout = nil, ltimeout = def_read_loop_timeout, opts = {})
		# For those people who are used to being able to use a negative timeout!
		if (timeout and timeout.to_i < 0)
			timeout = nil
		end

		# No data in the first place? bust.
		if (has_read_data?(timeout) == false)
			return nil
		end

		buf = ""
		lps = 0
		eof = false

		# Keep looping until there is no more data to be gotten..
		while (has_read_data?(ltimeout) == true)
			# Catch EOF errors so that we can handle them properly.
			begin
				temp = read(def_block_size)
			rescue EOFError
				eof = true
			end
		
			# If we read zero bytes and we had data, then we've hit EOF
			if (temp and temp.length == 0)
				eof = true
			end

			# If we reached EOF and there are no bytes in the buffer we've been
			# reading into, then throw an EOF error.
			if (eof)
				# If we've already read at least some data, then it's time to
				# break out and let it be processed before throwing an EOFError.
				if (buf.length > 0)
					break
				else
					raise EOFError
				end
			end

			break if (temp == nil or temp.empty? == true)

			buf += temp
			lps += 1
			
			break if (lps >= def_max_loops)
		end

		# Return the entire buffer we read in
		return buf
	end

	##
	#
	# Defaults
	#
	##

	#
	# The default number of seconds to wait for a write operation to timeout.
	#
	def def_write_timeout
		10
	end

	#
	# The default number of seconds to wait for a read operation to timeout.
	#
	def def_read_timeout
		10
	end

	#
	# The default number of seconds to wait while in a read loop after read
	# data has been found.
	#
	def def_read_loop_timeout
		0.1
	end

	#
	# The maximum number of read loops to perform before returning to the
	# caller.
	#
	def def_max_loops
		1024
	end

	#
	# The default block size to read in chunks from the wire.
	#
	def def_block_size
		16384
	end

	#
	# This flag indicates whether or not an abortive close has been issued.
	#
	attr_accessor :abortive_close

protected

end

end end