lib/py/src/transport/TSSLSocket.py

#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
#

import logging
import os
import socket
import ssl
import sys
import warnings

from thrift.transport import TSocket
from thrift.transport.TTransport import TTransportException

logger = logging.getLogger(__name__)
warnings.filterwarnings('default', category=DeprecationWarning, module=__name__)


class TSSLBase(object):
  # SSLContext is not available for Python < 2.7.9
  _has_ssl_context = sys.hexversion >= 0x020709F0

  # ciphers argument is not available for Python < 2.7.0
  _has_ciphers = sys.hexversion >= 0x020700F0

  # For pythoon >= 2.7.9, use latest TLS that both client and server supports.
  # SSL 2.0 and 3.0 are disabled via ssl.OP_NO_SSLv2 and ssl.OP_NO_SSLv3.
  # For pythoon < 2.7.9, use TLS 1.0 since TLSv1_X nare OP_NO_SSLvX are unavailable.
  _default_protocol = ssl.PROTOCOL_SSLv23 if _has_ssl_context else ssl.PROTOCOL_TLSv1

  def _init_context(self, ssl_version):
    if self._has_ssl_context:
      self._context = ssl.SSLContext(ssl_version)
      if self._context.protocol == ssl.PROTOCOL_SSLv23:
        self._context.options |= ssl.OP_NO_SSLv2
        self._context.options |= ssl.OP_NO_SSLv3
    else:
      self._context = None
      self._ssl_version = ssl_version

  @property
  def ssl_version(self):
    if self._has_ssl_context:
      return self.ssl_context.protocol
    else:
      return self._ssl_version

  @property
  def ssl_context(self):
    return self._context

  SSL_VERSION = _default_protocol
  """
  Default SSL version.
  For backword compatibility, it can be modified.
  Use __init__ keywoard argument "ssl_version" instead.
  """

  def _deprecated_arg(self, args, kwargs, pos, key):
    if len(args) <= pos:
      return
    real_pos = pos + 3
    warnings.warn(
        '%dth positional argument is deprecated. Use keyward argument insteand.' % real_pos,
        DeprecationWarning)
    if key in kwargs:
      raise TypeError('Duplicate argument: %dth argument and %s keyward argument.', (real_pos, key))
    kwargs[key] = args[pos]

  def _unix_socket_arg(self, host, port, args, kwargs):
    key = 'unix_socket'
    if host is None and port is None and len(args) == 1 and key not in kwargs:
      kwargs[key] = args[0]
      return True
    return False

  def __getattr__(self, key):
    if key == 'SSL_VERSION':
      warnings.warn('Use ssl_version attribute instead.', DeprecationWarning)
      return self.ssl_version

  def __init__(self, server_side, host, ssl_opts):
    self._server_side = server_side
    if TSSLBase.SSL_VERSION != self._default_protocol:
      warnings.warn('SSL_VERSION is deprecated. Use ssl_version keyward argument instead.', DeprecationWarning)
    self._context = ssl_opts.pop('ssl_context', None)
    self._server_hostname = None
    if not self._server_side:
      self._server_hostname = ssl_opts.pop('server_hostname', host)
    if self._context:
      self._custom_context = True
      if ssl_opts:
        raise ValueError('Incompatible arguments: ssl_context and %s' % ' '.join(ssl_opts.keys()))
      if not self._has_ssl_context:
        raise ValueError('ssl_context is not available for this version of Python')
    else:
      self._custom_context = False
      ssl_version = ssl_opts.pop('ssl_version', TSSLBase.SSL_VERSION)
      self._init_context(ssl_version)
      self.cert_reqs = ssl_opts.pop('cert_reqs', ssl.CERT_REQUIRED)
      self.ca_certs = ssl_opts.pop('ca_certs', None)
      self.keyfile = ssl_opts.pop('keyfile', None)
      self.certfile = ssl_opts.pop('certfile', None)
      self.ciphers = ssl_opts.pop('ciphers', None)

      if ssl_opts:
        raise ValueError('Unknown keyword arguments: ', ' '.join(ssl_opts.keys()))

      if self.cert_reqs != ssl.CERT_NONE:
        if not self.ca_certs:
          raise ValueError('ca_certs is needed when cert_reqs is not ssl.CERT_NONE')
        if not os.access(self.ca_certs, os.R_OK):
          raise IOError('Certificate Authority ca_certs file "%s" '
                        'is not readable, cannot validate SSL '
                        'certificates.' % (self.ca_certs))

  @property
  def certfile(self):
    return self._certfile

  @certfile.setter
  def certfile(self, certfile):
    if self._server_side and not certfile:
      raise ValueError('certfile is needed for server-side')
    if certfile and not os.access(certfile, os.R_OK):
      raise IOError('No such certfile found: %s' % (certfile))
    self._certfile = certfile

  def _wrap_socket(self, sock):
    if self._has_ssl_context:
      if not self._custom_context:
        self.ssl_context.verify_mode = self.cert_reqs
        if self.certfile:
          self.ssl_context.load_cert_chain(self.certfile, self.keyfile)
        if self.ciphers:
          self.ssl_context.set_ciphers(self.ciphers)
        if self.ca_certs:
          self.ssl_context.load_verify_locations(self.ca_certs)
      return self.ssl_context.wrap_socket(sock, server_side=self._server_side,
                                          server_hostname=self._server_hostname)
    else:
      ssl_opts = {
        'ssl_version': self._ssl_version,
        'server_side': self._server_side,
        'ca_certs': self.ca_certs,
        'keyfile': self.keyfile,
        'certfile': self.certfile,
        'cert_reqs': self.cert_reqs,
      }
      if self.ciphers:
        if self._has_ciphers:
          ssl_opts['ciphers'] = self.ciphers
        else:
          logger.warning('ciphers is specified but ignored due to old Python version')
      return ssl.wrap_socket(sock, **ssl_opts)


