require 'rex/proto/http'

module Rex
module Proto
module Http

DefaultProtocol = '1.1'

###
#
# This class represents an HTTP packet.
#
###
class Packet

	#
	# Parser processing codes
	#
	module ParseCode
		Completed = 1
		Partial   = 2
		Error     = 3
	end

	#
	# Parser states
	#
	module ParseState
		ProcessingHeader = 1
		ProcessingBody   = 2
		Completed        = 3
	end

	require 'rex/proto/http/header'

	#
	# Initializes an instance of an HTTP packet.
	#
	def initialize()
		self.headers = Header.new
		self.auto_cl = true

		reset
	end

	#
	# Return the associated header value, if any.
	#
	def [](key)
		if (self.headers.include?(key))
			return self.headers[key]
		end

		self.headers.each_pair do |k,v|
			if (k.downcase == key.downcase)
				return v
			end
		end

		return nil
	end

	#
	# Set the associated header value.
	#
	def []=(key, value)
		self.headers[key] = value
	end

	#
	# Parses the supplied buffer.  Returns one of the two parser processing
	# codes (Completed, Partial, or Error).
	#
	def parse(buf)

		# Append the incoming buffer to the buffer queue.
		self.bufq += buf.to_s

		begin

			# Process the header
			if(self.state == ParseState::ProcessingHeader)
				parse_header
			end

			# Continue on to the body if the header was processed
			if(self.state == ParseState::ProcessingBody)
				if (self.body_bytes_left == 0)
					self.state = ParseState::Completed
				else
					parse_body
				end
			end
		rescue
			# XXX: BUG: This rescue might be a problem because it will swallow TimeoutError
			self.error = $!
			return ParseCode::Error
		end

		# Return completed or partial to the parsing status to the caller
		(self.state == ParseState::Completed) ? ParseCode::Completed : ParseCode::Partial
	end

	#
	# Reset the parsing state and buffers.
	#
	def reset
		self.state = ParseState::ProcessingHeader
		self.transfer_chunked = false
		self.inside_chunk     = false
		self.headers.reset
		self.bufq  = ''
		self.body  = ''
	end

	#
	# Reset the parsing state but leave the buffers.
	#
	def reset_except_queue
		self.state = ParseState::ProcessingHeader
		self.transfer_chunked = false
		self.inside_chunk     = false
		self.headers.reset
		self.body  = ''
	end

	#
	# Returns whether or not parsing has completed.
	#
	def completed?

		return true if self.state == ParseState::Completed

		# If the parser state is processing the body and there are an
		# undetermined number of bytes left to read, we just need to say that
		# things are completed as it's hard to tell whether or not they really
		# are.
		if (self.state == ParseState::ProcessingBody and self.body_bytes_left < 0)
			return true
		end

		false
	end

	#
	# Build a 'Transfer-Encoding: chunked' payload with random chunk sizes
	#
	def chunk(str, min_size = 1, max_size = 1000)
		chunked = ''

		# min chunk size is 1 byte
		if (min_size < 1); min_size = 1; end

		# don't be dumb
		if (max_size < min_size); max_size = min_size; end

		while (str.size > 0)
			chunk = str.slice!(0, rand(max_size - min_size) + min_size)
			chunked += sprintf("%x", chunk.size) + "\r\n" + chunk + "\r\n"
		end
		chunked += "0\r\n\r\n"
	end

	#
	# Converts the packet to a string.
	#
	def to_s
		# Duplicate and make sure this is 8BIT safe for Ruby 1.9
		content = self.body.unpack("C*").pack("C*")
		
		# Update the content length field in the header with the body length.
		if (content)
			if !self.compress.nil?
				case self.compress
					when 'gzip'
						self.headers['Content-Encoding'] = 'gzip'
						content = Rex::Text.gzip(content)
					when 'deflate'
						self.headers['Content-Encoding'] = 'deflate'
						content = Rex::Text.zlib_deflate(content)
					when 'none'
					# this one is fine...
					# when 'compress'
					else
						raise RuntimeError, 'Invalid Content-Encoding'
				end
			end

			if (self.auto_cl == true && self.transfer_chunked == true)
				raise RuntimeError, "'Content-Length' and 'Transfer-Encoding: chunked' are incompatible"
			elsif self.auto_cl == true
				self.headers['Content-Length'] = content.length
			elsif self.transfer_chunked == true
				if self.proto != '1.1'
					raise RuntimeError, 'Chunked encoding is only available via 1.1'
				end
				self.headers['Transfer-Encoding'] = 'chunked'
				content = self.chunk(content, self.chunk_min_size, self.chunk_max_size)
			end
		end

		str  = self.headers.to_s(cmd_string)
		str += content || ''
	end

	#
	# Converts the packet from a string.
	#
	def from_s(str)
		reset
		parse(str)
	end

	#
	# Returns the command string, such as:
	#
	# HTTP/1.0 200 OK for a response
	#
	# or
	#
	# GET /foo HTTP/1.0 for a request
	#
	def cmd_string
		self.headers.cmd_string
	end

	attr_accessor :headers
	attr_accessor :error
	attr_accessor :state
	attr_accessor :bufq
	attr_accessor :body
	attr_accessor :auto_cl
	attr_accessor :max_data
	attr_accessor :transfer_chunked
	attr_accessor :compress
	attr_reader   :incomplete

	attr_accessor :chunk_min_size
	attr_accessor :chunk_max_size

