diff options
Diffstat (limited to 'Lib/ssl.py')
| -rw-r--r-- | Lib/ssl.py | 361 |
1 files changed, 105 insertions, 256 deletions
@@ -74,7 +74,7 @@ from _ssl import \ SSL_ERROR_EOF, \ SSL_ERROR_INVALID_ERROR_CODE -from socket import socket +from socket import socket, _fileobject from socket import getnameinfo as _getnameinfo import base64 # for DER-to-PEM translation @@ -86,8 +86,16 @@ class SSLSocket (socket): def __init__(self, sock, keyfile=None, certfile=None, server_side=False, cert_reqs=CERT_NONE, - ssl_version=PROTOCOL_SSLv23, ca_certs=None): + ssl_version=PROTOCOL_SSLv23, ca_certs=None, + do_handshake_on_connect=True, + suppress_ragged_eofs=True): socket.__init__(self, _sock=sock._sock) + # the initializer for socket trashes the methods (tsk, tsk), so... + self.send = lambda x, flags=0: SSLSocket.send(self, x, flags) + self.recv = lambda x, flags=0: SSLSocket.recv(self, x, flags) + self.sendto = lambda data, addr, flags=0: SSLSocket.sendto(self, data, addr, flags) + self.recvfrom = lambda addr, buflen, flags: SSLSocket.recvfrom(self, addr, buflen, flags) + if certfile and not keyfile: keyfile = certfile # see if it's connected @@ -101,18 +109,34 @@ class SSLSocket (socket): self._sslobj = _ssl.sslwrap(self._sock, server_side, keyfile, certfile, cert_reqs, ssl_version, ca_certs) + if do_handshake_on_connect: + timeout = self.gettimeout() + try: + self.settimeout(None) + self.do_handshake() + finally: + self.settimeout(timeout) self.keyfile = keyfile self.certfile = certfile self.cert_reqs = cert_reqs self.ssl_version = ssl_version self.ca_certs = ca_certs + self.do_handshake_on_connect = do_handshake_on_connect + self.suppress_ragged_eofs = suppress_ragged_eofs + self._makefile_refs = 0 def read(self, len=1024): """Read up to LEN bytes and return them. Return zero-length string on EOF.""" - return self._sslobj.read(len) + try: + return self._sslobj.read(len) + except SSLError, x: + if x.args[0] == SSL_ERROR_EOF and self.suppress_ragged_eofs: + return '' + else: + raise def write(self, data): @@ -143,16 +167,27 @@ class SSLSocket (socket): raise ValueError( "non-zero flags not allowed in calls to send() on %s" % self.__class__) - return self._sslobj.write(data) + 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: return socket.send(self, data, flags) - def send_to (self, data, addr, flags=0): + def sendto (self, data, addr, flags=0): if self._sslobj: - raise ValueError("send_to not allowed on instances of %s" % + raise ValueError("sendto not allowed on instances of %s" % self.__class__) else: - return socket.send_to(self, data, addr, flags) + return socket.sendto(self, data, addr, flags) def sendall (self, data, flags=0): if self._sslobj: @@ -160,7 +195,12 @@ class SSLSocket (socket): raise ValueError( "non-zero flags not allowed in calls to sendall() on %s" % self.__class__) - return self._sslobj.write(data) + amount = len(data) + count = 0 + while (count < amount): + v = self.send(data[count:]) + count += v + return amount else: return socket.sendall(self, data, flags) @@ -170,25 +210,51 @@ class SSLSocket (socket): raise ValueError( "non-zero flags not allowed in calls to sendall() on %s" % self.__class__) - return self._sslobj.read(data, buflen) + while True: + try: + return self.read(buflen) + except SSLError, x: + if x.args[0] == SSL_ERROR_WANT_READ: + continue + else: + raise x else: return socket.recv(self, buflen, flags) - def recv_from (self, addr, buflen=1024, flags=0): + def recvfrom (self, addr, buflen=1024, flags=0): if self._sslobj: - raise ValueError("recv_from not allowed on instances of %s" % + raise ValueError("recvfrom not allowed on instances of %s" % self.__class__) else: - return socket.recv_from(self, addr, buflen, flags) + return socket.recvfrom(self, addr, buflen, flags) - def shutdown(self, how): + def pending (self): + if self._sslobj: + return self._sslobj.pending() + else: + return 0 + + def shutdown (self, how): self._sslobj = None socket.shutdown(self, how) - def close(self): + def close (self): self._sslobj = None socket.close(self) + def close (self): + if self._makefile_refs < 1: + self._sslobj = None + socket.close(self) + else: + self._makefile_refs -= 1 + + def do_handshake (self): + + """Perform a TLS/SSL handshake.""" + + self._sslobj.do_handshake() + def connect(self, addr): """Connects to remote ADDR, and then wraps the connection in @@ -202,6 +268,8 @@ class SSLSocket (socket): 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): @@ -210,260 +278,39 @@ class SSLSocket (socket): SSL channel, and the address of the remote client.""" newsock, addr = socket.accept(self) - return (SSLSocket(newsock, True, self.keyfile, self.certfile, - self.cert_reqs, self.ssl_version, - self.ca_certs), addr) - + return (SSLSocket(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), + addr) def makefile(self, mode='r', bufsize=-1): """Ouch. Need to make and return a file-like object that works with the SSL connection.""" - if self._sslobj: - return SSLFileStream(self._sslobj, mode, bufsize) - else: - return socket.makefile(self, mode, bufsize) - - -class SSLFileStream: - - """A class to simulate a file stream on top of a socket. - Most of this is just lifted from the socket module, and - adjusted to work with an SSL stream instead of a socket.""" - - - default_bufsize = 8192 - name = "<SSL stream>" - - __slots__ = ["mode", "bufsize", "softspace", - # "closed" is a property, see below - "_sslobj", "_rbufsize", "_wbufsize", "_rbuf", "_wbuf", - "_close", "_fileno"] - - def __init__(self, sslobj, mode='rb', bufsize=-1, close=False): - self._sslobj = sslobj - self.mode = mode # Not actually used in this version - if bufsize < 0: - bufsize = self.default_bufsize - self.bufsize = bufsize - self.softspace = False - if bufsize == 0: - self._rbufsize = 1 - elif bufsize == 1: - self._rbufsize = self.default_bufsize - else: - self._rbufsize = bufsize - self._wbufsize = bufsize - self._rbuf = "" # A string - self._wbuf = [] # A list of strings - self._close = close - self._fileno = -1 - - def _getclosed(self): - return self._sslobj is None - closed = property(_getclosed, doc="True if the file is closed") - - def fileno(self): - return self._fileno - - def close(self): - try: - if self._sslobj: - self.flush() - finally: - if self._close and self._sslobj: - self._sslobj.close() - self._sslobj = None - - def __del__(self): - try: - self.close() - except: - # close() may fail if __init__ didn't complete - pass - - def flush(self): - if self._wbuf: - buffer = "".join(self._wbuf) - self._wbuf = [] - count = 0 - while (count < len(buffer)): - written = self._sslobj.write(buffer) - count += written - buffer = buffer[written:] - - def write(self, data): - data = str(data) # XXX Should really reject non-string non-buffers - if not data: - return - self._wbuf.append(data) - if (self._wbufsize == 0 or - self._wbufsize == 1 and '\n' in data or - self._get_wbuf_len() >= self._wbufsize): - self.flush() - - def writelines(self, list): - # XXX We could do better here for very long lists - # XXX Should really reject non-string non-buffers - self._wbuf.extend(filter(None, map(str, list))) - if (self._wbufsize <= 1 or - self._get_wbuf_len() >= self._wbufsize): - self.flush() - - def _get_wbuf_len(self): - buf_len = 0 - for x in self._wbuf: - buf_len += len(x) - return buf_len - - def read(self, size=-1): - data = self._rbuf - if size < 0: - # Read until EOF - buffers = [] - if data: - buffers.append(data) - self._rbuf = "" - if self._rbufsize <= 1: - recv_size = self.default_bufsize - else: - recv_size = self._rbufsize - while True: - data = self._sslobj.read(recv_size) - if not data: - break - buffers.append(data) - return "".join(buffers) - else: - # Read until size bytes or EOF seen, whichever comes first - buf_len = len(data) - if buf_len >= size: - self._rbuf = data[size:] - return data[:size] - buffers = [] - if data: - buffers.append(data) - self._rbuf = "" - while True: - left = size - buf_len - recv_size = max(self._rbufsize, left) - data = self._sslobj.read(recv_size) - if not data: - break - buffers.append(data) - n = len(data) - if n >= left: - self._rbuf = data[left:] - buffers[-1] = data[:left] - break - buf_len += n - return "".join(buffers) - - def readline(self, size=-1): - data = self._rbuf - if size < 0: - # Read until \n or EOF, whichever comes first - if self._rbufsize <= 1: - # Speed up unbuffered case - assert data == "" - buffers = [] - while data != "\n": - data = self._sslobj.read(1) - if not data: - break - buffers.append(data) - return "".join(buffers) - nl = data.find('\n') - if nl >= 0: - nl += 1 - self._rbuf = data[nl:] - return data[:nl] - buffers = [] - if data: - buffers.append(data) - self._rbuf = "" - while True: - data = self._sslobj.read(self._rbufsize) - if not data: - break - buffers.append(data) - nl = data.find('\n') - if nl >= 0: - nl += 1 - self._rbuf = data[nl:] - buffers[-1] = data[:nl] - break - return "".join(buffers) - else: - # Read until size bytes or \n or EOF seen, whichever comes first - nl = data.find('\n', 0, size) - if nl >= 0: - nl += 1 - self._rbuf = data[nl:] - return data[:nl] - buf_len = len(data) - if buf_len >= size: - self._rbuf = data[size:] - return data[:size] - buffers = [] - if data: - buffers.append(data) - self._rbuf = "" - while True: - data = self._sslobj.read(self._rbufsize) - if not data: - break - buffers.append(data) - left = size - buf_len - nl = data.find('\n', 0, left) - if nl >= 0: - nl += 1 - self._rbuf = data[nl:] - buffers[-1] = data[:nl] - break - n = len(data) - if n >= left: - self._rbuf = data[left:] - buffers[-1] = data[:left] - break - buf_len += n - return "".join(buffers) - - def readlines(self, sizehint=0): - total = 0 - list = [] - while True: - line = self.readline() - if not line: - break - list.append(line) - total += len(line) - if sizehint and total >= sizehint: - break - return list - - # Iterator protocols - - def __iter__(self): - return self - - def next(self): - line = self.readline() - if not line: - raise StopIteration - return line - + self._makefile_refs += 1 + return _fileobject(self, mode, bufsize) def wrap_socket(sock, keyfile=None, certfile=None, server_side=False, cert_reqs=CERT_NONE, - ssl_version=PROTOCOL_SSLv23, ca_certs=None): + ssl_version=PROTOCOL_SSLv23, ca_certs=None, + do_handshake_on_connect=True, + suppress_ragged_eofs=True): return SSLSocket(sock, keyfile=keyfile, certfile=certfile, server_side=server_side, cert_reqs=cert_reqs, - ssl_version=ssl_version, ca_certs=ca_certs) + ssl_version=ssl_version, ca_certs=ca_certs, + do_handshake_on_connect=do_handshake_on_connect, + suppress_ragged_eofs=suppress_ragged_eofs) + # some utility functions @@ -549,5 +396,7 @@ def sslwrap_simple (sock, keyfile=None, certfile=None): for compability with Python 2.5 and earlier. Will disappear in Python 3.0.""" - return _ssl.sslwrap(sock._sock, 0, keyfile, certfile, CERT_NONE, - PROTOCOL_SSLv23, None) + ssl_sock = _ssl.sslwrap(sock._sock, 0, keyfile, certfile, CERT_NONE, + PROTOCOL_SSLv23, None) + ssl_sock.do_handshake() + return ssl_sock |