DevAPI: Core TLS/SSL options for the mysqlx URI scheme

amitab · amitab · commit 1c8dfb1c58e6 · 2016-10-05T15:18:26.000+05:30
This patch adds the ability to connect to a MySQLx enabled server with SSL.
New connection options added:
1. ssl-enable
2. ssl-ca
3. ssl-cert
4. ssl-key

ssl-enable can be used by itself to use default SSL settings.

Tests were added for regression.
diff --git a/lib/mysqlx/__init__.py b/lib/mysqlx/__init__.py
@@ -25,9 +25,9 @@
 
 import re
 
-import constants
+from . import constants
 
-from .compat import STRING_TYPES, urlparse
+from .compat import STRING_TYPES, urlparse, parse_qsl
 from .connection import XSession, NodeSession
 from .crud import Schema, Collection, Table, View
 from .dbdoc import DbDoc
@@ -108,12 +108,16 @@ def _parse_connection_uri(uri):
     if parsed.hostname is None or parsed.username is None \
        or parsed.password is None:
         raise InterfaceError("Malformed URI '{0}'".format(uri))
+
     settings = {
         "user": parsed.username,
         "password": parsed.password,
         "schema": parsed.path.lstrip("/")
     }
 
+    query = dict(parse_qsl(parsed.query, True))
+    for opt, val in query.items():
+        settings[opt] = val.strip("() ") or True
     settings.update(_parse_address_list(parsed.netloc.split("@")[-1]))
     return settings
 
@@ -172,6 +176,10 @@ def _get_connection_settings(*args, **kwargs):
     else:
         _validate_settings(settings)
 
+    ssl_opts = ["ssl-key", "ssl-cert", "ssl-ca", "ssl-crl"]
+    if any(key in settings for key in ssl_opts):
+        settings["ssl-enable"] = True
+
     return settings
 
 def get_session(*args, **kwargs):
diff --git a/lib/mysqlx/compat.py b/lib/mysqlx/compat.py
@@ -32,7 +32,7 @@
 
 
 if PY3:
-    from urllib.parse import urlparse
+    from urllib.parse import urlparse, parse_qsl
 
     def hexlify(data):
         return binascii.hexlify(data).decode("utf-8")
@@ -45,7 +45,7 @@ def hexlify(data):
 
 
 else:
-    from urlparse import urlparse
+    from urlparse import urlparse, parse_qsl
 
     def hexlify(data):
         return data.encode("hex")
diff --git a/lib/mysqlx/connection.py b/lib/mysqlx/connection.py
@@ -23,6 +23,12 @@
 
 """Implementation of communication for MySQL X servers."""
 
+try:
+    import ssl
+    SSL_AVAILABLE = True
+except:
+    SSL_AVAILABLE = False
+
 import socket
 
 from functools import wraps
@@ -42,6 +48,7 @@
 class SocketStream(object):
     def __init__(self):
         self._socket = None
+        self._is_ssl = False
 
     def connect(self, host, port):
         self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
@@ -68,6 +75,34 @@ def close(self):
         self._socket.close()
         self._socket = None
 
+    def set_ssl(self, ssl_opts={}):
+        if not SSL_AVAILABLE:
+            raise RuntimeError("Python installation has no SSL support.")
+
+        context = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
+        context.load_default_certs()
+        if "ssl-ca" in ssl_opts:
+            try:
+                context.load_verify_locations(ssl_opts["ssl-ca"])
+                context.verify_mode = ssl.CERT_REQUIRED
+            except (IOError, SSLError):
+                raise InterfaceError("Invalid CA certificate.")
+        if "ssl-crl" in ssl_opts:
+            try:
+                context.load_verify_locations(ssl_opts["ssl-crl"])
+                context.verify_flags = ssl.VERIFY_CRL_CHECK_CHAIN
+            except (IOError, SSLError):
+                raise InterfaceError("Invalid CRL.")
+        if "ssl-cert" in ssl_opts:
+            try:
+                context.load_cert_chain(ssl_opts["ssl-cert"],
+                    ssl_opts.get("ssl-key", None))
+            except (IOError, SSLError):
+                raise InterfaceError("Invalid Client Certificate/Key.")
+
+        self._socket = context.wrap_socket(self._socket)
+        self._is_ssl = True
+
 
 def catch_network_exception(func):
     @wraps(func)
@@ -104,10 +139,16 @@ def connect(self, host, port):
         self._authenticate()
 
     def _handle_capabilities(self):
