Modular sasl mechanisms (wip)

Author: dpkp

kafka/conn.py

@@ -0,0 +1,26 @@
+from __future__ import absolute_import
+
+from kafka.sasl.gssapi import SaslMechanismGSSAPI
+from kafka.sasl.oauth import SaslMechanismOAuth
+from kafka.sasl.plain import SaslMechanismPlain
+from kafka.sasl.scram import SaslMechanismScram
+
+
+SASL_MECHANISMS = {}
+
+
+def register_sasl_mechanism(name, klass, overwrite=False):
+    if not overwrite and name in SASL_MECHANISMS:
+        raise ValueError('Sasl mechanism %s already defined!' % name)
+    SASL_MECHANISMS[name] = klass
+
+
+def get_sasl_mechanism(name):
+    return SASL_MECHANISMS[name]
+
+
+register_sasl_mechanism('GSSAPI', SaslMechanismGSSAPI)
+register_sasl_mechanism('OAUTHBEARER', SaslMechanismOAuth)
+register_sasl_mechanism('PLAIN', SaslMechanismPlain)
+register_sasl_mechanism('SCRAM-SHA-256', SaslMechanismScram)
+register_sasl_mechanism('SCRAM-SHA-512', SaslMechanismScram)

kafka/sasl/__init__.py

@@ -0,0 +1,27 @@
+from __future__ import absolute_import
+
+import abc
+
+
+class SaslMechanism(object):
+    __metaclass__ = abc.ABCMeta
+
+    @abc.abstractmethod
+    def __init__(self, **config):
+        pass
+
+    @abc.abstractmethod
+    def auth_bytes(self):
+        pass
+
+    @abc.abstractmethod
+    def receive(self, auth_bytes):
+        pass
+
+    @abc.abstractmethod
+    def is_done(self):
+        pass
+
+    @abc.abstractmethod
+    def is_authenticated(self):
+        pass

kafka/sasl/abc.py

@@ -0,0 +1,73 @@
+from __future__ import absolute_import
+
+# needed for SASL_GSSAPI authentication:
+try:
+    import gssapi
+    from gssapi.raw.misc import GSSError
+except (ImportError, OSError):
+    #no gssapi available, will disable gssapi mechanism
+    gssapi = None
+    GSSError = None
+
+from kafka.sasl.abc import SaslMechanism
+
+
+class SaslMechanismGSSAPI(SaslMechanism):
+    # Establish security context and negotiate protection level
+    # For reference RFC 2222, section 7.2.1
+
+    SASL_QOP_AUTH = 1
+    SASL_QOP_AUTH_INT = 2
+    SASL_QOP_AUTH_CONF = 4
+
+    def __init__(self, **config):
+        assert gssapi is not None, 'GSSAPI lib not available'
+        assert config['sasl_kerberos_service_name'] is not None, 'sasl_kerberos_service_name required for GSSAPI sasl'
+        self._is_done = False
+        self._is_authenticated = False
+        self.kerberos_damin_name = config['sasl_kerberos_domain_name'] or config['host']
+        self.auth_id = config['sasl_kerberos_service_name'] + '@' + kerberos_damin_name
+        self.gssapi_name = gssapi.Name(auth_id, name_type=gssapi.NameType.hostbased_service).canonicalize(gssapi.MechType.kerberos)
+        self._client_ctx = gssapi.SecurityContext(name=self.gssapi_name, usage='initiate')
+        self._next_token = self._client_ctx.step(None)
+
+    def auth_bytes(self):
+        # GSSAPI Auth does not have a final broker->client message
+        # so mark is_done after the final auth_bytes are provided
+        # in practice we'll still receive a response when using SaslAuthenticate
+        # but not when using the prior unframed approach.
+        if self._client_ctx.complete:
+            self._is_done = True
+            self._is_authenticated = True
+        return self._next_token or b''
+
+    def receive(self, auth_bytes):
+        if not self._client_ctx.complete:
+            # The server will send a token back. Processing of this token either
+            # establishes a security context, or it needs further token exchange.
+            # The gssapi will be able to identify the needed next step.
+            self._next_token = self._client_ctx.step(auth_bytes)
+        elif self._is_done:
+            # The final step of gssapi is send, so we do not expect any additional bytes
+            # however, allow an empty message to support SaslAuthenticate response
+            if auth_bytes != b'':
+                raise ValueError("Unexpected receive auth_bytes after sasl/gssapi completion")
+        else:
+            # unwraps message containing supported protection levels and msg size
+            msg = client_ctx.unwrap(received_token).message
+            # Kafka currently doesn't support integrity or confidentiality security layers, so we
+            # simply set QoP to 'auth' only (first octet). We reuse the max message size proposed
+            # by the server
+            message_parts = [
+                Int8.encode(self.SASL_QOP_AUTH & Int8.decode(io.BytesIO(msg[0:1]))),
+                msg[:1],
+                self.auth_id.encode(),
+            ]
+            # add authorization identity to the response, and GSS-wrap
+            self._next_token = self._client_ctx.wrap(b''.join(message_parts), False).message
+
+    def is_done(self):
+        return self._is_done
+
+    def is_authenticated(self):
+        return self._is_authenticated

