Repository: dpkp/kafka-python
Pull request #2519 (2025-03-11 19:26:25)

Support AWS_MSK_IAM authentication
Update @mattoberle PR #2255 to new sasl module.


Patch:
diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml
index bc0724e..1b0d71c 100644
--- a/.github/workflows/python-package.yml
+++ b/.github/workflows/python-package.yml
@@ -65,7 +65,6 @@ jobs:
           sudo apt install -y libsnappy-dev libzstd-dev
           python -m pip install --upgrade pip
           pip install -r requirements-dev.txt
-          pip install tox-gh-actions
       - name: Pylint
         run: pylint --recursive=y --errors-only --exit-zero kafka test
       - name: Setup java
diff --git a/kafka/sasl/__init__.py b/kafka/sasl/__init__.py
index e36d1df..8677f60 100644
--- a/kafka/sasl/__init__.py
+++ b/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 @@ register_sasl_mechanism('OAUTHBEARER', SaslMechanismOAuth)
 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 --git a/kafka/sasl/msk.py b/kafka/sasl/msk.py
new file mode 100644
index 0000000..db56b48
--- /dev/null
+++ b/kafka/sasl/msk.py
@@ -0,0 +1,233 @@
+from __future__ import absolute_import
+
+import datetime
+import hashlib
+import hmac
+import json
+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
+
+
+class SaslMechanismAwsMskIam(SaslMechanism):
+    def __init__(self, **config):
+        assert BotoSession is not None, 'AWS_MSK_IAM requires the "botocore" package'
+        assert config.get('security_protocol', '') == 'SASL_SSL', 'AWS_MSK_IAM requires SASL_SSL'
+        assert 'host' in config, 'AWS_MSK_IAM requires host configuration'
+        self.host = config['host']
+        self._auth = None
+        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''
+        self._auth = auth_bytes.deode('utf-8')
+
+    def is_done(self):
+        return self._is_done
+
+    def is_authenticated(self):
+        return self._is_authenticated
+
+    def auth_details(self):
+        if not self.is_authenticated:
+            raise RuntimeError('Not authenticated yet!')
+        return 'Authenticated via SASL / AWS_MSK_IAM %s' % (self._auth,)
+
+
+class AwsMskIamClient:
+    UNRESERVED_CHARS = string.ascii_letters + string.digits + '-._~'
+
+    def __init__(self, host, access_key, secret_key, region, token=None):
+        """
+        Arguments:
+            host (str): The hostname of the broker.
+            access_key (str): An AWS_ACCESS_KEY_ID.
+            secret_key (str): An AWS_SECRET_ACCESS_KEY.
+            region (str): An AWS_REGION.
+            token (Optional[str]): An AWS_SESSION_TOKEN if using temporary
+                credentials.
+        """
+        self.algorithm = 'AWS4-HMAC-SHA256'
+        self.expires = '900'
+        self.hashfunc = hashlib.sha256
+        self.headers = [
+            ('host', host)
+        ]
+        self.version = '2020_10_22'
+
+        self.service = 'kafka-cluster'
+        self.action = '{}:Connect'.format(self.service)
+
+        now = datetime.datetime.utcnow()
+        self.datestamp = now.strftime('%Y%m%d')
+        self.timestamp = now.strftime('%Y%m%dT%H%M%SZ')
+
+        self.host = host
+        self.access_key = access_key
+        self.secret_key = secret_key
+        self.region = region
+        self.token = token
+
+    @property
+    def _credential(self):
+        return '{0.access_key}/{0._scope}'.format(self)
+
+    @property
+    def _scope(self):
+        return '{0.datestamp}/{0.region}/{0.service}/aws4_request'.format(self)
+
+    @property
+    def _signed_headers(self):
+        """
+        Returns (str):
+            An alphabetically sorted, semicolon-delimited list of lowercase
+            request header names.
+        """
+        return ';'.join(sorted(k.lower() for k, _ in self.headers))
+
+    @property
+    def _canonical_headers(self):
+        """
+        Returns (str):
+            A newline-delited list of header names and values.
+            Header names are lowercased.
+        """
+        return '\n'.join(map(':'.join, self.headers)) + '\n'
+
+    @property
+    def _canonical_request(self):
+        """
+        Returns (str):
+            An AWS Signature Version 4 canonical request in the format:
+                <Method>\n
+                <Path>\n
+                <CanonicalQueryString>\n
+                <CanonicalHeaders>\n
+                <SignedHeaders>\n
+                <HashedPayload>
+        """
+        # The hashed_payload is always an empty string for MSK.
+        hashed_payload = self.hashfunc(b'').hexdigest()
+        return '\n'.join((
+            'GET',
+            '/',
+            self._canonical_querystring,
+            self._canonical_headers,
+            self._signed_headers,
+            hashed_payload,
+        ))
+
+    @property
+    def _canonical_querystring(self):
+        """
+        Returns (str):
+            A '&'-separated list of URI-encoded key/value pairs.
+        """
+        params = []
+        params.append(('Action', self.action))
+        params.append(('X-Amz-Algorithm', self.algorithm))
+        params.append(('X-Amz-Credential', self._credential))
+        params.append(('X-Amz-Date', self.timestamp))
+        params.append(('X-Amz-Expires', self.expires))
+        if self.token:
+            params.append(('X-Amz-Security-Token', self.token))
+        params.append(('X-Amz-SignedHeaders', self._signed_headers))
+
+        return '&'.join(self._uriencode(k) + '=' + self._uriencode(v) for k, v in params)
+
+    @property
+    def _signing_key(self):
+        """
+        Returns (bytes):
+            An AWS Signature V4 signing key generated from the secret_key, date,
+            region, service, and request type.
+        """
+        key = self._hmac(('AWS4' + self.secret_key).encode('utf-8'), self.datestamp)
+        key = self._hmac(key, self.region)
+        key = self._hmac(key, self.service)
+        key = self._hmac(key, 'aws4_request')
+        return key
+
+    @property
+    def _signing_str(self):
+        """
+        Returns (str):
+            A string used to sign the AWS Signature V4 payload in the format:
+                <Algorithm>\n
+                <Timestamp>\n
+                <Scope>\n
+                <CanonicalRequestHash>
+        """
+        canonical_request_hash = self.hashfunc(self._canonical_request.encode('utf-8')).hexdigest()
+        return '\n'.join((self.algorithm, self.timestamp, self._scope, canonical_request_hash))
+
+    def _uriencode(self, msg):
+        """
+        Arguments:
+            msg (str): A string to URI-encode.
+
+        Returns (str):
+            The URI-encoded version of the provided msg, following the encoding
+            rules specified: https://github.com/aws/aws-msk-iam-auth#uriencode
+        """
+        return urllib.parse.quote(msg, safe=self.UNRESERVED_CHARS)
+
+    def _hmac(self, key, msg):
+        """
+        Arguments:
+            key (bytes): A key to use for the HMAC digest.
+            msg (str): A value to include in the HMAC digest.
+        Returns (bytes):
+            An HMAC digest of the given key and msg.
+        """
+        return hmac.new(key, msg.encode('utf-8'), digestmod=self.hashfunc).digest()
+
+    def first_message(self):
+        """
+        Returns (bytes):
+            An encoded JSON authentication payload that can be sent to the
+            broker.
+        """
+        signature = hmac.new(
+            self._signing_key,
+            self._signing_str.encode('utf-8'),
+            digestmod=self.hashfunc,
+        ).hexdigest()
+        msg = {
+            'version':  self.version,
+            'host': self.host,
+            'user-agent': 'kafka-python',
+            'action': self.action,
+            'x-amz-algorithm': self.algorithm,
+            'x-amz-credential': self._credential,
+            'x-amz-date': self.timestamp,
+            'x-amz-signedheaders': self._signed_headers,
+            'x-amz-expires': self.expires,
+            'x-amz-signature': signature,
+        }
+        if self.token:
+            msg['x-amz-security-token'] = self.token
+
+        return json.dumps(msg, separators=(',', ':')).encode('utf-8')
diff --git a/pyproject.toml b/pyproject.toml
index 48be87f..ddd40a0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -37,7 +37,7 @@ crc32c = ["crc32c"]
 lz4 = ["lz4"]
 snappy = ["python-snappy"]
 zstd = ["zstandard"]
