306 lines
12 KiB
Python
306 lines
12 KiB
Python
__ssl = __import__('ssl')
|
|
|
|
for attr in dir(__ssl):
|
|
exec "%s = __ssl.%s" % (attr, attr)
|
|
|
|
import errno
|
|
import time
|
|
|
|
from eventlet.api import trampoline, getcurrent
|
|
from thread import get_ident
|
|
from eventlet.greenio import set_nonblocking, GreenSocket, SOCKET_CLOSED, CONNECT_ERR, CONNECT_SUCCESS
|
|
orig_socket = __import__('socket')
|
|
socket = orig_socket.socket
|
|
|
|
|
|
class GreenSSLSocket(__ssl.SSLSocket):
|
|
""" This is a green version of the SSLSocket class from the ssl module added
|
|
in 2.6. For documentation on it, please see the Python standard
|
|
documentation."""
|
|
# we are inheriting from SSLSocket because its constructor calls
|
|
# do_handshake whose behavior we wish to override
|
|
def __init__(self, sock, *args, **kw):
|
|
if not isinstance(sock, GreenSocket):
|
|
sock = GreenSocket(sock)
|
|
|
|
self.act_non_blocking = sock.act_non_blocking
|
|
self.timeout = sock.timeout
|
|
super(GreenSSLSocket, self).__init__(sock.fd, *args, **kw)
|
|
del sock
|
|
|
|
# the superclass initializer trashes the methods so...
|
|
self.send = lambda data, flags=0: GreenSSLSocket.send(self, data, flags)
|
|
self.sendto = lambda data, addr, flags=0: GreenSSLSocket.sendto(self, data, addr, flags)
|
|
self.recv = lambda buflen=1024, flags=0: GreenSSLSocket.recv(self, buflen, flags)
|
|
self.recvfrom = lambda addr, buflen=1024, flags=0: GreenSSLSocket.recvfrom(self, addr, buflen, flags)
|
|
self.recv_into = lambda buffer, nbytes=None, flags=0: GreenSSLSocket.recv_into(self, buffer, nbytes, flags)
|
|
self.recvfrom_into = lambda buffer, nbytes=None, flags=0: GreenSSLSocket.recvfrom_into(self, buffer, nbytes, flags)
|
|
|
|
def settimeout(self, timeout):
|
|
self.timeout = timeout
|
|
|
|
def gettimeout(self):
|
|
return self.timeout
|
|
|
|
def setblocking(self, flag):
|
|
if flag:
|
|
self.act_non_blocking = False
|
|
self.timeout = None
|
|
else:
|
|
self.act_non_blocking = True
|
|
self.timeout = 0.0
|
|
|
|
def _call_trampolining(self, func, *a, **kw):
|
|
if self.act_non_blocking:
|
|
return func(*a, **kw)
|
|
else:
|
|
while True:
|
|
try:
|
|
return func(*a, **kw)
|
|
except SSLError, exc:
|
|
if exc[0] == SSL_ERROR_WANT_READ:
|
|
trampoline(self.fileno(),
|
|
read=True,
|
|
timeout=self.gettimeout(),
|
|
timeout_exc=SSLError)
|
|
elif exc[0] == SSL_ERROR_WANT_WRITE:
|
|
trampoline(self.fileno(),
|
|
write=True,
|
|
timeout=self.gettimeout(),
|
|
timeout_exc=SSLError)
|
|
else:
|
|
raise
|
|
|
|
|
|
def write(self, data):
|
|
"""Write DATA to the underlying SSL channel. Returns
|
|
number of bytes of DATA actually transmitted."""
|
|
return self._call_trampolining(
|
|
super(GreenSSLSocket, self).write, data)
|
|
|
|
def read(self, len=1024):
|
|
"""Read up to LEN bytes and return them.
|
|
Return zero-length string on EOF."""
|
|
return self._call_trampolining(
|
|
super(GreenSSLSocket, self).read,len)
|
|
|
|
def send (self, data, flags=0):
|
|
# *NOTE: gross, copied code from ssl.py becase it's not factored well enough to be used as-is
|
|
if self._sslobj:
|
|
if flags != 0:
|
|
raise ValueError(
|
|
"non-zero flags not allowed in calls to send() on %s" %
|
|
self.__class__)
|
|
while True:
|
|
try:
|
|
v = self._sslobj.write(data)
|
|
except SSLError, x:
|
|
if x.args[0] == SSL_ERROR_WANT_READ:
|
|
return 0
|
|
elif x.args[0] == SSL_ERROR_WANT_WRITE:
|
|
return 0
|
|
else:
|
|
raise
|
|
else:
|
|
return v
|
|
else:
|
|
while True:
|
|
try:
|
|
return socket.send(self, data, flags)
|
|
except orig_socket.error, e:
|
|
if self.act_non_blocking:
|
|
raise
|
|
if e[0] == errno.EWOULDBLOCK or \
|
|
e[0] == errno.ENOTCONN:
|
|
return 0
|
|
raise
|
|
|
|
def sendto (self, data, addr, flags=0):
|
|
# *NOTE: gross, copied code from ssl.py becase it's not factored well enough to be used as-is
|
|
if self._sslobj:
|
|
raise ValueError("sendto not allowed on instances of %s" %
|
|
self.__class__)
|
|
else:
|
|
trampoline(self.fileno(), write=True, timeout_exc=orig_socket.timeout)
|
|
return socket.sendto(self, data, addr, flags)
|
|
|
|
def sendall (self, data, flags=0):
|
|
# *NOTE: gross, copied code from ssl.py becase it's not factored well enough to be used as-is
|
|
if self._sslobj:
|
|
if flags != 0:
|
|
raise ValueError(
|
|
"non-zero flags not allowed in calls to sendall() on %s" %
|
|
self.__class__)
|
|
amount = len(data)
|
|
count = 0
|
|
while (count < amount):
|
|
v = self.send(data[count:])
|
|
count += v
|
|
return amount
|
|
else:
|
|
while True:
|
|
try:
|
|
return socket.sendall(self, buflen, flags)
|
|
except orig_socket.error, e:
|
|
if self.act_non_blocking:
|
|
raise
|
|
if e[0] == errno.EWOULDBLOCK:
|
|
trampoline(self.fileno(), write=True,
|
|
timeout=self.gettimeout(), timeout_exc=orig_socket.timeout)
|
|
if e[0] in SOCKET_CLOSED:
|
|
return ''
|
|
raise
|
|
|
|
def recv(self, buflen=1024, flags=0):
|
|
# *NOTE: gross, copied code from ssl.py becase it's not factored well enough to be used as-is
|
|
if self._sslobj:
|
|
if flags != 0:
|
|
raise ValueError(
|
|
"non-zero flags not allowed in calls to recv() on %s" %
|
|
self.__class__)
|
|
read = self.read(buflen)
|
|
return read
|
|
else:
|
|
while True:
|
|
try:
|
|
return socket.recv(self, buflen, flags)
|
|
except orig_socket.error, e:
|
|
if self.act_non_blocking:
|
|
raise
|
|
if e[0] == errno.EWOULDBLOCK:
|
|
trampoline(self.fileno(), read=True,
|
|
timeout=self.gettimeout(), timeout_exc=orig_socket.timeout)
|
|
if e[0] in SOCKET_CLOSED:
|
|
return ''
|
|
raise
|
|
|
|
|
|
def recv_into (self, buffer, nbytes=None, flags=0):
|
|
if not self.act_non_blocking:
|
|
trampoline(self.fileno(), read=True, timeout=self.gettimeout(), timeout_exc=orig_socket.timeout)
|
|
return super(GreenSSLSocket, self).recv_into(buffer, nbytes, flags)
|
|
|
|
def recvfrom (self, addr, buflen=1024, flags=0):
|
|
if not self.act_non_blocking:
|
|
trampoline(self.fileno(), read=True, timeout=self.gettimeout(), timeout_exc=orig_socket.timeout)
|
|
return super(GreenSSLSocket, self).recvfrom(addr, buflen, flags)
|
|
|
|
def recvfrom_into (self, buffer, nbytes=None, flags=0):
|
|
if not self.act_non_blocking:
|
|
trampoline(self.fileno(), read=True, timeout=self.gettimeout(), timeout_exc=orig_socket.timeout)
|
|
return super(GreenSSLSocket, self).recvfrom_into(buffer, nbytes, flags)
|
|
|
|
def unwrap(self):
|
|
return GreenSocket(super(GreenSSLSocket, self).unwrap())
|
|
|
|
def do_handshake(self):
|
|
"""Perform a TLS/SSL handshake."""
|
|
return self._call_trampolining(
|
|
super(GreenSSLSocket, self).do_handshake)
|
|
|
|
def _socket_connect(self, addr):
|
|
real_connect = socket.connect
|
|
if self.act_non_blocking:
|
|
return real_connect(self, addr)
|
|
else:
|
|
# *NOTE: gross, copied code from greenio because it's not factored
|
|
# well enough to reuse
|
|
if self.gettimeout() is None:
|
|
while True:
|
|
try:
|
|
return real_connect(self, addr)
|
|
except orig_socket.error, exc:
|
|
if exc[0] in CONNECT_ERR:
|
|
trampoline(self.fileno(), write=True)
|
|
elif exc[0] in CONNECT_SUCCESS:
|
|
return
|
|
else:
|
|
raise
|
|
else:
|
|
end = time.time() + self.gettimeout()
|
|
while True:
|
|
try:
|
|
real_connect(self, addr)
|
|
except orig_socket.error, exc:
|
|
if exc[0] in CONNECT_ERR:
|
|
trampoline(self.fileno(), write=True,
|
|
timeout=end-time.time(), timeout_exc=orig_socket.timeout)
|
|
elif exc[0] in CONNECT_SUCCESS:
|
|
return
|
|
else:
|
|
raise
|
|
if time.time() >= end:
|
|
raise orig_socket.timeout
|
|
|
|
|
|
def connect(self, addr):
|
|
"""Connects to remote ADDR, and then wraps the connection in
|
|
an SSL channel."""
|
|
# *NOTE: grrrrr copied this code from ssl.py because of the reference
|
|
# to socket.connect which we don't want to call directly
|
|
if self._sslobj:
|
|
raise ValueError("attempt to connect already-connected SSLSocket!")
|
|
self._socket_connect(addr)
|
|
self._sslobj = _ssl.sslwrap(self._sock, False, self.keyfile, self.certfile,
|
|
self.cert_reqs, self.ssl_version,
|
|
self.ca_certs)
|
|
if self.do_handshake_on_connect:
|
|
self.do_handshake()
|
|
|
|
def accept(self):
|
|
"""Accepts a new connection from a remote client, and returns
|
|
a tuple containing that new connection wrapped with a server-side
|
|
SSL channel, and the address of the remote client."""
|
|
# RDW grr duplication of code from greenio
|
|
if self.act_non_blocking:
|
|
newsock, addr = socket.accept(self)
|
|
else:
|
|
while True:
|
|
try:
|
|
newsock, addr = socket.accept(self)
|
|
set_nonblocking(newsock)
|
|
break
|
|
except orig_socket.error, e:
|
|
if e[0] != errno.EWOULDBLOCK:
|
|
raise
|
|
trampoline(self.fileno(), read=True, timeout=self.gettimeout(),
|
|
timeout_exc=orig_socket.timeout)
|
|
|
|
new_ssl = type(self)(newsock,
|
|
keyfile=self.keyfile,
|
|
certfile=self.certfile,
|
|
server_side=True,
|
|
cert_reqs=self.cert_reqs,
|
|
ssl_version=self.ssl_version,
|
|
ca_certs=self.ca_certs,
|
|
do_handshake_on_connect=self.do_handshake_on_connect,
|
|
suppress_ragged_eofs=self.suppress_ragged_eofs)
|
|
return (new_ssl, addr)
|
|
|
|
|
|
SSLSocket = GreenSSLSocket
|
|
|
|
def wrap_socket(sock, keyfile=None, certfile=None,
|
|
server_side=False, cert_reqs=CERT_NONE,
|
|
ssl_version=PROTOCOL_SSLv23, ca_certs=None,
|
|
do_handshake_on_connect=True,
|
|
suppress_ragged_eofs=True):
|
|
return GreenSSLSocket(sock, keyfile=keyfile, certfile=certfile,
|
|
server_side=server_side, cert_reqs=cert_reqs,
|
|
ssl_version=ssl_version, ca_certs=ca_certs,
|
|
do_handshake_on_connect=do_handshake_on_connect,
|
|
suppress_ragged_eofs=suppress_ragged_eofs)
|
|
|
|
|
|
if hasattr(__ssl, 'sslwrap_simple'):
|
|
def sslwrap_simple(sock, keyfile=None, certfile=None):
|
|
"""A replacement for the old socket.ssl function. Designed
|
|
for compability with Python 2.5 and earlier. Will disappear in
|
|
Python 3.0."""
|
|
ssl_sock = GreenSSLSocket(sock, keyfile=keyfile, certfile=certfile,
|
|
server_side=False,
|
|
cert_reqs=CERT_NONE,
|
|
ssl_version=PROTOCOL_SSLv23,
|
|
ca_certs=None)
|
|
return ssl_sock
|