import socket, base64, string
import httplib
from OpenSSL import SSL

class OpenSSLFile(httplib.SSLFile):
    def _read(self):
        buf = ''

        while True:
            try:
                buf = self._ssl.recv(self._bufsize)
	    except SSL.SysCallError:
		break
	    except SSL.ZeroReturnError:
        	break
            except socket.sslerror, err:
                if (err[0] == socket.SSL_ERROR_WANT_READ
                    or err[0] == socket.SSL_ERROR_WANT_WRITE):
                    continue
                if (err[0] == socket.SSL_ERROR_ZERO_RETURN
                    or err[0] == socket.SSL_ERROR_EOF):
                    break
                raise
            except socket.error, err:
                if err[0] == errno.EINTR:
                    continue
                if err[0] == errno.EBADF:
                    # XXX socket was closed?
                    break
                raise
            else:
                break
        return buf

class FakeOpenSSLSocket(httplib.FakeSocket):
#    def __init__(self, sock, ssl):
#	httplib.FakeSocket.__init__(self,sock,ssl)
#	self.resp_sock = OpenSSLFile(self._shared, self._ssl)

    def makefile(self, mode, bufsize=None):
#+	return self.resp_sock
        if mode != 'r' and mode != 'rb':
            raise UnimplementedFileMode()
        return OpenSSLFile(self._shared, self._ssl, bufsize)

class ProxyOpenSSLConnection(httplib.HTTPSConnection):
    "This class allows communication via SSL."

    default_proxy_port = '3128'

    def __init__(self, host, port=None, key_file=None, cert_file=None, strict=None,
	proxy_host=None, proxy_port=None, proxy_user=None, proxy_pass=None):
        httplib.HTTPSConnection.__init__(self, host, port, key_file, cert_file, strict)
	self._set_proxy_hostport(proxy_host, proxy_port, proxy_user, proxy_pass)

#    def __del__(self):
#	self.real_close()

    def _set_proxy_hostport(self, host, port, user, password):
	i = host.find('://')
	if i >= 0:
	    host = host[i+3:]

	i = host.find('@')
	if i >= 0:
	    auth = host[:i]
	    host = host[i+1:]
	    if user is None or password is None:
	        i = auth.find(':')
		if i >= 0:
		    if user is None:
			user = auth[:i]
		    if password is None:
			password = auth[i+1:]
		else:
		    if user is None:
			user = auth
	
        i = host.find(':')
        if port is None:
            if i >= 0:
                try:
                    port = int(host[i+1:])
                except ValueError:
                    raise httplib.InvalidURL("nonnumeric port: '%s'" % host[i+1:])
            else:
                port = self.default_proxy_port
        host = host[:i]
	
        self.proxy_host = host
        self.proxy_port = port
	self.proxy_user = user
	self.proxy_pass = password


    def connect(self):
	if self.sock:
	    return
	
	s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
	try:
    	    s.connect((self.proxy_host, self.proxy_port))
	    
	    proxy_connect = "CONNECT %s:%s HTTP/1.0\r\n" % (self.host, self.port)
#	    proxy_connect = "CONNECT %s:%s HTTP/1.1\r\n" % (self.host, self.port)
#	    proxy_connect = proxy_connect + "Proxy-Connection: keep-alive\r\n"
	    if self.proxy_user:
		if self.proxy_pass:
		    proxy_authorization='Proxy-Authorization: Basic '+ base64.encodestring(self.proxy_user+':'+self.proxy_pass)+'\r\n'
		else:
		    proxy_authorization='Proxy-Authorization: Basic '+ base64.encodestring(self.proxy_user)+'\r\n'
		proxy_request=proxy_connect+proxy_authorization
	    else:
		proxy_request=proxy_connect + "\r\n"

	    # Use the CONNECT method to get a connection to the actual server
	    s.send(proxy_request)
	    resp=s.recv(1024)

	    if self.debuglevel > 0:
	        print "Proxy response: %s" % string.strip(resp)
	
	    ctx = SSL.Context(SSL.SSLv23_METHOD)
	    conn = SSL.Connection(ctx, s)
	    
	    # Go to client mode
	    conn.set_connect_state()
                
	    if self.debuglevel > 0:
        	print "proxy connect: (%s, %s)" % (self.proxy_host, self.proxy_port)

	    self.sock = FakeOpenSSLSocket(s,conn)

	except socket.error, msg:
	    if self.debuglevel > 0:
        	print "proxy fail: (%s, %s)" % (self.proxy_host, self.proxy_port)

            raise socket.error, msg

