feat(core): add SASL DIGEST-MD5 support

This adds the possibility to connect to Zookeeper using DIGEST-MD5 SASL.
It uses the pure-sasl library to connect using SASL. In case the library
is missing, connection to Zookeeper will be done without any
authentification and a warning message will be displayed. Tests have
been added for this feature. Documentation also has been updated.

7 files changed, 154 insertions, 9 deletions

diff --git a/kazoo/client.py b/kazoo/client.py
index 3391b71..afff0dd 100644
--- a/kazoo/client.py
+++ b/kazoo/client.py
@@ -303,6 +303,15 @@ class KazooClient(object):
         self.Semaphore = partial(Semaphore, self)
         self.ShallowParty = partial(ShallowParty, self)
 
+        # Managing SASL client
+        self.use_sasl = False
+        for scheme, auth in self.auth_data:
+            if scheme == "sasl":
+                self.use_sasl = True
+                # Could be used later for GSSAPI implementation
+                self.sasl_server_principal = "zk-sasl-md5"
+                break
+
         # If we got any unhandled keywords, complain like Python would
         if kwargs:
             raise TypeError('__init__() got unexpected keyword arguments: %s'
@@ -728,8 +737,12 @@ class KazooClient(object):
         """Send credentials to server.
 
         :param scheme: authentication scheme (default supported:
-                       "digest").
+                       "digest", "sasl"). Note that "sasl" scheme is
+                       requiring "pure-sasl" library to be
+                       installed.
         :param credential: the credential -- value depends on scheme.
+                           "digest": user:password
+                           "sasl": user:password
 
         :returns: True if it was successful.
         :rtype: bool
diff --git a/kazoo/protocol/connection.py b/kazoo/protocol/connection.py
index 32c7398..9d0e087 100644
--- a/kazoo/protocol/connection.py
+++ b/kazoo/protocol/connection.py
@@ -25,6 +25,7 @@ from kazoo.protocol.serialization import (
     Ping,
     PingInstance,
     ReplyHeader,
+    SASL,
     Transaction,
     Watch,
     int_struct
@@ -39,6 +40,11 @@ from kazoo.retry import (
     ForceRetryError,
     RetryFailedError
 )
+try:
+    from puresasl.client import SASLClient
+    PURESASL_AVAILABLE = True
+except ImportError:
+    PURESASL_AVAILABLE = False
 
 
 log = logging.getLogger(__name__)
@@ -154,6 +160,8 @@ class ConnectionHandler(object):
 
         self._connection_routine = None
 
+        self.sasl_cli = None
+
     # This is instance specific to avoid odd thread bug issues in Python
     # during shutdown global cleanup
     @contextmanager
@@ -416,6 +424,24 @@ class ConnectionHandler(object):
                 async_object.set(True)
         elif header.xid == WATCH_XID:
             self._read_watch_event(buffer, offset)
+        elif self.sasl_cli and not self.sasl_cli.complete:
+            # SASL authentication is not yet finished, this can only
+            # be a SASL packet
+            self.logger.log(BLATHER, 'Received SASL')
+            try:
+                challenge, _ = SASL.deserialize(buffer, offset)
+            except Exception:
+                raise ConnectionDropped('error while SASL authentication.')
+            response = self.sasl_cli.process(challenge)
+            if response:
+                # authentication not yet finished, answering the challenge
+                self._send_sasl_request(challenge=response,
+                                        timeout=client._session_timeout)
+            else:
+                # authentication is ok, state is CONNECTED
+                # remove sensible information from the object
+                client._session_callback(KeeperState.CONNECTED)
+                self.sasl_cli.dispose()
         else:
             self.logger.log(BLATHER, 'Reading for header %r', header)
 
@@ -544,11 +570,11 @@ class ConnectionHandler(object):
             client._session_callback(KeeperState.CONNECTING)
 
         try:
+            self._xid = 0
             read_timeout, connect_timeout = self._connect(host, port)
             read_timeout = read_timeout / 1000.0
             connect_timeout = connect_timeout / 1000.0
             retry.reset()
-            self._xid = 0
             self.ping_outstanding.clear()
             with self._socket_error_handling():
                 while not close_connection:
@@ -660,13 +686,53 @@ class ConnectionHandler(object):
             client._session_callback(KeeperState.CONNECTED_RO)
             self._ro_mode = iter(self._server_pinger())
         else:
-            client._session_callback(KeeperState.CONNECTED)
             self._ro_mode = None
-
-        for scheme, auth in client.auth_data:
-            ap = Auth(0, scheme, auth)
-            zxid = self._invoke(connect_timeout / 1000.0, ap, xid=AUTH_XID)
-            if zxid:
-                client.last_zxid = zxid
+            if client.use_sasl and self.sasl_cli is None:
+                if PURESASL_AVAILABLE:
+                    for scheme, auth in client.auth_data:
+                        if scheme == 'sasl':
+                            username, password = auth.split(":")
+                            self.sasl_cli = SASLClient(
+                                host=client.sasl_server_principal,
+                                service='zookeeper',
+                                mechanism='DIGEST-MD5',
+                                username=username,
+                                password=password
+                            )
+                            break
+
+                    # As described in rfc
+                    # https://tools.ietf.org/html/rfc2831#section-2.1
+                    # sending empty challenge
+                    self._send_sasl_request(challenge=b'',
+                                            timeout=connect_timeout)
+                else:
+                    self.logger.warn('Pure-sasl library is missing while sasl'
+                                     ' authentification is configured. Please'
+                                     ' install pure-sasl library to connect '
+                                     'using sasl. Now falling back '
+                                     'connecting WITHOUT any '
+                                     'authentification.')
+                    client.use_sasl = False
+                    client._session_callback(KeeperState.CONNECTED)
+            else:
+                client._session_callback(KeeperState.CONNECTED)
+                for scheme, auth in client.auth_data:
+                    if scheme == "digest":
+                        ap = Auth(0, scheme, auth)
+                        zxid = self._invoke(
+                            connect_timeout / 1000.0,
+                            ap,
+                            xid=AUTH_XID
+                        )
+                        if zxid:
+                            client.last_zxid = zxid
 
         return read_timeout, connect_timeout
+
+    def _send_sasl_request(self, challenge, timeout):
+        """ Called when sending a SASL request, xid needs be to incremented """
+        sasl_request = SASL(challenge)
+        self._xid = (self._xid % 2147483647) + 1
+        xid = self._xid
+        self._submit(sasl_request, timeout / 1000.0, xid)
diff --git a/kazoo/protocol/serialization.py b/kazoo/protocol/serialization.py
index 046a62e..75c6abe 100644
--- a/kazoo/protocol/serialization.py
+++ b/kazoo/protocol/serialization.py
@@ -378,6 +378,19 @@ class Auth(namedtuple('Auth', 'auth_type scheme auth')):
                 write_string(self.auth))
 
 
+class SASL(namedtuple('SASL', 'challenge')):
+    type = 102
+
+    def serialize(self):
+        b = bytearray()
+        b.extend(write_buffer(self.challenge))
+        return b
+
+    @classmethod
+    def deserialize(cls, bytes, offset):
+        challenge, offset = read_buffer(bytes, offset)
+        return challenge, offset
+
 class Watch(namedtuple('Watch', 'type state path')):
     @classmethod
     def deserialize(cls, bytes, offset):
diff --git a/kazoo/testing/common.py b/kazoo/testing/common.py
index ca92cc5..c4aa0b9 100644
--- a/kazoo/testing/common.py
+++ b/kazoo/testing/common.py
@@ -96,6 +96,7 @@ class ManagedZooKeeper(object):
         if self.running:
             return
         config_path = os.path.join(self.working_path, "zoo.cfg")
+        jass_config_path = os.path.join(self.working_path, "jaas.conf")
         log_path = os.path.join(self.working_path, "log")
         log4j_path = os.path.join(self.working_path, "log4j.properties")
         data_path = os.path.join(self.working_path, "data")
@@ -115,6 +116,7 @@ dataDir=%s
 clientPort=%s
 maxClientCnxns=0
 admin.serverPort=%s
+authProvider.1=org.apache.zookeeper.server.auth.SASLAuthenticationProvider
 """ % (to_java_compatible_path(data_path),
        self.server_info.client_port,
        self.server_info.admin_port))  # NOQA