kafka/sasl/gssapi.py

@@ -0,0 +1,40 @@
+from __future__ import absolute_import
+
+from kafka.sasl.abc import SaslMechanism
+
+
+class SaslMechanismOAuth(SaslMechanism):
+
+    def __init__(self, **config):
+        self.token_provider = config['sasl_oauth_token_provider']
+        assert self.token_provider is not None, 'sasl_oauth_token_provider required for OAUTHBEARER sasl'
+        assert callable(getattr(self.token_provider, 'token', None)), 'sasl_oauth_token_provider must implement method #token()'
+        self._is_done = False
+        self._is_authenticated = False
+
+    def auth_bytes(self):
+        token = self.token_provider.token()
+        extensions = self._token_extensions()
+        return "n,,\x01auth=Bearer {}{}\x01\x01".format(token, extensions).encode('utf-8')
+
+    def receive(self, auth_bytes):
+        if auth_bytes != b'':
+            self._is_authenticated = False
+            self._is_done = True
+
+    def is_done(self):
+        return self._is_done
+
+    def is_authenticated(self):
+        return self._is_authenticated
+
+    def _token_extensions(self):
+        """
+        Return a string representation of the OPTIONAL key-value pairs that can be sent with an OAUTHBEARER
+        initial request.
+        """
+        # Only run if the #extensions() method is implemented by the clients Token Provider class
+        # Builds up a string separated by \x01 via a dict of key value pairs
+        extensions = getattr(self.token_provider, 'extensions', lambda: [])()
+        msg = '\x01'.join(['{}={}'.format(k, v) for k, v in extensions.items()])
+        return '\x01' + msg if msg else ''

kafka/sasl/oauth.py

@@ -0,0 +1,36 @@
+from __future__ import absolute_import
+
+import logging
+
+from kafka.sasl.abc import SaslMechanism
+
+
+log = logging.getLogger(__name__)
+
+
+class SaslMechanismPlain(SaslMechanism):
+
+    def __init__(self, **config):
+        if config['security_protocol'] == 'SASL_PLAINTEXT':
+            log.warning('Sending username and password in the clear')
+        assert config['sasl_plain_username'] is not None, 'sasl_plain_username required for PLAIN sasl'
+        assert config['sasl_plain_password'] is not None, 'sasl_plain_password required for PLAIN sasl'
+
+        self.username = config['sasl_plain_username']
+        self.password = config['sasl_plain_password']
+        self._is_done = False
+        self._is_authenticated = False
+
+    def auth_bytes(self):
+        # Send PLAIN credentials per RFC-4616
+        return bytes('\0'.join([self.username, self.username, self.password]).encode('utf-8'))
+
+    def receive(self, auth_bytes):
+        self._is_done = True
+        self._is_authenticated = auth_bytes == b''
+
+    def is_done(self):
+        return self._is_done
+
+    def is_authenticated(self):
+        return self._is_authenticated

kafka/sasl/plain.py

@@ -3,11 +3,17 @@
 import base64
 import hashlib
 import hmac
+import logging
 import uuid
 
+
+from kafka.sasl.abc import SaslMechanism
 from kafka.vendor import six
 
 
+log = logging.getLogger(__name__)
+
+
 if six.PY2:
     def xor_bytes(left, right):
         return bytearray(ord(lb) ^ ord(rb) for lb, rb in zip(left, right))
@@ -16,17 +22,58 @@ def xor_bytes(left, right):
         return bytes(lb ^ rb for lb, rb in zip(left, right))
 
 
