Connecting to a Found cluster in Python with Thrift and SSL

found_thrift_client.py

import os
import json
import base64

from elasticsearch.connection.thrift import *

from elasticsearch import Elasticsearch
from elasticsearch.connection import thrift
from elasticsearch.connection.esthrift import Rest

class FoundTSSLSocket(TSSLSocket.TSSLSocket):
    """Added support for wildcard certificates."""

    def _match_common_name(self, host, certificate):
        if host == certificate:
            return True
        hostParts = host.split('.')
        certificateParts = certificate.split('.')
        if (len(hostParts) != len(certificateParts)):
            return False
        for hostPart, certificatePart in zip(hostParts, certificateParts):
            if not (hostPart == certificatePart or certificatePart == '*'):
                return False
        return True

    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
          if self._match_common_name(self.host, certhost):
            # 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 FoundThriftConnection(thrift.ThriftConnection):
    """ A custom thrift connection class that allows setting some default headers for all requests and supports wildcard certificates when enablings ssl. """

    def __init__(self, default_thrift_headers=None, ca_certs=None, authentication=None, *a, **kw):
        super(FoundThriftConnection, self).__init__(*a, **kw)
        self.default_thrift_headers = default_thrift_headers or {'X-Found-Cluster-Name': kw['host'], 'Authorization': "Basic " + base64.b64encode(authentication)}
        if kw['use_ssl']:
            self._tsocket_args = (kw['host'], kw['port'], True, ca_certs)
            self._tsocket_class = FoundTSSLSocket
    
    def _make_connection(self):
        """ Overrides the client returned. """
        socket = self._tsocket_class(*self._tsocket_args)
        socket.setTimeout(self.timeout * 1000.0)
        if self._framed_transport:
            transport = TTransport.TFramedTransport(socket)
        else:
            transport = TTransport.TBufferedTransport(socket)

        protocol = TBinaryProtocol.TBinaryProtocolAccelerated(transport)
        client = FoundThriftClient(self.default_thrift_headers, protocol)
        transport.open()
        return client

class FoundThriftClient(Rest.Client):

    def __init__(self, default_headers, *a, **kw):
        super(FoundThriftClient, self).__init__(*a, **kw)
        self.default_headers = default_headers

    def execute(self, request):
        headers = self.default_headers.copy()
        headers.update(request.headers or {})
        request.headers = headers
        return super(FoundThriftClient, self).execute(request)

test.py

from found_thrift_client import *

es = Elasticsearch(connection_class=FoundThriftConnection, host='myclusterid.foundcluster.com', port=9543, framed_transport=True, use_ssl=True, ca_certs='/Users/beiske/Downloads/cacert.pem', authentication="username:password")
res = es.search(body={"query": { "match_all": { } }})
print res