Lib/ssl.py

   1 # Wrapper module for _ssl, providing some additional facilities
   2 # implemented in Python.  Written by Bill Janssen.
   3
   4 """This module provides some more Pythonic support for SSL.
   5
   6 Object types:
   7
   8   SSLSocket -- subtype of socket.socket which does SSL over the socket
   9
  10 Exceptions:
  11
  12   SSLError -- exception raised for I/O errors
  13
  14 Functions:
  15
  16   cert_time_to_seconds -- convert time string used for certificate
  17                           notBefore and notAfter functions to integer
  18                           seconds past the Epoch (the time values
  19                           returned from time.time())
  20
  21   fetch_server_certificate (HOST, PORT) -- fetch the certificate provided
  22                           by the server running on HOST at port PORT.  No
  23                           validation of the certificate is performed.
  24
  25 Integer constants:
  26
  27 SSL_ERROR_ZERO_RETURN
  28 SSL_ERROR_WANT_READ
  29 SSL_ERROR_WANT_WRITE
  30 SSL_ERROR_WANT_X509_LOOKUP
  31 SSL_ERROR_SYSCALL
  32 SSL_ERROR_SSL
  33 SSL_ERROR_WANT_CONNECT
  34
  35 SSL_ERROR_EOF
  36 SSL_ERROR_INVALID_ERROR_CODE
  37
  38 The following group define certificate requirements that one side is
  39 allowing/requiring from the other side:
  40
  41 CERT_NONE - no certificates from the other side are required (or will
  42             be looked at if provided)
  43 CERT_OPTIONAL - certificates are not required, but if provided will be
  44                 validated, and if validation fails, the connection will
  45                 also fail
  46 CERT_REQUIRED - certificates are required, and will be validated, and
  47                 if validation fails, the connection will also fail
  48
  49 The following constants identify various SSL protocol variants:
  50
  51 PROTOCOL_SSLv2
  52 PROTOCOL_SSLv3
  53 PROTOCOL_SSLv23
  54 PROTOCOL_TLSv1
  55 """
  56
  57 import textwrap
  58
  59 import _ssl             # if we can't import it, let the error propagate
  60
  61 from _ssl import SSLError
  62 from _ssl import CERT_NONE, CERT_OPTIONAL, CERT_REQUIRED
  63 from _ssl import (PROTOCOL_SSLv2, PROTOCOL_SSLv3, PROTOCOL_SSLv23,
  64                   PROTOCOL_TLSv1)
  65 from _ssl import RAND_status, RAND_egd, RAND_add
  66 from _ssl import (
  67     SSL_ERROR_ZERO_RETURN,
  68     SSL_ERROR_WANT_READ,
  69     SSL_ERROR_WANT_WRITE,
  70     SSL_ERROR_WANT_X509_LOOKUP,
  71     SSL_ERROR_SYSCALL,
  72     SSL_ERROR_SSL,
  73     SSL_ERROR_WANT_CONNECT,
  74     SSL_ERROR_EOF,
  75     SSL_ERROR_INVALID_ERROR_CODE,
  76     )
  77
  78 from socket import getnameinfo as _getnameinfo
  79 from socket import error as socket_error
  80 from socket import dup as _dup
  81 from socket import socket, AF_INET, SOCK_STREAM
  82 import base64        # for DER-to-PEM translation
  83 import traceback
  84
  85 class SSLSocket(socket):
  86
  87     """This class implements a subtype of socket.socket that wraps
  88     the underlying OS socket in an SSL context when necessary, and
  89     provides read and write methods over that channel."""
  90
  91     def __init__(self, sock=None, keyfile=None, certfile=None,
  92                  server_side=False, cert_reqs=CERT_NONE,
  93                  ssl_version=PROTOCOL_SSLv23, ca_certs=None,
  94                  do_handshake_on_connect=True,
  95                  family=AF_INET, type=SOCK_STREAM, proto=0, fileno=None,
  96                  suppress_ragged_eofs=True):
  97
  98         if sock is not None:
  99             socket.__init__(self,
 100                             family=sock.family,
 101                             type=sock.type,
 102                             proto=sock.proto,
 103                             fileno=_dup(sock.fileno()))
 104             sock.close()
 105         elif fileno is not None:
 106             socket.__init__(self, fileno=fileno)
 107         else:
 108             socket.__init__(self, family=family, type=type, proto=proto)
 109
 110         self._closed = False
 111
 112         if certfile and not keyfile:
 113             keyfile = certfile
 114         # see if it's connected
 115         try:
 116             socket.getpeername(self)
 117         except socket_error:
 118             # no, no connection yet
 119             self._sslobj = None
 120         else:
 121             # yes, create the SSL object
 122             try:
 123                 self._sslobj = _ssl.sslwrap(self, server_side,
 124                                             keyfile, certfile,
 125                                             cert_reqs, ssl_version, ca_certs)
 126                 if do_handshake_on_connect:
 127                     timeout = self.gettimeout()
 128                     if timeout == 0.0:
 129                         # non-blocking
 130                         raise ValueError("do_handshake_on_connect should not be specified for non-blocking sockets")
 131                     self.do_handshake()
 132
 133             except socket_error as x:
 134                 self.close()
 135                 raise x
 136
 137         self.keyfile = keyfile
 138         self.certfile = certfile
 139         self.cert_reqs = cert_reqs
 140         self.ssl_version = ssl_version
 141         self.ca_certs = ca_certs
 142         self.do_handshake_on_connect = do_handshake_on_connect
 143         self.suppress_ragged_eofs = suppress_ragged_eofs
 144
 145     def dup(self):
 146         raise NotImplemented("Can't dup() %s instances" %
 147                              self.__class__.__name__)
 148
 149     def _checkClosed(self, msg=None):
 150         # raise an exception here if you wish to check for spurious closes
 151         pass
 152
 153     def read(self, len=0, buffer=None):
 154         """Read up to LEN bytes and return them.
 155         Return zero-length string on EOF."""
 156
 157         self._checkClosed()
 158         try:
 159             if buffer:
 160                 v = self._sslobj.read(buffer, len)
 161             else:
 162                 v = self._sslobj.read(len or 1024)
 163             return v
 164         except SSLError as x:
 165             if x.args[0] == SSL_ERROR_EOF and self.suppress_ragged_eofs:
 166                 if buffer:
 167                     return 0
 168                 else:
 169                     return b''
 170             else:
 171                 raise
 172
 173     def write(self, data):
 174         """Write DATA to the underlying SSL channel.  Returns
 175         number of bytes of DATA actually transmitted."""
 176
 177         self._checkClosed()
 178         return self._sslobj.write(data)
 179
 180     def getpeercert(self, binary_form=False):
 181         """Returns a formatted version of the data in the
 182         certificate provided by the other end of the SSL channel.
 183         Return None if no certificate was provided, {} if a
 184         certificate was provided, but not validated."""
 185
 186         self._checkClosed()
 187         return self._sslobj.peer_certificate(binary_form)
 188
 189     def cipher(self):
 190         self._checkClosed()
 191         if not self._sslobj:
 192             return None
 193         else:
 194             return self._sslobj.cipher()
 195
 196     def send(self, data, flags=0):
 197         self._checkClosed()
 198         if self._sslobj:
 199             if flags != 0:
 200                 raise ValueError(
 201                     "non-zero flags not allowed in calls to send() on %s" %
 202                     self.__class__)
 203             while True:
 204                 try:
 205                     v = self._sslobj.write(data)
 206                 except SSLError as x:
 207                     if x.args[0] == SSL_ERROR_WANT_READ:
 208                         return 0
 209                     elif x.args[0] == SSL_ERROR_WANT_WRITE:
 210                         return 0
 211                     else:
 212                         raise
 213                 else:
 214                     return v
 215         else:
 216             return socket.send(self, data, flags)
 217
 218     def sendto(self, data, addr, flags=0):
 219         self._checkClosed()
 220         if self._sslobj:
 221             raise ValueError("sendto not allowed on instances of %s" %
 222                              self.__class__)
 223         else:
 224             return socket.sendto(self, data, addr, flags)
 225
 226     def sendall(self, data, flags=0):
 227         self._checkClosed()
 228         if self._sslobj:
 229             amount = len(data)
 230             count = 0
 231             while (count < amount):
 232                 v = self.send(data[count:])
 233                 count += v
 234             return amount
 235         else:
 236             return socket.sendall(self, data, flags)
 237
 238     def recv(self, buflen=1024, flags=0):
 239         self._checkClosed()
 240         if self._sslobj:
 241             if flags != 0:
 242                 raise ValueError(
 243                   "non-zero flags not allowed in calls to recv_into() on %s" %
 244                   self.__class__)
 245             while True:
 246                 try:
 247                     return self.read(buflen)
 248                 except SSLError as x:
 249                     if x.args[0] == SSL_ERROR_WANT_READ:
 250                         continue
 251                     else:
 252                         raise x
 253         else:
 254             return socket.recv(self, buflen, flags)
 255
 256     def recv_into(self, buffer, nbytes=None, flags=0):
 257         self._checkClosed()
 258         if buffer and (nbytes is None):
 259             nbytes = len(buffer)
 260         elif nbytes is None:
 261             nbytes = 1024
 262         if self._sslobj:
 263             if flags != 0:
 264                 raise ValueError(
 265                   "non-zero flags not allowed in calls to recv_into() on %s" %
 266                   self.__class__)
 267             while True:
 268                 try:
 269                     v = self.read(nbytes, buffer)
 270                     return v
 271                 except SSLError as x:
 272                     if x.args[0] == SSL_ERROR_WANT_READ:
 273                         continue
 274                     else:
 275                         raise x
 276         else:
 277             return socket.recv_into(self, buffer, nbytes, flags)
 278
 279     def recvfrom(self, addr, buflen=1024, flags=0):
 280         self._checkClosed()
 281         if self._sslobj:
 282             raise ValueError("recvfrom not allowed on instances of %s" %
 283                              self.__class__)
 284         else:
 285             return socket.recvfrom(self, addr, buflen, flags)
 286
 287     def recvfrom_into(self, buffer, nbytes=None, flags=0):
 288         self._checkClosed()
 289         if self._sslobj:
 290             raise ValueError("recvfrom_into not allowed on instances of %s" %
 291                              self.__class__)
 292         else:
 293             return socket.recvfrom_into(self, buffer, nbytes, flags)
 294
 295     def pending(self):
 296         self._checkClosed()
 297         if self._sslobj:
 298             return self._sslobj.pending()
 299         else:
 300             return 0
 301
 302     def shutdown(self, how):
 303         self._checkClosed()
 304         self._sslobj = None
 305         socket.shutdown(self, how)
 306
 307     def unwrap (self):
 308         if self._sslobj:
 309             s = self._sslobj.shutdown()
 310             self._sslobj = None
 311             return s
 312         else:
 313             raise ValueError("No SSL wrapper around " + str(self))
 314
 315     def _real_close(self):
 316         self._sslobj = None
 317         # self._closed = True
 318         socket._real_close(self)
 319
 320     def do_handshake(self, block=False):
 321         """Perform a TLS/SSL handshake."""
 322
 323         timeout = self.gettimeout()
 324         try:
 325             if timeout == 0.0 and block:
 326                 self.settimeout(None)
 327             self._sslobj.do_handshake()
 328         finally:
 329             self.settimeout(timeout)
 330
 331     def connect(self, addr):
 332         """Connects to remote ADDR, and then wraps the connection in
 333         an SSL channel."""
 334
 335         # Here we assume that the socket is client-side, and not
 336         # connected at the time of the call.  We connect it, then wrap it.
 337         if self._sslobj:
 338             raise ValueError("attempt to connect already-connected SSLSocket!")
 339         socket.connect(self, addr)
 340         self._sslobj = _ssl.sslwrap(self, False, self.keyfile, self.certfile,
 341                                     self.cert_reqs, self.ssl_version,
 342                                     self.ca_certs)
 343         try:
 344             if self.do_handshake_on_connect:
 345                 self.do_handshake()
 346         except:
 347             self._sslobj = None
 348             raise
 349
 350     def accept(self):
 351         """Accepts a new connection from a remote client, and returns
 352         a tuple containing that new connection wrapped with a server-side
 353         SSL channel, and the address of the remote client."""
 354
 355         newsock, addr = socket.accept(self)
 356         return (SSLSocket(sock=newsock,
 357                           keyfile=self.keyfile, certfile=self.certfile,
 358                           server_side=True,
 359                           cert_reqs=self.cert_reqs,
 360                           ssl_version=self.ssl_version,
 361                           ca_certs=self.ca_certs,
 362                           do_handshake_on_connect=
 363                               self.do_handshake_on_connect),
 364                 addr)
 365
 366     def __del__(self):
 367         # sys.stderr.write("__del__ on %s\n" % repr(self))
 368         self._real_close()
 369
 370
 371 def wrap_socket(sock, keyfile=None, certfile=None,
 372                 server_side=False, cert_reqs=CERT_NONE,
 373                 ssl_version=PROTOCOL_SSLv23, ca_certs=None,
 374                 do_handshake_on_connect=True,
 375                 suppress_ragged_eofs=True):
 376
 377     return SSLSocket(sock=sock, keyfile=keyfile, certfile=certfile,
 378                      server_side=server_side, cert_reqs=cert_reqs,
 379                      ssl_version=ssl_version, ca_certs=ca_certs,
 380                      do_handshake_on_connect=do_handshake_on_connect,
 381                      suppress_ragged_eofs=suppress_ragged_eofs)
 382
 383 # some utility functions
 384
 385 def cert_time_to_seconds(cert_time):
 386     """Takes a date-time string in standard ASN1_print form
 387     ("MON DAY 24HOUR:MINUTE:SEC YEAR TIMEZONE") and return
 388     a Python time value in seconds past the epoch."""
 389
 390     import time
 391     return time.mktime(time.strptime(cert_time, "%b %d %H:%M:%S %Y GMT"))
 392
 393 PEM_HEADER = "-----BEGIN CERTIFICATE-----"
 394 PEM_FOOTER = "-----END CERTIFICATE-----"
 395
 396 def DER_cert_to_PEM_cert(der_cert_bytes):
 397     """Takes a certificate in binary DER format and returns the
 398     PEM version of it as a string."""
 399
 400     f = str(base64.standard_b64encode(der_cert_bytes), 'ASCII', 'strict')
 401     return (PEM_HEADER + '\n' +
 402             textwrap.fill(f, 64) + '\n' +
 403             PEM_FOOTER + '\n')
 404
 405 def PEM_cert_to_DER_cert(pem_cert_string):
 406     """Takes a certificate in ASCII PEM format and returns the
 407     DER-encoded version of it as a byte sequence"""
 408
 409     if not pem_cert_string.startswith(PEM_HEADER):
 410         raise ValueError("Invalid PEM encoding; must start with %s"
 411                          % PEM_HEADER)
 412     if not pem_cert_string.strip().endswith(PEM_FOOTER):
 413         raise ValueError("Invalid PEM encoding; must end with %s"
 414                          % PEM_FOOTER)
 415     d = pem_cert_string.strip()[len(PEM_HEADER):-len(PEM_FOOTER)]
 416     return base64.decodebytes(d.encode('ASCII', 'strict'))
 417
 418 def get_server_certificate(addr, ssl_version=PROTOCOL_SSLv3, ca_certs=None):
 419     """Retrieve the certificate from the server at the specified address,
 420     and return it as a PEM-encoded string.
 421     If 'ca_certs' is specified, validate the server cert against it.
 422     If 'ssl_version' is specified, use it in the connection attempt."""
 423
 424     host, port = addr
 425     if (ca_certs is not None):
 426         cert_reqs = CERT_REQUIRED
 427     else:
 428         cert_reqs = CERT_NONE
 429     s = wrap_socket(socket(), ssl_version=ssl_version,
 430                     cert_reqs=cert_reqs, ca_certs=ca_certs)
 431     s.connect(addr)
 432     dercert = s.getpeercert(True)
 433     s.close()
 434     return DER_cert_to_PEM_cert(dercert)
 435
 436 def get_protocol_name(protocol_code):
 437     if protocol_code == PROTOCOL_TLSv1:
 438         return "TLSv1"
 439     elif protocol_code == PROTOCOL_SSLv23:
 440         return "SSLv23"
 441     elif protocol_code == PROTOCOL_SSLv2:
 442         return "SSLv2"
 443     elif protocol_code == PROTOCOL_SSLv3:
 444         return "SSLv3"
 445     else:
 446         return "<unknown>"