@@ -137,6 +139,14 @@ peerType=%s
         # Write server ids into datadir
         with open(os.path.join(data_path, "myid"), "w") as myid_file:
             myid_file.write(str(self.server_info.server_id))
+        # Write JAAS configuration
+        with open(jass_config_path, "w") as jaas_file:
+            jaas_file.write("""
+Server {
+    org.apache.zookeeper.server.auth.DigestLoginModule required
+    user_super="super_secret"
+    user_jaasuser="jaas_password";
+};""")
 
         with open(log4j_path, "w") as log4j:
             log4j.write("""
@@ -163,6 +173,9 @@ log4j.appender.ROLLINGFILE.File=""" + to_java_compatible_path(  # NOQA
             # and from activation of the main workspace on run.
             "-Djava.awt.headless=true",
 
+            # JAAS configuration for SASL authentication
+            "-Djava.security.auth.login.config=%s" % jass_config_path,
+
             "org.apache.zookeeper.server.quorum.QuorumPeerMain",
             config_path,
         ]
diff --git a/kazoo/tests/test_client.py b/kazoo/tests/test_client.py
index 5bc7999..0dad864 100644
--- a/kazoo/tests/test_client.py
+++ b/kazoo/tests/test_client.py
@@ -179,6 +179,34 @@ class TestAuthentication(KazooTestCase):
             client.stop()
             client.close()
 
+    def test_connect_sasl_auth(self):
+        from kazoo.security import make_acl
+
+        if TRAVIS_ZK_VERSION:
+            version = TRAVIS_ZK_VERSION
+        else:
+            version = self.client.server_version()
+        if not version or version < (3, 4):
+            raise SkipTest("Must use Zookeeper 3.4 or above")
+
+        username = "jaasuser"
+        password = "jaas_password"
+        sasl_auth = "%s:%s" % (username, password)
+
+        acl = make_acl('sasl', credential=username, all=True)
+
+        client = self._get_client(auth_data=[('sasl', sasl_auth)])
+        client.start()
+        try:
+            client.create('/1', acl=(acl,))
+            # give ZK a chance to copy data to other node
+            time.sleep(0.1)
+            self.assertRaises(NoAuthError, self.client.get, "/1")
+        finally:
+            client.delete('/1')
+            client.stop()
+            client.close()
+
     def test_unicode_auth(self):
         username = u("xe4/\hm")
         password = u("/\xe4hm")
@@ -217,6 +245,16 @@ class TestAuthentication(KazooTestCase):
         self.assertRaises(TypeError, client.add_auth,
                           None, ('user', 'pass'))
 
+    def test_invalid_sasl_auth(self):
+        if TRAVIS_ZK_VERSION:
+            version = TRAVIS_ZK_VERSION
+        else:
+            version = self.client.server_version()
+        if not version or version < (3, 4):
+            raise SkipTest("Must use Zookeeper 3.4 or above")
+        client = self._get_client(auth_data=[('sasl', 'baduser:badpassword')])
+        self.assertRaises(ConnectionLoss, client.start)
+
     def test_async_auth(self):
         client = self._get_client()
         client.start()
diff --git a/requirements.txt b/requirements.txt
index 87f71c2..d96b210 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,5 @@
 coverage==3.7.1
 mock==1.0.1
 nose==1.3.3
+pure-sasl==0.5.1
 flake8==2.3.0
diff --git a/setup.py b/setup.py
index 5edca74..21127c4 100644
--- a/setup.py
+++ b/setup.py
@@ -24,6 +24,7 @@ tests_require = install_requires + [
     'mock',
     'nose',
     'flake8',
+    'pure-sasl',
 ]
 
 if not (PYTHON3 or PYPY):