sfa/util/ssl_socket.py

   1 from ssl import SSLSocket
   2
   3 import textwrap
   4
   5 import _ssl             # if we can't import it, let the error propagate
   6
   7 from _ssl import SSLError
   8 from _ssl import CERT_NONE, CERT_OPTIONAL, CERT_REQUIRED
   9 from _ssl import PROTOCOL_SSLv2, PROTOCOL_SSLv3, PROTOCOL_SSLv23, PROTOCOL_TLSv1
  10 from _ssl import RAND_status, RAND_egd, RAND_add
  11 from _ssl import \
  12      SSL_ERROR_ZERO_RETURN, \
  13      SSL_ERROR_WANT_READ, \
  14      SSL_ERROR_WANT_WRITE, \
  15      SSL_ERROR_WANT_X509_LOOKUP, \
  16      SSL_ERROR_SYSCALL, \
  17      SSL_ERROR_SSL, \
  18      SSL_ERROR_WANT_CONNECT, \
  19      SSL_ERROR_EOF, \
  20      SSL_ERROR_INVALID_ERROR_CODE
  21
  22 from socket import socket, _fileobject
  23 from socket import getnameinfo as _getnameinfo
  24 import base64        # for DER-to-PEM translation
  25
  26 class SSLSocket(SSLSocket, socket):
  27
  28     """This class implements a subtype of socket.socket that wraps
  29     the underlying OS socket in an SSL context when necessary, and
  30     provides read and write methods over that channel."""
  31
  32     def __init__(self, sock, keyfile=None, certfile=None,
  33                  server_side=False, cert_reqs=CERT_NONE,
  34                  ssl_version=PROTOCOL_SSLv23, ca_certs=None,
  35                  do_handshake_on_connect=True,
  36                  suppress_ragged_eofs=True):
  37         socket.__init__(self, _sock=sock._sock)
  38         # the initializer for socket trashes the methods (tsk, tsk), so...
  39         self.send = lambda data, flags=0: SSLSocket.send(self, data, flags)
  40         self.sendto = lambda data, addr, flags=0: SSLSocket.sendto(self, data, addr, flags)
  41         self.recv = lambda buflen=1024, flags=0: SSLSocket.recv(self, buflen, flags)
  42         self.recvfrom = lambda addr, buflen=1024, flags=0: SSLSocket.recvfrom(self, addr, buflen, flags)
  43         self.recv_into = lambda buffer, nbytes=None, flags=0: SSLSocket.recv_into(self, buffer, nbytes, flags)
  44         self.recvfrom_into = lambda buffer, nbytes=None, flags=0: SSLSocket.recvfrom_into(self, buffer, nbytes, flags)
  45
  46         if certfile and not keyfile:
  47             keyfile = certfile
  48         # see if it's connected
  49         try:
  50             socket.getpeername(self)
  51         except:
  52             # no, no connection yet
  53             self._sslobj = None
  54         else:
  55             # yes, create the SSL object
  56             self._sslobj = _ssl.sslwrap(self._sock, server_side,
  57                                         keyfile, certfile,
  58                                         cert_reqs, ssl_version, ca_certs)
  59             if do_handshake_on_connect:
  60                 timeout = self.gettimeout()
  61                 try:
  62                     if timeout == 0:
  63                         self.settimeout(None)
  64                     self.do_handshake()
  65                 finally:
  66                     self.settimeout(timeout)
  67         self.keyfile = keyfile
  68         self.certfile = certfile
  69         self.cert_reqs = cert_reqs
  70         self.ssl_version = ssl_version
  71         self.ca_certs = ca_certs
  72         self.do_handshake_on_connect = do_handshake_on_connect
  73         self.suppress_ragged_eofs = suppress_ragged_eofs
  74         self._makefile_refs = 0
  75
  76