class TSSLSocket(TSocket.TSocket, TSSLBase):
  """
  SSL implementation of TSocket

  This class creates outbound sockets wrapped using the
  python standard ssl module for encrypted connections.
  """

  # New signature
  # def __init__(self, host='localhost', port=9090, unix_socket=None, **ssl_args):
  # Deprecated signature
  # def __init__(self, host='localhost', port=9090, validate=True, ca_certs=None, keyfile=None, certfile=None, unix_socket=None, ciphers=None):
  def __init__(self, host='localhost', port=9090, *args, **kwargs):
    """Positional arguments: ``host``, ``port``, ``unix_socket``

    Keyword arguments: ``keyfile``, ``certfile``, ``cert_reqs``, ``ssl_version``,
                       ``ca_certs``, ``ciphers`` (Python 2.7.0 or later),
                       ``server_hostname`` (Python 2.7.9 or later)
    Passed to ssl.wrap_socket. See ssl.wrap_socket documentation.

    Alternative keywoard arguments: (Python 2.7.9 or later)
      ``ssl_context``: ssl.SSLContext to be used for SSLContext.wrap_socket
      ``server_hostname``: Passed to SSLContext.wrap_socket
    """
    self.is_valid = False
    self.peercert = None

    if args:
      if len(args) > 6:
        raise TypeError('Too many positional argument')
      if not self._unix_socket_arg(host, port, args, kwargs):
        self._deprecated_arg(args, kwargs, 0, 'validate')
      self._deprecated_arg(args, kwargs, 1, 'ca_certs')
      self._deprecated_arg(args, kwargs, 2, 'keyfile')
      self._deprecated_arg(args, kwargs, 3, 'certfile')
      self._deprecated_arg(args, kwargs, 4, 'unix_socket')
      self._deprecated_arg(args, kwargs, 5, 'ciphers')

    validate = kwargs.pop('validate', None)
    if validate is not None:
      cert_reqs_name = 'CERT_REQUIRED' if validate else 'CERT_NONE'
      warnings.warn(
          'validate is deprecated. Use cert_reqs=ssl.%s instead' % cert_reqs_name,
          DeprecationWarning)
      if 'cert_reqs' in kwargs:
        raise TypeError('Cannot specify both validate and cert_reqs')
      kwargs['cert_reqs'] = ssl.CERT_REQUIRED if validate else ssl.CERT_NONE

    unix_socket = kwargs.pop('unix_socket', None)
    TSSLBase.__init__(self, False, host, kwargs)
    TSocket.TSocket.__init__(self, host, port, unix_socket)

  @property
  def validate(self):
    warnings.warn('Use cert_reqs instead', DeprecationWarning)
    return self.cert_reqs != ssl.CERT_NONE

  @validate.setter
  def validate(self, value):
    warnings.warn('Use cert_reqs instead', DeprecationWarning)
    self.cert_reqs = ssl.CERT_REQUIRED if value else ssl.CERT_NONE

  def open(self):
    try:
      res0 = self._resolveAddr()
      for res in res0:
        sock_family, sock_type = res[0:2]
        ip_port = res[4]
        plain_sock = socket.socket(sock_family, sock_type)
        self.handle = self._wrap_socket(plain_sock)
        self.handle.settimeout(self._timeout)
        try:
          self.handle.connect(ip_port)
        except socket.error as e:
          if res is not res0[-1]:
            logger.warning('Error while connecting with %s. Trying next one.', ip_port, exc_info=True)
            continue
          else:
            raise
        break
    except socket.error as e:
      if self._unix_socket:
        message = 'Could not connect to secure socket %s: %s' \
                  % (self._unix_socket, e)
      else:
        message = 'Could not connect to %s:%d: %s' % (self.host, self.port, e)
      logger.error('Error while connecting with %s.', ip_port, exc_info=True)
      raise TTransportException(type=TTransportException.NOT_OPEN,
                                message=message)
    if self.validate:
      self._validate_cert()

  def _validate_cert(self):
    """internal method to validate the peer's SSL certificate, and to check the
    commonName of the certificate to ensure it matches the hostname we
    used to make this connection.  Does not support subjectAltName records
    in certificates.

    raises TTransportException if the certificate fails validation.
    """
    cert = self.handle.getpeercert()
    self.peercert = cert
    if 'subject' not in cert:
      raise TTransportException(
        type=TTransportException.NOT_OPEN,
        message='No SSL certificate found from %s:%s' % (self.host, self.port))
    fields = cert['subject']
    for field in fields:
      # ensure structure we get back is what we expect
      if not isinstance(field, tuple):
        continue
      cert_pair = field[0]
      if len(cert_pair) < 2:
        continue
      cert_key, cert_value = cert_pair[0:2]
      if cert_key != 'commonName':
        continue
      certhost = cert_value
      # this check should be performed by some sort of Access Manager
      if certhost == self.host:
        # success, cert commonName matches desired hostname
        self.is_valid = True
        return
      else:
        raise TTransportException(
          type=TTransportException.UNKNOWN,
          message='Hostname we connected to "%s" doesn\'t match certificate '
                  'provided commonName "%s"' % (self.host, certhost))
    raise TTransportException(
      type=TTransportException.UNKNOWN,
      message='Could not validate SSL certificate from '
              'host "%s".  Cert=%s' % (self.host, cert))