#   def close(self):
#	pass
#	
#    def real_close(self):
#	httplib.HTTPSConnection.close(self)



class OpenSSLProxy(httplib.HTTPS):
    _connection_class = ProxyOpenSSLConnection

    def __init__(self, host='', port=None, key_file=None, cert_file=None, strict=None,
	proxy_host = None , proxy_port = None, proxy_user = None, proxy_pass = None):
    
        if port == 0:
            port = None
        if proxy_port == 0:
	    proxy_port = None		
		
    	self._setup(self._connection_class(host, port, key_file, cert_file, strict,
		proxy_host,proxy_port,proxy_user,proxy_pass))





import urllib2
from urllib import splittype, splithost, addinfourl, getproxies

class OpenSSLProxyConnection:
    def __init__(self, host, connection, age):
	self.host = host
	self.connection = connection
	self.age = age
    
    def set_age(self, age):
	self.age = age
    
    def get_age(self):
	return self.age
    
    def get_host(self):
	return self.host
	
    def get_connection(self):
	return self.connection
    
    

class OpenSSLProxyHandler(urllib2.HTTPSHandler):
    def __init__(self, proxy_host=None, proxy_port=None, proxy_user=None, proxy_pass=None):
	self.age = 0
	self.persistent = 0
	self.connections = {}
	self.proxy_host = proxy_host
	self.proxy_port = proxy_port
	self.proxy_user = proxy_user
	self.proxy_pass = proxy_pass
	self.cert = None
	self.pkey = None

    def set_persistent_connections(self, max_conn=5):
	self.persistent = max_conn
	    
    def set_x509(self,pkey,cert):
	self.pkey = pkey
	self.cert = cert
    
    def https_open(self, req):
        host = req.get_host()
        if not host:
            raise URLError('no host given')
	
	if self.proxy_host:
	    if self.persistent:
		self.age=self.age+1
		try:
		    connection = self.connections[host]
		    h = connection.get_connection()
		    connection.set_age(self.age)
		except:
		    if len(self.connections)>=self.persistent:
			minage = 1000000
			for h,conn in self.connections.iteritems():
			    age = conn.get_age()
			    if age<minage:
				minage = age
				connection = h
			del self.connections[connection]
		    		    
		    h = OpenSSLProxy(host, None, self.pkey, self.cert, None,
			self.proxy_host, self.proxy_port, self.proxy_user, self.proxy_pass)
		    connection = OpenSSLProxyConnection(host,h,self.age)
		    self.connections[host] = connection
	    else:		    	
		h = OpenSSLProxy(host, None, self.pkey, self.cert, None,
		    self.proxy_host, self.proxy_port, self.proxy_user, self.proxy_pass)
	else:
	    return self.do_open(httplib.HTTPS,req)

        if req.has_data():
            data = req.get_data()
            h.putrequest('POST', req.get_selector())
            if not 'Content-type' in req.headers:
                h.putheader('Content-type',
                            'application/x-www-form-urlencoded')
            if not 'Content-length' in req.headers:
                h.putheader('Content-length', '%d' % len(data))
        else:
            h.putrequest('GET', req.get_selector())

        scheme, sel = splittype(req.get_selector())
        sel_host, sel_path = splithost(sel)
        h.putheader('Host', sel_host or host)
        for name, value in self.parent.addheaders:
            name = name.capitalize()
            if name not in req.headers:
                h.putheader(name, value)
        for k, v in req.headers.items():
            h.putheader(k, v)
        # httplib will attempt to connect() here.  be prepared
        # to convert a socket error to a URLError.
        try:
            h.endheaders()
        except socket.error, err:
            raise URLError(err)
        if req.has_data():
            h.send(data)

        code, msg, hdrs = h.getreply()
        fp = h.getfile()
        if code == 200:
            return addinfourl(fp, hdrs, req.get_full_url())
        else:
            return self.parent.error('http', req, fp, code, msg, hdrs)




def OpenSSLInstallHandler(proxy_host = None, proxy_port = None, proxy_user = None, proxy_pass = None, persistent_connections = 0):
	proxies=getproxies()
	if proxies.has_key('https'):
	    if not proxy_host:
		proxy_host = proxies['https']
	    del proxies['https']
	
	proxy = urllib2.ProxyHandler(proxies)
    	openssl = OpenSSLProxyHandler(proxy_host, proxy_port, proxy_user, proxy_pass)

	if persistent_connections==0:
	    openssl.set_persistent_connections()
	elif persistent_connections:
    	    openssl.set_persistent_connections(persistent_connections)
	    
	opener = urllib2.build_opener(proxy, openssl)
	urllib2.install_opener(opener)