1 # Copyright 2010, Google Inc.

     2 # All rights reserved.

     3 #

     4 # Redistribution and use in source and binary forms, with or without

     5 # modification, are permitted provided that the following conditions are

     6 # met:

     7 #

     8 #     * Redistributions of source code must retain the above copyright

     9 # notice, this list of conditions and the following disclaimer.

    10 #     * Redistributions in binary form must reproduce the above

    11 # copyright notice, this list of conditions and the following disclaimer

    12 # in the documentation and/or other materials provided with the

    13 # distribution.

    14 #     * Neither the name of Google Inc. nor the names of its

    15 # contributors may be used to endorse or promote products derived from

    16 # this software without specific prior written permission.

    17 

    18 # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS

    19 # "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT

    20 # LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR

    21 # A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT

    22 # OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,

    23 # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT

    24 # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,

    25 # DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY

    26 # THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT

    27 # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE

    28 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

    29 """Abstraction to simplify socket use for Python < 2.6

    30 

    31 This will attempt to use the ssl module and the new

    32 socket.create_connection method, but fall back to the old

    33 methods if those are unavailable.

    34 """

    35 import logging

    36 import socket

    37 

    38 logger = logging.getLogger(__name__)

    39 

    40 try:

    41     import ssl

    42     ssl.wrap_socket  # make demandimporters load the module

    43     have_ssl = True

    44 except ImportError:

    45     import httplib

    46     import urllib2

    47     have_ssl = getattr(urllib2, 'HTTPSHandler', False)

    48     ssl = False

    49 

    50 

    51 try:

    52     create_connection = socket.create_connection

    53 except AttributeError:

    54     def create_connection(address):

    55         host, port = address

    56         msg = "getaddrinfo returns an empty list"

    57         sock = None

    58         for res in socket.getaddrinfo(host, port, 0,

    59                                       socket.SOCK_STREAM):

    60             af, socktype, proto, _canonname, sa = res

    61             try:

    62                 sock = socket.socket(af, socktype, proto)

    63                 logger.info("connect: (%s, %s)", host, port)

    64                 sock.connect(sa)

    65             except socket.error, msg:

    66                 logger.info('connect fail: %s %s', host, port)

    67                 if sock:

    68                     sock.close()

    69                 sock = None

    70                 continue

    71             break

    72         if not sock:

    73             raise socket.error, msg

    74         return sock

    75 

    76 if ssl:

    77     wrap_socket = ssl.wrap_socket

    78     CERT_NONE = ssl.CERT_NONE

    79     CERT_OPTIONAL = ssl.CERT_OPTIONAL

    80     CERT_REQUIRED = ssl.CERT_REQUIRED

    81     PROTOCOL_SSLv2 = ssl.PROTOCOL_SSLv2

    82     PROTOCOL_SSLv3 = ssl.PROTOCOL_SSLv3

    83     PROTOCOL_SSLv23 = ssl.PROTOCOL_SSLv23

    84     PROTOCOL_TLSv1 = ssl.PROTOCOL_TLSv1

    85 else:

    86     class FakeSocket(httplib.FakeSocket):

    87         """Socket wrapper that supports SSL.

    88         """

    89         # backport the behavior from Python 2.6, which is to busy wait

    90         # on the socket instead of anything nice. Sigh.

    91         # See http://bugs.python.org/issue3890 for more info.

    92         def recv(self, buflen=1024, flags=0):

    93             """ssl-aware wrapper around socket.recv

    94             """

    95             if flags != 0:

    96                 raise ValueError(

    97                     "non-zero flags not allowed in calls to recv() on %s" %

    98                     self.__class__)

    99             while True:

   100                 try:

   101                     return self._ssl.read(buflen)

   102                 except socket.sslerror, x:

   103                     if x.args[0] == socket.SSL_ERROR_WANT_READ:

   104                         continue

   105                     else:

   106                         raise x

   107 

   108     PROTOCOL_SSLv2 = 0

   109     PROTOCOL_SSLv3 = 1

   110     PROTOCOL_SSLv23 = 2

   111     PROTOCOL_TLSv1 = 3

   112 

   113     CERT_NONE = 0

   114     CERT_OPTIONAL = 1

   115     CERT_REQUIRED = 2

   116 

   117     def wrap_socket(sock, keyfile=None, certfile=None,

   118                 server_side=False, cert_reqs=CERT_NONE,

   119                 ssl_version=PROTOCOL_SSLv23, ca_certs=None,

   120                 do_handshake_on_connect=True,

   121                 suppress_ragged_eofs=True):

   122         if cert_reqs != CERT_NONE and ca_certs:

   123             raise CertificateValidationUnsupported(

   124                 'SSL certificate validation requires the ssl module'

   125                 '(included in Python 2.6 and later.)')

   126         sslob = socket.ssl(sock)

   127         # borrow httplib's workaround for no ssl.wrap_socket

   128         sock = FakeSocket(sock, sslob)

   129         return sock

   130 

   131 

   132 class CertificateValidationUnsupported(Exception):

   133     """Exception raised when cert validation is requested but unavailable."""

   134 # no-check-code