-testing = ["pytest", "mock", "pytest-mock"]
+testing = ["pytest", "mock; python_version < '3.3'", "pytest-mock"]
 
 [tool.setuptools]
 include-package-data = false
diff --git a/requirements-dev.txt b/requirements-dev.txt
index e272d1f..6cfb6d8 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -3,7 +3,7 @@ crc32c
 docker-py
 flake8
 lz4
-mock
+mock; python_version < '3.3'
 py
 pylint
 pytest
@@ -13,6 +13,5 @@ pytest-pylint
 python-snappy
 Sphinx
 sphinx-rtd-theme
-tox
 xxhash
 zstandard
diff --git a/tox.ini b/tox.ini
deleted file mode 100644
index 71e443d..0000000
--- a/tox.ini
+++ /dev/null
@@ -1,49 +0,0 @@
-[tox]
-envlist = py{38,39,310,311,312,py}, docs
-
-[pytest]
-testpaths = kafka test
-addopts = --durations=10
-log_format = %(created)f %(filename)-23s %(threadName)s %(message)s
-
-[gh-actions]
-python =
-    3.8: py38
-    3.9: py39
-    3.10: py310
-    3.11: py311
-    3.12: py312
-    pypy-3.9: pypy
-
-[testenv]
-deps =
-    pytest
-    pytest-cov
-    pylint
-    pytest-pylint
-    pytest-mock
-    mock
-    python-snappy
-    zstandard
-    lz4
-    xxhash
-    crc32c
-commands =
-    pytest {posargs:--cov=kafka --cov-config=.covrc}
-setenv =
-    CRC32C_SW_MODE = auto
-    PROJECT_ROOT = {toxinidir}
-passenv = KAFKA_*
-
-[testenv:pypy]
-# pylint is super slow on pypy...
-commands = pytest {posargs:--cov=kafka --cov-config=.covrc}
-
-[testenv:docs]
-deps =
-    sphinx_rtd_theme
-    sphinx
-
-commands =
-    sphinx-apidoc -o docs/apidoc/ kafka/
-    sphinx-build -b html docs/ docs/_build


