'''
proxy.py

Copyright 2006 Andres Riancho

This file is part of w3af, w3af.sourceforge.net .

w3af 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 version 2 of the License.

w3af 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 w3af; if not, write to the Free Software
Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA

'''

import os
from os import sep
from BaseHTTPServer import BaseHTTPRequestHandler, HTTPServer
import mimetypes
import time
import socket, signal, select
from core.controllers.threads.w3afThread import w3afThread
import core.controllers.outputManager as om
from core.data.parsers.urlParser import *
from core.controllers.threads.threadManager import threadManagerObj as tm
from core.controllers.w3afException import w3afException

class proxy(w3afThread):
	'''
	This class defines a simple HTTP proxy, it is mainly used for "complex" plugins.
	
	You should call it like this:
		ws = proxy( '127.0.0.1', 8080, urlOpener )
		ws.start2()
	
	Or like this, if you want to override the proxyHandler (most times you want to do it...):
		ws = proxy( '127.0.0.1', 8080, urlOpener, proxyHandler=pH )
		ws.start2()
	
	Where pH is a class like this:
		class w3afProxyHandler(BaseHTTPRequestHandler):
			def do_GET(self):
				...
				...
	
	@author: Andres Riancho ( andres.riancho@gmail.com )	
	'''

	def __init__( self, ip, port, urlOpener, proxyHandler=None ):
		w3afThread.__init__(self)

		# Internal vars
		self._server = None
		self._proxyHandler = proxyHandler
		self._go = True
		self._running = False
		self._urlOpener = urlOpener
		self._tm = tm
		
		# User configured parameters
		self._ip = ip
		self._port = port
		
	def stop(self):
		om.out.debug('Calling stop of proxy daemon.')
		if self._running:
			self._server.server_close()
			self._go = False
			s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
			try:
				s.connect((self._ip, self._port))
				s.close()
			except:
				pass
			self._running = False
	
	def isRunning( self ):
		return self._running
	
	def run(self):
		'''
		Starts the proxy daemon.
		'''
		if self._proxyHandler == None:
			self._proxyHandler = w3afProxyHandler
		
		# Timeout to wait for thread starting
		time.sleep(0.1)
		
		om.out.debug( 'Using proxy handler: ' + str(self._proxyHandler) )
		self._proxyHandler._urlOpener = self._urlOpener
		try:
			self._server = HTTPServer( (self._ip, self._port), self._proxyHandler )
		except Exception, e:
			om.out.error('Failed to start proxy server, error: ' + str(e) )
		else:
			message = 'HTTP Proxy server listening on '+ self._ip + ':'+ str(self._port)
			om.out.debug( message )
			self._running = True
			
			while self._go:
				try:
					self._server.handle_request()
				except:
					self._server.server_close()
				
class w3afProxyHandler(BaseHTTPRequestHandler):

	def _sendToServer( self ):
		self.headers['Connection'] = 'close'
		
		# Do the request to the remote server
		if self.headers.dict.has_key('content-length'):
			# POST
			cl = int( self.headers['content-length'] )
			postData = self.rfile.read( cl )
			try:
				res = self._urlOpener.POST( self.path, data=postData, headers=self.headers )
			except w3afException, w:
				om.out.error('The proxy request failed, error: ' + str(w) )
				raise w
			except:
				raise
			return res
			
		else:
			# GET
			url = uri2url( self.path )
			qs = getQueryString( self.path )
			try:
				res = self._urlOpener.GET( url, data=str(qs), headers=self.headers )
			except w3afException, w:
				om.out.error('The proxy request failed, error: ' + str(w) )
				raise w
			except:
				raise
			return res
	
	def _sendError( self, exceptionObj ):
		'''
		Send an error to the browser.
		'''
		self.send_response( 400 )
		self.send_header( 'Connection', 'close')
		self.send_header( 'Content-type', 'text/html')		
		self.end_headers()
		self.wfile.write( 'Proxy error: ' + str(exceptionObj) )
		self.wfile.close()
	
	def _sendToBrowser( self, res ):
		'''
		Send a response to the browser
		'''
		# return the response to the browser
		try:
			self.send_response( res.getCode() )
			
			for header in res.getHeaders():
				self.send_header( header, res.getHeaders()[header] )
			self.send_header( 'Connection', 'close')
			self.end_headers()
			
			self.wfile.write( res.getBody() )
			self.wfile.close()
		except Exception, e:
			om.out.debug('Failed to send the data to the browser: ' + str(e) )		
	
	def doAll( self ):
		try:
			# Send the request to the remove webserver
			res = self._sendToServer()
		except Exception, e:
			self._sendError( e )
		else:
			self._sendToBrowser( res )
	
	do_GET = do_POST = do_HEAD = doAll
	
	class TimeoutError (Exception): pass
	def SIGALRM_handler(sig, stack): raise TimeoutError()
	# Windows signal.SIGALRM doesn't exist
	try:
		signal.signal(signal.SIGALRM, SIGALRM_handler)
	except:
		pass
	
	def _connect_to( self, netloc, sock):
		i = netloc.find(':')
		if i >= 0:
			host, port = netloc[:i], int(netloc[i+1:])
		else:
			host, port = netloc, 80
		signal.alarm(10)
		try:
			sock.connect((host, port))
		finally:
			signal.alarm(0)
		return True
			
	def do_CONNECT(self):
		soc = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
		try:
			if self._connect_to(self.path, soc):
				self.log_request(200)
				self.wfile.write(self.protocol_version + " 200 Connection established\r\n")
				self.wfile.write("Proxy-agent: %s\r\n" % self.version_string())
				self.wfile.write("\r\n")
				self._read_write(soc, 300)
		finally:
			om.out.debug('Closed proxy connection.')
			soc.close()
			self.connection.close()
	
	def _read_write(self, sock, max_idling=20):
		iw = [self.connection, sock]
		ow = []
		count = 0
		while 1:
			count += 1
			(ins, _, exs) = select.select(iw, ow, iw, 3)
			if exs: break
			
			if ins:
				for i in ins:
					if i is sock:
						out = self.connection
					else:
						out = sock
					data = i.recv(8192)
					if data:
						out.send(data)
						count = 0

			if count == max_idling: 
				om.out.debug('Closing Idle HTTP connection.')
				break

	def log_message( self, format, *args):
		'''
		I dont want messages written to stderr, please write them to the om.
		'''
		message = "Local proxy daemon handling request: %s - %s" % (self.address_string(),format%args) 
		om.out.debug( message )