class TSSLServerSocket(TSocket.TServerSocket, TSSLBase):
  """SSL implementation of TServerSocket

  This uses the ssl module's wrap_socket() method to provide SSL
  negotiated encryption.
  """

  # New signature
  # def __init__(self, host='localhost', port=9090, unix_socket=None, **ssl_args):
  # Deprecated signature
  # def __init__(self, host=None, port=9090, certfile='cert.pem', unix_socket=None, ciphers=None):
  def __init__(self, host=None, port=9090, *args, **kwargs):
    """Positional arguments: ``host``, ``port``, ``unix_socket``

    Keyword arguments: ``keyfile``, ``certfile``, ``cert_reqs``, ``ssl_version``,
                       ``ca_certs``, ``ciphers`` (Python 2.7.0 or later)
    See ssl.wrap_socket documentation.

    Alternative keywoard arguments: (Python 2.7.9 or later)
      ``ssl_context``: ssl.SSLContext to be used for SSLContext.wrap_socket
      ``server_hostname``: Passed to SSLContext.wrap_socket
    """
    if args:
      if len(args) > 3:
        raise TypeError('Too many positional argument')
      if not self._unix_socket_arg(host, port, args, kwargs):
        self._deprecated_arg(args, kwargs, 0, 'certfile')
      self._deprecated_arg(args, kwargs, 1, 'unix_socket')
      self._deprecated_arg(args, kwargs, 2, 'ciphers')

    if 'ssl_context' not in kwargs:
      # Preserve existing behaviors for default values
      if 'cert_reqs' not in kwargs:
        kwargs['cert_reqs'] = ssl.CERT_NONE
      if'certfile' not in kwargs:
        kwargs['certfile'] = 'cert.pem'

    unix_socket = kwargs.pop('unix_socket', None)
    TSSLBase.__init__(self, True, None, kwargs)
    TSocket.TServerSocket.__init__(self, host, port, unix_socket)

  def setCertfile(self, certfile):
    """Set or change the server certificate file used to wrap new connections.

    @param certfile: The filename of the server certificate,
                     i.e. '/etc/certs/server.pem'
    @type certfile: str

    Raises an IOError exception if the certfile is not present or unreadable.
    """
    warnings.warn('Use certfile property instead.', DeprecationWarning)
    self.certfile = certfile

  def accept(self):
    plain_client, addr = self.handle.accept()
    try:
      client = self._wrap_socket(plain_client)
    except ssl.SSLError:
      logger.error('Error while accepting from %s', addr, exc_info=True)
      # failed handshake/ssl wrap, close socket to client
      plain_client.close()
      # raise
      # We can't raise the exception, because it kills most TServer derived
      # serve() methods.
      # Instead, return None, and let the TServer instance deal with it in
      # other exception handling.  (but TSimpleServer dies anyway)
      return None
    result = TSocket.TSocket()
    result.setHandle(client)
    return result