-        # TODO: To implement
-        # caps = mysqlx_connection_pb2.CapabilitiesGet()
-        # data = caps.SerializeToString()
-        pass
+        if not self.settings.get("ssl-enable", False):
+            return
+
+        data = self.protocol.get_capabilites()
+        if not next((True for cap in data.capabilities \
+            if cap.name == "tls"), False):
+            raise OperationalError("SSL is not enabled on server.")
+
+        self.protocol.set_capabilities(tls=True)
+        self.stream.set_ssl(self.settings)
 
     def _authenticate(self):
         plugin = MySQL41AuthPlugin(self._user, self._password)
diff --git a/lib/mysqlx/protocol.py b/lib/mysqlx/protocol.py
@@ -33,6 +33,7 @@
 from .protobuf import mysqlx_resultset_pb2 as MySQLxResultset
 from .protobuf import mysqlx_crud_pb2 as MySQLxCrud
 from .protobuf import mysqlx_expr_pb2 as MySQLxExpr
+from .protobuf import mysqlx_connection_pb2 as MySQLxConnection
 from .result import ColumnMetaData
 from .compat import STRING_TYPES, INT_TYPES
 from .dbdoc import DbDoc
@@ -56,6 +57,7 @@
     (MySQLx.ServerMessages.RESULTSET_FETCH_DONE_MORE_RESULTSETS,
      MySQLxResultset.FetchDoneMoreResultsets),
     (MySQLx.ServerMessages.OK, MySQLx.Ok),
+    (MySQLx.ServerMessages.CONN_CAPABILITIES, MySQLxConnection.Capabilities),
 ]
 
 
@@ -105,6 +107,24 @@ def __init__(self, reader_writer):
         self._writer = reader_writer
         self._message = None
 
+    def get_capabilites(self):
+        msg = MySQLxConnection.CapabilitiesGet()
+        self._writer.write_message(
+            MySQLx.ClientMessages.CON_CAPABILITIES_GET, msg)
+        return self._reader.read_message()
+
+    def set_capabilities(self, **kwargs):
+        msg = MySQLxConnection.CapabilitiesSet()
+        for key, value in kwargs.items():
+            value = self._create_any(value)
+            capability = MySQLxConnection.Capability(name=key, value=value)
+            msg.capabilities.capabilities.extend([capability])
+
+        self._writer.write_message(
+            MySQLx.ClientMessages.CON_CAPABILITIES_SET, msg)
+
+        return self.read_ok()
+
     def send_auth_start(self, method):
         msg = MySQLxSession.AuthenticateStart(mech_name=method)
         self._writer.write_message(
diff --git a/tests/test_mysqlx_connection.py b/tests/test_mysqlx_connection.py
@@ -273,6 +273,26 @@ def test_bind_to_default_shard(self):
                 tests.MYSQL_SERVERS[0].start()
                 tests.MYSQL_SERVERS[0].wait_up()
 
+    def test_ssl_connection(self):
+        config = {}
+        config.update(self.connect_kwargs)
+        config["ssl-ca"] = tests.SSL_CA
+        config["ssl-cert"] = tests.SSL_CERT
+        config["ssl-key"] = tests.SSL_KEY
+
+        session = mysqlx.get_session(config)
+
+        res = mysqlx.statement.SqlStatement(session._connection,
+            "SHOW STATUS LIKE 'Mysqlx_ssl_active'").execute().fetch_all()
+        self.assertEqual("ON", res[0][1])
+
+        res = mysqlx.statement.SqlStatement(session._connection,
+            "SHOW STATUS LIKE 'Mysqlx_ssl_version'").execute().fetch_all()
+        self.assertTrue("TLS" in res[0][1])
+
+        session.close()
+
+
 @unittest.skipIf(tests.MYSQL_VERSION < (5, 7, 12), "XPlugin not compatible")
 class MySQLxNodeSessionTests(tests.MySQLxTests):
 
@@ -384,3 +404,22 @@ def test_commit(self):
         self.assertEqual(table.count(), 1)
 
         schema.drop_table(table_name)
+
+    def test_ssl_connection(self):
+        config = {}
+        config.update(self.connect_kwargs)
+        config["ssl-ca"] = tests.SSL_CA
+        config["ssl-cert"] = tests.SSL_CERT
+        config["ssl-key"] = tests.SSL_KEY
+
+        session = mysqlx.get_node_session(config)
+
+        res = session.sql("SHOW STATUS LIKE 'Mysqlx_ssl_active'") \
+                     .execute().fetch_all()
+        self.assertEqual("ON", res[0][1])
+
+        res = session.sql("SHOW STATUS LIKE 'Mysqlx_ssl_version'") \
+            .execute().fetch_all()
+        self.assertTrue("TLS" in res[0][1])
+
+        session.close()
