Convert msk to sasl module

Author: dpkp

Diff for: kafka/sasl/__init__.py

@@ -1,6 +1,7 @@
 from __future__ import absolute_import
 
 from kafka.sasl.gssapi import SaslMechanismGSSAPI
+from kafka.sasl.msk import SaslMechanismAwsMskIam
 from kafka.sasl.oauth import SaslMechanismOAuth
 from kafka.sasl.plain import SaslMechanismPlain
 from kafka.sasl.scram import SaslMechanismScram
@@ -24,3 +25,4 @@ def get_sasl_mechanism(name):
 register_sasl_mechanism('PLAIN', SaslMechanismPlain)
 register_sasl_mechanism('SCRAM-SHA-256', SaslMechanismScram)
 register_sasl_mechanism('SCRAM-SHA-512', SaslMechanismScram)
+register_sasl_mechanism('AWS_MSK_IAM', SaslMechanismAwsMskIam)

Diff for: kafka/msk.py renamed to kafka/sasl/msk.py

@@ -1,12 +1,57 @@
+from __future__ import absolute_import
+
 import datetime
 import hashlib
 import hmac
 import json
+import logging
 import string
 
+# needed for AWS_MSK_IAM authentication:
+try:
+    from botocore.session import Session as BotoSession
+except ImportError:
+    # no botocore available, will disable AWS_MSK_IAM mechanism
+    BotoSession = None
+
+from kafka.sasl.abc import SaslMechanism
 from kafka.vendor.six.moves import urllib
 
 
+log = logging.getLogger(__name__)
+
+
+class SaslMechanismAwsMskIam(SaslMechanism):
+    def __init__(self, **config):
+        assert BotoSession is not None, 'AWS_MSK_IAM requires the "botocore" package'
+        assert config['security_protocol'] == 'SASL_SSL', 'AWS_MSK_IAM requires SASL_SSL'
+        self._is_done = False
+        self._is_authenticated = False
+
+    def auth_bytes(self):
+        session = BotoSession()
+        credentials = session.get_credentials().get_frozen_credentials()
+        client = AwsMskIamClient(
+            host=self.host,
+            access_key=credentials.access_key,
+            secret_key=credentials.secret_key,
+            region=session.get_config_variable('region'),
+            token=credentials.token,
+        )
+        return client.first_message()
+
+    def receive(self, auth_bytes):
+        self._is_done = True
+        self._is_authenticated = auth_bytes != b''
+        log.info('Authenticated via AWS_MSK_IAM %s', auth_bytes.decode('utf-8'))
+
+    def is_done(self):
+        return self._is_done
+
+    def is_authenticated(self):
+        return self._is_authenticated
+
+
 class AwsMskIamClient:
     UNRESERVED_CHARS = string.ascii_letters + string.digits + '-._~'
 

Diff for: test/test_msk.py renamed to test/sasl/test_msk.py

@@ -1,7 +1,7 @@
 import datetime
 import json
 
-from kafka.msk import AwsMskIamClient
+from kafka.sasl.msk import AwsMskIamClient
 
 try:
     from unittest import mock
@@ -11,7 +11,7 @@
 
 def client_factory(token=None):
     now = datetime.datetime.utcfromtimestamp(1629321911)
-    with mock.patch('kafka.msk.datetime') as mock_dt:
+    with mock.patch('kafka.sasl.msk.datetime') as mock_dt:
         mock_dt.datetime.utcnow = mock.Mock(return_value=now)
         return AwsMskIamClient(
             host='localhost',