def test_deprecation(self):
with warnings.catch_warnings(record=True) as w:
warnings.filterwarnings('always',
category=DeprecationWarning,
module='thrift.*SSL.*')
TSSLSocket('localhost',
TEST_PORT,
validate=True,
ca_certs=SERVER_CERT)
self.assertEqual(len(w), 1)
with warnings.catch_warnings(record=True) as w:
warnings.filterwarnings('always',
category=DeprecationWarning,
module='thrift.*SSL.*')
# Deprecated signature
# def __init__(self, host='localhost', port=9090, validate=True, ca_certs=None, keyfile=None, certfile=None, unix_socket=None, ciphers=None):
client = TSSLSocket('localhost', TEST_PORT, True, SERVER_CERT,
CLIENT_KEY, CLIENT_CERT, None, TEST_CIPHERS)
self.assertEqual(len(w), 7)
with warnings.catch_warnings(record=True) as w:
warnings.filterwarnings('always',
category=DeprecationWarning,
module='thrift.*SSL.*')
# Deprecated signature
# def __init__(self, host=None, port=9090, certfile='cert.pem', unix_socket=None, ciphers=None):
server = TSSLServerSocket(None, TEST_PORT, SERVER_PEM, None,
TEST_CIPHERS)
self.assertEqual(len(w), 3)
self._assert_connection_success(server, client)
def __init__(self, *args, **kw):
# Curse 2.6.x + PEP-3102
self.__delay_handshake = kw.pop('delay_handshake', False)
self.__socket_factory = kw.pop('socket_factory', socket)
if 'unix_socket' in kw:
raise ValueError('%s does not support unix_sockets!' % self.__class__.__name__)
TSSLSocket.__init__(self, *args, **kw) # thrift does not support super()
def test_server_cert(self):
server = TSSLServerSocket(port=TEST_PORT,
keyfile=SERVER_KEY,
certfile=SERVER_CERT)
client = TSSLSocket('localhost',
TEST_PORT,
cert_reqs=ssl.CERT_REQUIRED,
ca_certs=SERVER_CERT)
self._assert_connection_success(server, client)
server = TSSLServerSocket(port=TEST_PORT,
keyfile=SERVER_KEY,
certfile=SERVER_CERT)
# server cert on in ca_certs
client = TSSLSocket('localhost',
TEST_PORT,
cert_reqs=ssl.CERT_REQUIRED,
ca_certs=CLIENT_CERT)
self._assert_connection_failure(server, client)
server = TSSLServerSocket(port=TEST_PORT,
keyfile=SERVER_KEY,
certfile=SERVER_CERT)
client = TSSLSocket('localhost', TEST_PORT, cert_reqs=ssl.CERT_NONE)
self._assert_connection_success(server, client)
def test_ciphers(self):
server = TSSLServerSocket(port=TEST_PORT,
keyfile=SERVER_KEY,
certfile=SERVER_CERT,
ciphers=TEST_CIPHERS)
client = TSSLSocket('localhost',
TEST_PORT,
ca_certs=SERVER_CERT,
ciphers=TEST_CIPHERS)
self._assert_connection_success(server, client)
if not TSSLSocket._has_ciphers:
# unittest.skip is not available for Python 2.6
print('skipping test_ciphers')
return
server = TSSLServerSocket(port=TEST_PORT,
keyfile=SERVER_KEY,
certfile=SERVER_CERT)
client = TSSLSocket('localhost',
TEST_PORT,
ca_certs=SERVER_CERT,
ciphers='NULL')
self._assert_connection_failure(server, client)
server = TSSLServerSocket(port=TEST_PORT,
keyfile=SERVER_KEY,
certfile=SERVER_CERT,
ciphers=TEST_CIPHERS)
client = TSSLSocket('localhost',
TEST_PORT,
ca_certs=SERVER_CERT,
ciphers='NULL')
self._assert_connection_failure(server, client)
def connect_to_thrift(conf):
"""
Connect to a thrift endpoint as determined by the 'conf' parameter.
Note that this does *not* open the transport.
Returns a tuple of (service, protocol, transport)
"""
if conf.use_ssl:
sock = TSSLSocket(conf.host, conf.port, validate=conf.validate, ca_certs=conf.ca_certs, keyfile=conf.keyfile, certfile=conf.certfile)
else:
sock = TSocket(conf.host, conf.port)
if conf.timeout_seconds:
# Thrift trivia: You can do this after the fact with
# _grab_transport_from_wrapper(self.wrapped.transport).setTimeout(seconds*1000)
sock.setTimeout(conf.timeout_seconds * 1000.0)
if conf.use_sasl:
def sasl_factory():
saslc = sasl.Client()
saslc.setAttr("host", str(conf.host))
saslc.setAttr("service", str(conf.kerberos_principal))
if conf.mechanism == 'PLAIN':
saslc.setAttr("username", str(conf.username))
saslc.setAttr("password", str(conf.password)) # defaults to hue for a non-empty string unless using ldap
saslc.init()
return saslc
transport = TSaslClientTransport(sasl_factory, conf.mechanism, sock)
elif conf.transport == 'framed':
transport = TFramedTransport(sock)
else:
transport = TBufferedTransport(sock)
protocol = TBinaryProtocol(transport)
service = conf.klass(protocol)
return service, protocol, transport
def test_newer_tls(self):
if not TSSLSocket._has_ssl_context:
# unittest.skip is not available for Python 2.6
print('skipping test_newer_tls')
return
if not hasattr(ssl, 'PROTOCOL_TLSv1_2'):
print('PROTOCOL_TLSv1_2 is not available')
else:
server = TSSLServerSocket(port=TEST_PORT, keyfile=SERVER_KEY, certfile=SERVER_CERT, ssl_version=ssl.PROTOCOL_TLSv1_2)
client = TSSLSocket('localhost', TEST_PORT, ca_certs=SERVER_CERT, ssl_version=ssl.PROTOCOL_TLSv1_2)
self._assert_connection_success(server, client)
if not hasattr(ssl, 'PROTOCOL_TLSv1_1'):
print('PROTOCOL_TLSv1_1 is not available')
else:
server = TSSLServerSocket(port=TEST_PORT, keyfile=SERVER_KEY, certfile=SERVER_CERT, ssl_version=ssl.PROTOCOL_TLSv1_1)
client = TSSLSocket('localhost', TEST_PORT, ca_certs=SERVER_CERT, ssl_version=ssl.PROTOCOL_TLSv1_1)
self._assert_connection_success(server, client)
if not hasattr(ssl, 'PROTOCOL_TLSv1_1') or not hasattr(ssl, 'PROTOCOL_TLSv1_2'):
print('PROTOCOL_TLSv1_1 and/or PROTOCOL_TLSv1_2 is not available')
else:
server = TSSLServerSocket(port=TEST_PORT, keyfile=SERVER_KEY, certfile=SERVER_CERT, ssl_version=ssl.PROTOCOL_TLSv1_2)
client = TSSLSocket('localhost', TEST_PORT, ca_certs=SERVER_CERT, ssl_version=ssl.PROTOCOL_TLSv1_1)
self._assert_connection_failure(server, client)
def test_deprecation(self):
with warnings.catch_warnings(record=True) as w:
warnings.filterwarnings('always',
category=DeprecationWarning,
module=self.__module__)
TSSLSocket('localhost', 0, validate=True, ca_certs=SERVER_CERT)
self.assertEqual(len(w), 1)
with warnings.catch_warnings(record=True) as w:
warnings.filterwarnings('always',
category=DeprecationWarning,
module=self.__module__)
# Deprecated signature
# def __init__(self, host='localhost', port=9090, validate=True, ca_certs=None, keyfile=None, certfile=None, unix_socket=None, ciphers=None):
TSSLSocket('localhost', 0, True, SERVER_CERT, CLIENT_KEY,
CLIENT_CERT, None, TEST_CIPHERS)
self.assertEqual(len(w), 7)
with warnings.catch_warnings(record=True) as w:
warnings.filterwarnings('always',
category=DeprecationWarning,
module=self.__module__)
# Deprecated signature
# def __init__(self, host=None, port=9090, certfile='cert.pem', unix_socket=None, ciphers=None):
TSSLServerSocket(None, 0, SERVER_PEM, None, TEST_CIPHERS)
self.assertEqual(len(w), 3)
def test_set_validate_by_cert_reqs(self):
c1 = TSSLSocket('localhost', TEST_PORT, cert_reqs=ssl.CERT_NONE)
self.assertFalse(c1.validate)
c2 = TSSLSocket('localhost', TEST_PORT, cert_reqs=ssl.CERT_REQUIRED, ca_certs=SERVER_CERT)
self.assertTrue(c2.validate)
c3 = TSSLSocket('localhost', TEST_PORT, cert_reqs=ssl.CERT_OPTIONAL, ca_certs=SERVER_CERT)
self.assertTrue(c3.validate)
def test_set_cert_reqs_by_validate(self):
c1 = TSSLSocket('localhost',
TEST_PORT,
validate=True,
ca_certs=SERVER_CERT)
self.assertEqual(c1.cert_reqs, ssl.CERT_REQUIRED)
c1 = TSSLSocket('localhost', TEST_PORT, validate=False)
self.assertEqual(c1.cert_reqs, ssl.CERT_NONE)
def _refresh_thrift_client(self):
"""Refresh the Thrift socket, transport, and client."""
socket = TSSLSocket(self.host, self.port, False, KEY_FILE, KEY_FILE,
KEY_FILE)
if self.timeout is not None:
socket.setTimeout(self.timeout)
self.transport = self._transport_class(socket)
protocol = self._protocol_class(self.transport)
self.client = Hbase.Client(protocol)
def _refresh_thrift_client(self):
"""Refresh the Thrift socket, transport, and client."""
socket = TSSLSocket(self.host, self.port, False, KEY_FILE, KEY_FILE,
KEY_FILE)
if self.timeout is not None:
socket.setTimeout(self.timeout)
self.transport = self._transport_class(socket)
protocol = self._protocol_class(self.transport)
self.client = Hbase.Client(protocol)
def __init__(self, *args, **kw):
# Curse 2.6.x + PEP-3102
self.__delay_handshake = kw.pop('delay_handshake', False)
self.__socket_factory = kw.pop('socket_factory', socket)
if 'unix_socket' in kw:
raise ValueError('%s does not support unix_sockets!' %
self.__class__.__name__)
TSSLSocket.__init__(self, *args,
**kw) # thrift does not support super()
def _get_socket(host, port, use_ssl, ca_cert):
if use_ssl:
from thrift.transport.TSSLSocket import TSSLSocket
if ca_cert is None:
return TSSLSocket(host, port, validate=False)
else:
return TSSLSocket(host, port, validate=True, ca_certs=ca_cert)
else:
return TSocket(host, port)
def transport_factory(host, port, username, password, **kwargs):
"""
Creates a thrift_sasl transport for use with HIVE/PyHive. Only pass this object
to PyHive connection.
:param host: str - host name
:param port: int/str - port number - hive default is 10000
:param username: str -
:param password: str
:param kwargs: optional
use_ssl=True will use a SSL socket with validate=False, default is False
socket_kwargs={}, pass custom kwargs to SSL socket
use_sasl=optionally select to use sasl library instead of PureSASL
:return:
"""
sasl_auth = 'PLAIN'
use_ssl = kwargs.get('use_ssl', False)
socket_kwargs = kwargs.get('socket_kwargs', {})
kerberos_service_name = kwargs.get('kerberos_service_name', None)
use_sasl = kwargs.get('use_sasl', False)
if use_ssl:
if socket_kwargs:
socket = TSSLSocket(host, port, **socket_kwargs)
else:
socket = TSSLSocket(host, port, cert_reqs=ssl.CERT_NONE)
else:
socket = TSocket(host, port) # basic socket
if use_sasl:
import sasl
def sasl_factory():
sasl_client = sasl.Client()
sasl_client.setAttr('host', host)
if sasl_auth == 'GSSAPI':
sasl_client.setAttr('service', kerberos_service_name)
elif sasl_auth == 'PLAIN':
sasl_client.setAttr('username', username)
sasl_client.setAttr('password', password)
else:
raise AssertionError
sasl_client.init()
return sasl_client
else:
from sasl_compat import PureSASLClient
def sasl_factory():
return PureSASLClient(host,
username=username,
password=password,
service=kerberos_service_name,
mechanism=sasl_auth)
transport = thrift_sasl.TSaslClientTransport(sasl_factory, sasl_auth,
socket)
return transport
def test_set_cert_reqs_by_validate(self):
with warnings.catch_warnings(record=True) as w:
warnings.filterwarnings('always', category=DeprecationWarning, module=self.__module__)
c1 = TSSLSocket('localhost', 0, validate=True, ca_certs=SERVER_CERT)
self.assertEqual(c1.cert_reqs, ssl.CERT_REQUIRED)
c1 = TSSLSocket('localhost', 0, validate=False)
self.assertEqual(c1.cert_reqs, ssl.CERT_NONE)
self.assertEqual(len(w), 2)
def connect_to_thrift(conf):
"""
Connect to a thrift endpoint as determined by the 'conf' parameter.
Note that this does *not* open the transport.
Returns a tuple of (service, protocol, transport)
"""
if conf.use_ssl:
sock = TSSLSocket(conf.host,
conf.port,
validate=conf.validate,
ca_certs=conf.ca_certs,
keyfile=conf.keyfile,
certfile=conf.certfile)
else:
sock = TSocket(conf.host, conf.port)
if conf.timeout_seconds:
# Thrift trivia: You can do this after the fact with
# _grab_transport_from_wrapper(self.wrapped.transport).setTimeout(seconds*1000)
sock.setTimeout(conf.timeout_seconds * 1000.0)
if conf.transport_mode == 'http':
sock = THttpClient(conf.http_url, cert_validate=conf.validate)
if conf.use_sasl and conf.mechanism != 'PLAIN':
sock.set_kerberos_auth()
else:
sock.set_basic_auth(conf.username, conf.password)
transport = TBufferedTransport(sock)
elif conf.use_sasl:
def sasl_factory():
saslc = sasl.Client()
saslc.setAttr("host", str(conf.host))
saslc.setAttr("service", str(conf.kerberos_principal))
if conf.mechanism == 'PLAIN':
saslc.setAttr("username", str(conf.username))
saslc.setAttr(
"password", str(conf.password)
) # defaults to hue for a non-empty string unless using ldap
saslc.init()
return saslc
transport = TSaslClientTransport(sasl_factory, conf.mechanism, sock)
elif conf.transport_mode == 'framed':
transport = TFramedTransport(sock)
else:
transport = TBufferedTransport(sock)
protocol = TBinaryProtocol(transport)
if conf.multiple:
protocol = TMultiplexedProtocol(protocol, conf.service_name)
service = conf.klass(protocol)
return service, protocol, transport
def get_socket(host, port, use_ssl, ca_cert):
# based on the Impala shell impl
log.debug('get_socket: host=%s port=%s use_ssl=%s ca_cert=%s', host, port,
use_ssl, ca_cert)
if use_ssl:
from thrift.transport.TSSLSocket import TSSLSocket
if ca_cert is None:
return TSSLSocket(host, port, validate=False)
else:
return TSSLSocket(host, port, validate=True, ca_certs=ca_cert)
else:
return TSocket(host, port)
def test_set_validate_by_cert_reqs(self):
with warnings.catch_warnings(record=True) as w:
warnings.filterwarnings('always', category=DeprecationWarning, module=self.__module__)
c1 = TSSLSocket('localhost', 0, cert_reqs=ssl.CERT_NONE)
self.assertFalse(c1.validate)
c2 = TSSLSocket('localhost', 0, cert_reqs=ssl.CERT_REQUIRED, ca_certs=SERVER_CERT)
self.assertTrue(c2.validate)
c3 = TSSLSocket('localhost', 0, cert_reqs=ssl.CERT_OPTIONAL, ca_certs=SERVER_CERT)
self.assertTrue(c3.validate)
self.assertEqual(len(w), 3)
def __init__(self,
host=None,
port=10000,
ssl=True,
authMechanism=None,
user=None,
password=None,
database=None,
configuration=None,
timeout=None):
authMechanisms = set(['NOSASL', 'PLAIN', 'KERBEROS', 'LDAP'])
if authMechanism not in authMechanisms:
raise NotImplementedError(
'authMechanism is either not supported or not implemented')
#Must set a password for thrift, even if it doesn't need one
#Open issue with python-sasl
if authMechanism == 'PLAIN' and (password is None
or len(password) == 0):
password = '******'
socket = TSSLSocket(host, port, False) if ssl else TSocket(host, port)
socket.setTimeout(timeout)
if authMechanism == 'NOSASL':
transport = TBufferedTransport(socket)
else:
sasl_mech = 'PLAIN'
saslc = sasl.Client()
saslc.setAttr("username", user)
saslc.setAttr("password", password)
if authMechanism == 'KERBEROS':
krb_host, krb_service = self._get_krb_settings(
host, configuration)
sasl_mech = 'GSSAPI'
saslc.setAttr("host", krb_host)
saslc.setAttr("service", krb_service)
saslc.init()
transport = TSaslClientTransport(saslc, sasl_mech, socket)
self.client = TCLIService.Client(TBinaryProtocol(transport))
transport.open()
res = self.client.OpenSession(
TOpenSessionReq(username=user,
password=password,
configuration=configuration))
self.session = res.sessionHandle
if database is not None:
with self.cursor() as cur:
query = "USE {0}".format(database)
cur.execute(query)
def test_client_cert(self):
server = TSSLServerSocket(port=TEST_PORT,
cert_reqs=ssl.CERT_REQUIRED,
keyfile=SERVER_KEY,
certfile=SERVER_CERT,
ca_certs=CLIENT_CERT)
client = TSSLSocket('localhost',
TEST_PORT,
cert_reqs=ssl.CERT_NONE,
certfile=SERVER_CERT,
keyfile=SERVER_KEY)
self._assert_connection_failure(server, client)
server = TSSLServerSocket(port=TEST_PORT,
cert_reqs=ssl.CERT_REQUIRED,
keyfile=SERVER_KEY,
certfile=SERVER_CERT,
ca_certs=CLIENT_CA)
client = TSSLSocket('localhost',
TEST_PORT,
cert_reqs=ssl.CERT_NONE,
certfile=CLIENT_CERT_NO_IP,
keyfile=CLIENT_KEY_NO_IP)
self._assert_connection_failure(server, client)
server = TSSLServerSocket(port=TEST_PORT,
cert_reqs=ssl.CERT_REQUIRED,
keyfile=SERVER_KEY,
certfile=SERVER_CERT,
ca_certs=CLIENT_CA)
client = TSSLSocket('localhost',
TEST_PORT,
cert_reqs=ssl.CERT_NONE,
certfile=CLIENT_CERT,
keyfile=CLIENT_KEY)
self._assert_connection_success(server, client)
server = TSSLServerSocket(port=TEST_PORT,
cert_reqs=ssl.CERT_OPTIONAL,
keyfile=SERVER_KEY,
certfile=SERVER_CERT,
ca_certs=CLIENT_CA)
client = TSSLSocket('localhost',
TEST_PORT,
cert_reqs=ssl.CERT_NONE,
certfile=CLIENT_CERT,
keyfile=CLIENT_KEY)
self._assert_connection_success(server, client)
def test_unix_domain_socket(self):
if platform.system() == 'Windows':
print('skipping test_unix_domain_socket')
return
server = TSSLServerSocket(unix_socket=TEST_ADDR, keyfile=SERVER_KEY, certfile=SERVER_CERT)
client = TSSLSocket(None, None, TEST_ADDR, cert_reqs=ssl.CERT_NONE)
self._assert_connection_success(server, client)
def get_client(saas, host=SAAS_HOST, port=SAAS_PORT):
sock = TSSLSocket(host, port, ca_certs=None, validate=False)
# sock = TSocket(SAAS_HOST, SAAS_PORT)
# transport = TTransport.TBufferedTransport(sock)
transport = TTransport.TFramedTransport(sock)
protocol = TBinaryProtocol.TBinaryProtocol(transport)
client = saas.Client(protocol)
return transport, client
def test_set_server_cert(self):
server = TSSLServerSocket(port=TEST_PORT, keyfile=SERVER_KEY, certfile=CLIENT_CERT)
with self._assert_raises(Exception):
server.certfile = 'foo'
with self._assert_raises(Exception):
server.certfile = None
server.certfile = SERVER_CERT
client = TSSLSocket('localhost', TEST_PORT, cert_reqs=ssl.CERT_REQUIRED, ca_certs=SERVER_CERT)
self._assert_connection_success(server, client)
def _connectable_client(self, server, expect_failure=False, path=None, **client_kwargs):
acc = ServerAcceptor(server, expect_failure)
try:
acc.start()
acc.await_listening()
host, port = ('localhost', acc.port) if path is None else (None, None)
client = TSSLSocket(host, port, unix_socket=path, **client_kwargs)
yield acc, client
finally:
acc.close()
def test_ssl2_and_ssl3_disabled(self):
if not hasattr(ssl, 'PROTOCOL_SSLv3'):
print('PROTOCOL_SSLv3 is not available')
else:
server = TSSLServerSocket(port=TEST_PORT, keyfile=SERVER_KEY, certfile=SERVER_CERT)
client = TSSLSocket('localhost', TEST_PORT, ca_certs=SERVER_CERT, ssl_version=ssl.PROTOCOL_SSLv3)
self._assert_connection_failure(server, client)
server = TSSLServerSocket(port=TEST_PORT, keyfile=SERVER_KEY, certfile=SERVER_CERT, ssl_version=ssl.PROTOCOL_SSLv3)
client = TSSLSocket('localhost', TEST_PORT, ca_certs=SERVER_CERT)
self._assert_connection_failure(server, client)
if not hasattr(ssl, 'PROTOCOL_SSLv2'):
print('PROTOCOL_SSLv2 is not available')
else:
server = TSSLServerSocket(port=TEST_PORT, keyfile=SERVER_KEY, certfile=SERVER_CERT)
client = TSSLSocket('localhost', TEST_PORT, ca_certs=SERVER_CERT, ssl_version=ssl.PROTOCOL_SSLv2)
self._assert_connection_failure(server, client)
server = TSSLServerSocket(port=TEST_PORT, keyfile=SERVER_KEY, certfile=SERVER_CERT, ssl_version=ssl.PROTOCOL_SSLv2)
client = TSSLSocket('localhost', TEST_PORT, ca_certs=SERVER_CERT)
self._assert_connection_failure(server, client)
def connect_to_thrift(conf):
"""
Connect to a thrift endpoint as determined by the 'conf' parameter.
Note that this does *not* open the transport.
Returns a tuple of (service, protocol, transport)
"""
if conf.use_ssl:
sock = TSSLSocket(conf.host,
conf.port,
validate=conf.validate,
ca_certs=conf.ca_certs,
keyfile=conf.keyfile,
certfile=conf.certfile)
else:
sock = TSocket(conf.host, conf.port)
if conf.timeout_seconds:
# Thrift trivia: You can do this after the fact with
# _grab_transport_from_wrapper(self.wrapped.transport).setTimeout(seconds*1000)
sock.setTimeout(conf.timeout_seconds * 1000.0)
if conf.use_sasl:
def sasl_factory():
saslc = sasl.Client()
saslc.setAttr("host", str(conf.host))
saslc.setAttr("service", str(conf.kerberos_principal))
if conf.mechanism == 'PLAIN':
saslc.setAttr("username", str(conf.username))
saslc.setAttr("password", 'hue') # Just a non empty string
saslc.init()
return saslc
transport = TSaslClientTransport(sasl_factory, conf.mechanism, sock)
else:
transport = TBufferedTransport(sock)
protocol = TBinaryProtocol(transport)
service = conf.klass(protocol)
return service, protocol, transport
def test_ssl_context(self):
if not TSSLSocket._has_ssl_context:
# unittest.skip is not available for Python 2.6
print('skipping test_ssl_context')
return
server_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
server_context.load_cert_chain(SERVER_CERT, SERVER_KEY)
server_context.load_verify_locations(CLIENT_CERT)
client_context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
client_context.load_cert_chain(CLIENT_CERT, CLIENT_KEY)
client_context.load_verify_locations(SERVER_CERT)
server = TSSLServerSocket(port=TEST_PORT, ssl_context=server_context)
client = TSSLSocket('localhost', TEST_PORT, ssl_context=client_context)
self._assert_connection_success(server, client)
def build_ssl_transport(host,
port,
username,
auth,
kerberos_service_name,
password=None):
transport = None
if port is None:
port = 10000
if auth is None:
auth = 'NONE'
socket = TSSLSocket(host, port, cert_reqs=ssl.CERT_NONE)
if auth == 'NOSASL':
# NOSASL corresponds to hive.server2.authentication=NOSASL
# in hive-site.xml
transport = thrift.transport.TTransport.TBufferedTransport(socket)
elif auth in ('LDAP', 'KERBEROS', 'NONE', 'CUSTOM'):
# Defer import so package dependency is optional
if auth == 'KERBEROS':
# KERBEROS mode in hive.server2.authentication is GSSAPI
# in sasl library
sasl_auth = 'GSSAPI'
else:
sasl_auth = 'PLAIN'
if password is None:
# Password doesn't matter in NONE mode, just needs
# to be nonempty.
password = '******'
def sasl_factory():
sasl_client = sasl.Client()
sasl_client.setAttr('host', host)
if sasl_auth == 'GSSAPI':
sasl_client.setAttr('service', kerberos_service_name)
elif sasl_auth == 'PLAIN':
sasl_client.setAttr('username', username)
sasl_client.setAttr('password', password)
else:
raise AssertionError
sasl_client.init()
return sasl_client
transport = thrift_sasl.TSaslClientTransport(sasl_factory, sasl_auth,
socket)
return transport
def _rebuild_protocol(self):
"""
Rebuild the transport, protocol from the configuration.
Should not be used directly by users.
Returns:
None
"""
if self.use_http:
# if use http transport,
prefix = 'https://' if self.use_ssl else 'http://'
self.transport = THttpClient(uri_or_host=prefix + self.host + ':' +
str(self.port))
self.protocol = TBinaryProtocol.TBinaryProtocol(self.transport)
return
if self.use_ssl:
from thrift.transport.TSSLSocket import TSSLSocket
socket = TSSLSocket(host=self.host,
port=self.port,
validate=False,
socket_keepalive=self.keep_alive)
else:
socket = TSocket(host=self.host,
port=self.port,
socket_keepalive=self.keep_alive)
if self.authentication:
socket = TSaslClientTransport(
socket,
host=self.host,
service=self.authentication.service,
mechanism=self.authentication.mechanism,
username=self.authentication.username,
password=self.authentication.password,
)
self.transport = self._transport_type(socket)
self.protocol = self._protocol_type(self.transport)
def connect(self, conn_timeout_ms, req_timeout_ms):
"""Connect to the endpoint specified in self.host and self.port.
.. note:: It makes connection only if it's not already connected.
"""
if self.connected:
self._socket.setTimeout(req_timeout_ms)
return
# Socket.
if self.is_ssl:
self._socket = TSSLSocket(self.host, self.port, self.validate,
self.ca_certs)
else:
self._socket = TNoDelaySocket(self.host, self.port)
# Set socket timeout
self._socket.setTimeout(conn_timeout_ms)
# Transport.
self._transport = TBufferedTransport(self._socket)
# open() may throw TTransportException() if fail to connect.
self._transport.open()
# Protocol.
self._protocol = self.protocol_factory.getProtocol(self._transport)
# Client.
# Need to get the parent class of the client class
# Ex: <class
# 'data_clients.follower_service_client.FollowerServiceClient'>
# => <class services.follower.thrift_libs.FollowerService.Client>.
self._client = self.__class__.__bases__[0](self._protocol)
self.connected = True
self.connected_at = time.time()
self._socket.setTimeout(req_timeout_ms)
if side in ("rx", "RX"):
is_tx = False
elif side in ("tx", "TX"):
is_tx = True
else:
print "Wrong side provided"
exit(-1)
is_ok = True
print "Connecting to TX QKD server at {}:{}...".format(tx_host, tx_port)
# Raw sockets are very slow, use buffered transport
transp_tx = TBufferedTransport(TSSLSocket(tx_host,
tx_port,
validate=False,
certfile='ssl/tx_client.crt',
keyfile='ssl/tx_client.key',
ca_certs='ssl/pair_ca_bundle.crt'))
client_tx = Client(TBinaryProtocol(transp_tx))
transp_tx.open()
print "OK: Connected"
print "Connecting to RX QKD server at {}:{}...".format(rx_host, rx_port)
# Raw sockets are very slow, use buffered transport
transp_rx = TBufferedTransport(TSSLSocket(rx_host,
rx_port,
validate=False,
certfile='ssl/rx_client.crt',
keyfile='ssl/rx_client.key',
ca_certs='ssl/pair_ca_bundle.crt'))
client_rx = Client(TBinaryProtocol(transp_rx))
from qkd_client_api.v1.QkdApiService import Client
from qkd_client_api.v1 import ttypes
if __name__ == "__main__":
_, api_host, api_port = sys.argv
api_port = int(api_port)
print "Connecting to QKD server at {}:{}...".format(api_host, api_port)
# Raw sockets are very slow, use buffered transport
transp_tx = TBufferedTransport(
TSSLSocket(api_host,
api_port,
validate=False,
certfile='ssl/tx_client.crt',
keyfile='ssl/tx_client.key',
ca_certs='ssl/pair_ca_bundle.crt'))
client_tx = Client(TBinaryProtocol(transp_tx))
transp_tx.open()
print "Connected"
start_time = now()
total_key_size = 0
i = 0
while 1:
try:
def transport_factory():
# return THttpClient(("https" if ssl else "http") + "://" + client_host + ":" + str(blob_port) + "/thrift")
if ssl:
return TSSLSocket(host=client_host, port=client_port, validate=False)
else:
return TSocket(host=client_host, port=client_port)
