Merge pull request #219 from aiven/alex-tls-cert-key-support

ssl: add support for private key/certificate file
mymarilyn · May 21, 2021 · 7856756 · 7856756
2 parents b9a7410 + 70e8fa7
commit 7856756
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 1 deletion.
diff --git a/clickhouse_driver/client.py b/clickhouse_driver/client.py
@@ -653,7 +653,7 @@ def from_url(cls, url):
                 kwargs[name] = asbool(value)
             elif name == 'ssl_version':
                 kwargs[name] = getattr(ssl, value)
-            elif name in ['ca_certs', 'ciphers']:
+            elif name in ['ca_certs', 'ciphers', 'keyfile', 'certfile']:
                 kwargs[name] = value
             elif name == 'alt_hosts':
                 kwargs['alt_hosts'] = value

diff --git a/clickhouse_driver/connection.py b/clickhouse_driver/connection.py
@@ -110,6 +110,8 @@ class Connection(object):
     :param ssl_version: see :func:`ssl.wrap_socket` docs.
     :param ca_certs: see :func:`ssl.wrap_socket` docs.
     :param ciphers: see :func:`ssl.wrap_socket` docs.
+    :param keyfile: see :func:`ssl.wrap_socket` docs.
+    :param certfile: see :func:`ssl.wrap_socket` docs.
     :param alt_hosts: list of alternative hosts for connection.
                       Example: alt_hosts=host1:port1,host2:port2.
     :param settings_is_important: ``False`` means unknown settings will be
@@ -131,6 +133,7 @@ def __init__(
             secure=False,
             # Secure socket parameters.
             verify=True, ssl_version=None, ca_certs=None, ciphers=None,
+            keyfile=None, certfile=None,
             alt_hosts=None,
             settings_is_important=False,
     ):
@@ -165,6 +168,10 @@ def __init__(
             ssl_options['ca_certs'] = ca_certs
         if ciphers is not None:
             ssl_options['ciphers'] = ciphers
+        if keyfile is not None:
+            ssl_options['keyfile'] = keyfile
+        if certfile is not None:
+            ssl_options['certfile'] = certfile
 
         self.ssl_options = ssl_options
 
@@ -265,6 +272,10 @@ def _create_ssl_context(self, ssl_options):
         if 'cert_reqs' in ssl_options:
             context.options = ssl_options['cert_reqs']
 
+        if 'certfile' in ssl_options:
+            keyfile = ssl_options.get('keyfile')
+            context.load_cert_chain(ssl_options['certfile'], keyfile=keyfile)
+
         return context
 
     def _init_connection(self, host, port):

diff --git a/tests/test_client.py b/tests/test_client.py
@@ -163,6 +163,42 @@ def test_ssl(self):
             'ciphers': 'HIGH:-aNULL:-eNULL:-PSK:RC4-SHA:RC4-MD5'
         })
 
+    def test_ssl_key_cert(self):
+        base_url = (
+            'clickhouses://host?'
+            'verify=true&'
+            'ssl_version=PROTOCOL_SSLv23&'
+            'ca_certs=/tmp/certs&'
+            'ciphers=HIGH:-aNULL:-eNULL:-PSK:RC4-SHA:RC4-MD5&'
+        )
+        base_expected = {
+            'ssl_version': ssl.PROTOCOL_SSLv23,
+            'ca_certs': '/tmp/certs',
+            'ciphers': 'HIGH:-aNULL:-eNULL:-PSK:RC4-SHA:RC4-MD5'
+        }
+
+        c = Client.from_url(
+            base_url +
+            'keyfile=/tmp/client.key&'
+            'certfile=/tmp/client.cert'
+        )
+        expected = base_expected.copy()
+        expected.update({
+            'keyfile': '/tmp/client.key',
+            'certfile': '/tmp/client.cert'
+        })
+        self.assertEqual(c.connection.ssl_options, expected)
+
+        c = Client.from_url(
+            base_url +
+            'certfile=/tmp/client.cert'
+        )
+        expected = base_expected.copy()
+        expected.update({
+            'certfile': '/tmp/client.cert'
+        })
+        self.assertEqual(c.connection.ssl_options, expected)
+
     def test_alt_hosts(self):
         c = Client.from_url('clickhouse://host?alt_hosts=host2:1234')
         self.assertHostsEqual(c, [('host', 9000), ('host2', 1234)])