protected

	attr_writer   :incomplete
	attr_accessor :body_bytes_left
	attr_accessor :inside_chunk
	attr_accessor :keepalive

	##
	#
	# Overridable methods
	#
	##

	#
	# Allows derived classes to split apart the command string.
	#
	def update_cmd_parts(str)
	end

	##
	#
	# Parsing
	#
	##

	def parse_header

		head,data = self.bufq.split(/\r?\n\r?\n/, 2)
		return if not data

		self.headers.from_s(head)
		self.bufq = data || ""

		# Set the content-length to -1 as a placeholder (read until EOF)
		self.body_bytes_left = -1

		# Extract the content length if it was specified
		if (self.headers['Content-Length'])
			self.body_bytes_left = self.headers['Content-Length'].to_i
		end

		# Look for a chunked transfer header
		if (self.headers['Transfer-Encoding'].to_s.downcase == 'chunked')
			self.transfer_chunked = true
			self.auto_cl = false
		end

		# Determine how to handle data when there is no length header
		if (self.body_bytes_left == -1)
			if (not self.transfer_chunked)
				if (self.headers['Connection'].to_s.downcase.include?('keep-alive'))
					# If we are using keep-alive, but have no content-length and
					# no chunked transfer header, pretend this is the entire
					# buffer and call it done
					self.body_bytes_left = self.bufq.length
				elsif (not self.headers['Content-Length'] and self.class == Rex::Proto::Http::Request)
					# RFC 2616 says: "The presence of a message-body in a request
					# is signaled by the inclusion of a Content-Length or
					# Transfer-Encoding header field in the request's
					# message-headers."
					#
					# So if we haven't seen either a Content-Length or a
					# Transfer-Encoding header, there shouldn't be a message body.
					self.body_bytes_left = 0
				#else
				# Otherwise we need to keep reading until EOF
				end
			end
		end

		# Throw an error if we didnt parse the header properly
		if !self.headers.cmd_string
			raise RuntimeError, "Invalid command string", caller
		end

		# Move the state into body processing
		self.state = ParseState::ProcessingBody

		# Allow derived classes to update the parts of the command string
		self.update_cmd_parts(self.headers.cmd_string)
	end

	#
	# Parses the body portion of the request.
	#
	def parse_body
		# Just return if the buffer is empty
		if (self.bufq.length == 0)
			return
		end

		# Handle chunked transfer-encoding responses
		if (self.transfer_chunked and self.inside_chunk != 1 and self.bufq.length)

			# Remove any leading newlines or spaces
			self.bufq.lstrip!

			# If we didn't get a newline, then this might not be the full
			# length, go back and get more.
			# e.g. 
			#  first packet: "200"
			#  second packet: "0\r\n\r\n<html>..."
			if not bufq.index("\n")
				return
			end

			# Extract the actual hexadecimal length value
			clen = self.bufq.slice!(/^[a-fA-F0-9]+\r?\n/)

			clen.rstrip! if (clen)

			# if we happen to fall upon the end of the buffer for the next chunk len and have no data left, go get some more...
			if clen.nil? and self.bufq.length == 0
				return
			end

			# Invalid chunk length, exit out early
			if clen.nil?
				self.state = ParseState::Completed
				return
			end

			self.body_bytes_left = clen.hex

			if (self.body_bytes_left == 0)
				self.bufq.sub!(/^\r?\n/s,'')
				self.state = ParseState::Completed
				return
			end

			self.inside_chunk = 1
		end

		# If there are bytes remaining, slice as many as we can and append them
		# to our body state.
		if (self.body_bytes_left > 0)
			part = self.bufq.slice!(0, self.body_bytes_left)

			self.body += part
			self.body_bytes_left -= part.length
		# Otherwise, just read it all.
		else
			self.body += self.bufq
			self.bufq  = ''
		end

		# Finish this chunk and move on to the next one
		if (self.transfer_chunked and self.body_bytes_left == 0)
			self.inside_chunk = 0
			self.parse_body
			return
		end

		# If there are no more bytes left, then parsing has completed and we're
		# ready to go.
		if (not self.transfer_chunked and self.body_bytes_left == 0)
			self.state = ParseState::Completed
			return
		end
	end

end

end
end
end