Test patch:
diff --git a/docs/tests.rst b/docs/tests.rst
index 561179c..988afca 100644
--- a/docs/tests.rst
+++ b/docs/tests.rst
@@ -6,12 +6,14 @@ Tests
 .. image:: https://travis-ci.org/dpkp/kafka-python.svg?branch=master
     :target: https://travis-ci.org/dpkp/kafka-python
 
-Test environments are managed via tox. The test suite is run via pytest.
+The test suite is run via pytest.
 
-Linting is run via pylint, but is generally skipped on pypy due to pylint
-compatibility / performance issues.
+Linting is run via pylint, but is currently skipped during CI/CD due to
+accumulated debt. We'd like to transition to ruff!
 
 For test coverage details, see https://coveralls.io/github/dpkp/kafka-python
+Coverage reporting is currently disabled as we have transitioned from travis
+to GH Actions and have not yet re-enabled coveralls integration.
 
 The test suite includes unit tests that mock network interfaces, as well as
 integration tests that setup and teardown kafka broker (and zookeeper)
@@ -21,30 +23,21 @@ fixtures for client / consumer / producer testing.
 Unit tests
 ------------------
 
-To run the tests locally, install tox:
+To run the tests locally, install test dependencies:
 
 .. code:: bash
 
-     pip install tox
+     pip install -r requirements-dev.txt
 
-For more details, see https://tox.readthedocs.io/en/latest/install.html
-
-Then simply run tox, optionally setting the python environment.
-If unset, tox will loop through all environments.
+Then simply run pytest (or make test) from your preferred python + virtualenv.
 
 .. code:: bash
 
-    tox -e py27
-    tox -e py35
-
-    # run protocol tests only
-    tox -- -v test.test_protocol
-
-    # re-run the last failing test, dropping into pdb
-    tox -e py27 -- --lf --pdb
+    # run protocol tests only (via pytest)
+    pytest test/test_protocol.py
 
-    # see available (pytest) options
-    tox -e py27 -- --help
+    # Run conn tests only (via make)
+    PYTESTS=test/test_conn.py make test
 
 
 Integration tests
@@ -52,35 +45,8 @@ Integration tests
 
 .. code:: bash
 
-    KAFKA_VERSION=0.8.2.2 tox -e py27
-    KAFKA_VERSION=1.0.1 tox -e py36
-
-
-Integration tests start Kafka and Zookeeper fixtures. This requires downloading
-kafka server binaries:
-
-.. code:: bash
-
-    ./build_integration.sh
-
-By default, this will install the broker versions listed in build_integration.sh's `ALL_RELEASES`
-into the servers/ directory. To install a specific version, set the `KAFKA_VERSION` variable:
-
-.. code:: bash
-
-    KAFKA_VERSION=1.0.1 ./build_integration.sh
+    KAFKA_VERSION=3.9.0 make test
 
-Then to run the tests against a specific Kafka version, simply set the `KAFKA_VERSION`
-env variable to the server build you want to use for testing:
-
-.. code:: bash
-
-    KAFKA_VERSION=1.0.1 tox -e py36
-
-To test against the kafka source tree, set KAFKA_VERSION=trunk
-[optionally set SCALA_VERSION (defaults to the value set in `build_integration.sh`)]
-
-.. code:: bash
 
-    SCALA_VERSION=2.12 KAFKA_VERSION=trunk ./build_integration.sh
-    KAFKA_VERSION=trunk tox -e py36
+Integration tests start Kafka and Zookeeper fixtures. Make will download
+kafka server binaries automatically if needed.
diff --git a/pytest.ini b/pytest.ini
new file mode 100644
index 0000000..f545887
--- /dev/null
+++ b/pytest.ini
@@ -0,0 +1,2 @@
+[pytest]
+log_format = %(created)f %(filename)-23s %(threadName)s %(message)s
diff --git a/test/record/test_default_records.py b/test/record/test_default_records.py
index c3a7b02..e1c840f 100644
--- a/test/record/test_default_records.py
+++ b/test/record/test_default_records.py
@@ -1,7 +1,10 @@
 # -*- coding: utf-8 -*-
 from __future__ import unicode_literals
 import pytest