+class SaslMechanismScram(SaslMechanism):
+
+    def __init__(self, **config):
+        assert config['sasl_plain_username'] is not None, 'sasl_plain_username required for SCRAM sasl'
+        assert config['sasl_plain_password'] is not None, 'sasl_plain_password required for SCRAM sasl'
+        if config['security_protocol'] == 'SASL_PLAINTEXT':
+            log.warning('Exchanging credentials in the clear during Sasl Authentication')
+
+        self._scram_client = ScramClient(
+            config['sasl_plain_username'],
+            config['sasl_plain_password'],
+            config['sasl_mechanism']
+        )
+        self._state = 0
+
+    def auth_bytes(self):
+        if self._state == 0:
+            return self._scram_client.first_message()
+        elif self._state == 1:
+            return self._scram_client.final_message()
+        else:
+            raise ValueError('No auth_bytes for state: %s' % self._state)
+
+    def receive(self, auth_bytes):
+        if self._state == 0:
+            self._scram_client.process_server_first_message(auth_bytes)
+        elif self._state == 1:
+            self._scram_client.process_server_final_message(auth_bytes)
+        else:
+            raise ValueError('Cannot receive bytes in state: %s' % self._state)
+        self._state += 1
+        return self.is_done()
+
+    def is_done(self):
+        return self._state == 2
+
+    def is_authenticated(self):
+        # receive raises if authentication fails...?
+        return self._state == 2
+
+
 class ScramClient:
     MECHANISMS = {
         'SCRAM-SHA-256': hashlib.sha256,
         'SCRAM-SHA-512': hashlib.sha512
     }
 
     def __init__(self, user, password, mechanism):
-        self.nonce = str(uuid.uuid4()).replace('-', '')
-        self.auth_message = ''
+        self.nonce = str(uuid.uuid4()).replace('-', '').encode('utf-8')
+        self.auth_message = b''
         self.salted_password = None
-        self.user = user
+        self.user = user.encode('utf-8')
         self.password = password.encode('utf-8')
         self.hashfunc = self.MECHANISMS[mechanism]
         self.hashname = ''.join(mechanism.lower().split('-')[1:3])
@@ -38,29 +85,29 @@ def __init__(self, user, password, mechanism):
         self.server_signature = None
 
     def first_message(self):
-        client_first_bare = 'n={},r={}'.format(self.user, self.nonce)
+        client_first_bare = b'n=' + self.user + b',r=' + self.nonce
         self.auth_message += client_first_bare
-        return 'n,,' + client_first_bare
+        return b'n,,' + client_first_bare
 
     def process_server_first_message(self, server_first_message):
-        self.auth_message += ',' + server_first_message
-        params = dict(pair.split('=', 1) for pair in server_first_message.split(','))
-        server_nonce = params['r']
+        self.auth_message += b',' + server_first_message
+        params = dict(pair.split('=', 1) for pair in server_first_message.decode('utf-8').split(','))
+        server_nonce = params['r'].encode('utf-8')
         if not server_nonce.startswith(self.nonce):
             raise ValueError("Server nonce, did not start with client nonce!")
         self.nonce = server_nonce
-        self.auth_message += ',c=biws,r=' + self.nonce
+        self.auth_message += b',c=biws,r=' + self.nonce
 
         salt = base64.b64decode(params['s'].encode('utf-8'))
         iterations = int(params['i'])
         self.create_salted_password(salt, iterations)
 
         self.client_key = self.hmac(self.salted_password, b'Client Key')
         self.stored_key = self.hashfunc(self.client_key).digest()
-        self.client_signature = self.hmac(self.stored_key, self.auth_message.encode('utf-8'))
+        self.client_signature = self.hmac(self.stored_key, self.auth_message)
         self.client_proof = xor_bytes(self.client_key, self.client_signature)
         self.server_key = self.hmac(self.salted_password, b'Server Key')
-        self.server_signature = self.hmac(self.server_key, self.auth_message.encode('utf-8'))
+        self.server_signature = self.hmac(self.server_key, self.auth_message)
 
     def hmac(self, key, msg):
         return hmac.new(key, msg, digestmod=self.hashfunc).digest()
@@ -71,11 +118,9 @@ def create_salted_password(self, salt, iterations):
         )
 
     def final_message(self):
-        return 'c=biws,r={},p={}'.format(self.nonce, base64.b64encode(self.client_proof).decode('utf-8'))
+        return b'c=biws,r=' + self.nonce + b',p=' + base64.b64encode(self.client_proof)
 
     def process_server_final_message(self, server_final_message):
-        params = dict(pair.split('=', 1) for pair in server_final_message.split(','))
+        params = dict(pair.split('=', 1) for pair in server_final_message.decode('utf-8').split(','))
         if self.server_signature != base64.b64decode(params['v'].encode('utf-8')):
             raise ValueError("Server sent wrong signature!")
-
-