-from mock import patch
+try:
+    from unittest.mock import patch
+except ImportError:
+    from mock import patch
 import kafka.codec
 from kafka.record.default_records import (
     DefaultRecordBatch, DefaultRecordBatchBuilder
diff --git a/test/record/test_legacy_records.py b/test/record/test_legacy_records.py
index 43970f7..b15b537 100644
--- a/test/record/test_legacy_records.py
+++ b/test/record/test_legacy_records.py
@@ -1,6 +1,9 @@
 from __future__ import unicode_literals
 import pytest
-from mock import patch
+try:
+    from unittest.mock import patch
+except ImportError:
+    from mock import patch
 from kafka.record.legacy_records import (
     LegacyRecordBatch, LegacyRecordBatchBuilder
 )
diff --git a/test/sasl/test_msk.py b/test/sasl/test_msk.py
new file mode 100644
index 0000000..297ca84
--- /dev/null
+++ b/test/sasl/test_msk.py
@@ -0,0 +1,67 @@
+import datetime
+import json
+
+from kafka.sasl.msk import AwsMskIamClient
+
+try:
+    from unittest import mock
+except ImportError:
+    import mock
+
+
+def client_factory(token=None):
+    now = datetime.datetime.utcfromtimestamp(1629321911)
+    with mock.patch('kafka.sasl.msk.datetime') as mock_dt:
+        mock_dt.datetime.utcnow = mock.Mock(return_value=now)
+        return AwsMskIamClient(
+            host='localhost',
+            access_key='XXXXXXXXXXXXXXXXXXXX',
+            secret_key='XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX',
+            region='us-east-1',
+            token=token,
+        )
+
+
+def test_aws_msk_iam_client_permanent_credentials():
+    client = client_factory(token=None)
+    msg = client.first_message()
+    assert msg
+    assert isinstance(msg, bytes)
+    actual = json.loads(msg)
+
+    expected = {
+        'version': '2020_10_22',
+        'host': 'localhost',
+        'user-agent': 'kafka-python',
+        'action': 'kafka-cluster:Connect',
+        'x-amz-algorithm': 'AWS4-HMAC-SHA256',
+        'x-amz-credential': 'XXXXXXXXXXXXXXXXXXXX/20210818/us-east-1/kafka-cluster/aws4_request',
+        'x-amz-date': '20210818T212511Z',
+        'x-amz-signedheaders': 'host',
+        'x-amz-expires': '900',
+        'x-amz-signature': '0fa42ae3d5693777942a7a4028b564f0b372bafa2f71c1a19ad60680e6cb994b',
+    }
+    assert actual == expected
+
+
+def test_aws_msk_iam_client_temporary_credentials():
+    client = client_factory(token='XXXXX')
+    msg = client.first_message()
+    assert msg
+    assert isinstance(msg, bytes)
+    actual = json.loads(msg)
+
+    expected = {
+        'version': '2020_10_22',
+        'host': 'localhost',
+        'user-agent': 'kafka-python',
+        'action': 'kafka-cluster:Connect',
+        'x-amz-algorithm': 'AWS4-HMAC-SHA256',
+        'x-amz-credential': 'XXXXXXXXXXXXXXXXXXXX/20210818/us-east-1/kafka-cluster/aws4_request',
+        'x-amz-date': '20210818T212511Z',
+        'x-amz-signedheaders': 'host',
+        'x-amz-expires': '900',
+        'x-amz-signature': 'b0619c50b7ecb4a7f6f92bd5f733770df5710e97b25146f97015c0b1db783b05',
+        'x-amz-security-token': 'XXXXX',
+    }
+    assert actual == expected
diff --git a/test/test_conn.py b/test/test_conn.py
index 959cbb4..ea88fd0 100644
--- a/test/test_conn.py
+++ b/test/test_conn.py
@@ -4,7 +4,10 @@ from __future__ import absolute_import
 from errno import EALREADY, EINPROGRESS, EISCONN, ECONNRESET
 import socket
 
-import mock
+try:
+    from unittest import mock
+except ImportError:
+    import mock
 import pytest
 
 from kafka.conn import BrokerConnection, ConnectionStates, collect_hosts
diff --git a/test/test_consumer_integration.py b/test/test_consumer_integration.py
index 5aeb63d..af8ec68 100644
--- a/test/test_consumer_integration.py
+++ b/test/test_consumer_integration.py
@@ -1,7 +1,10 @@
 import logging
 import time
 
-from mock import patch, ANY
+try:
+    from unittest.mock import patch, ANY
+except ImportError:
+    from mock import patch, ANY
 import pytest
 from kafka.vendor.six.moves import range