AWS MSK IAM SASL Signer (Python)

Warnings

• gotcha When using `confluent-kafka-python`, the `sasl.oauthbearer.token.cb` (callback) expects the token expiry time in *seconds* since the epoch, while `MSKAuthTokenProvider.generate_auth_token` returns it in *milliseconds*. You must divide the returned expiry by 1000 before returning it from your callback.
  Fix: Divide `expiry_ms` by 1000 in your `oauth_cb` implementation before returning it (e.g., `return auth_token, expiry_ms / 1000`).
• breaking For users of `dpkp/kafka-python` client, versions 2.1.0 and later introduced a change in SASL module handling, including the `AWS_MSK_IAM` mechanism. The `aws-msk-iam-sasl-signer-python` library is designed for the `OAUTHBEARER` mechanism, and older examples might not work directly with `kafka-python` versions 2.1.0+.
  Fix: Ensure your Kafka client configuration specifies `sasl_mechanism='OAUTHBEARER'` and not `AWS_MSK_IAM` when using this signer library. Refer to the updated examples in the GitHub README or the AWS documentation. For `kafka-python`, use `sasl_oauth_token_provider` with your custom token provider class. For `confluent-kafka-python`, use `sasl.oauthbearer.token.cb`.
• gotcha Incorrect Kafka SASL mechanism: This library facilitates IAM authentication for MSK using the `SASL_OAUTHBEARER` mechanism. Some users mistakenly attempt to configure their Kafka clients with `sasl_mechanism='AWS_MSK_IAM'`, which is a custom mechanism primarily for Java clients or specific forks. Python clients `dpkp/kafka-python` and `confluent-kafka-python` do not natively support `AWS_MSK_IAM` and require `OAUTHBEARER` when used with this library.
  Fix: Always set `sasl_mechanism='OAUTHBEARER'` in your Python Kafka client configuration when using `aws-msk-iam-sasl-signer-python`.
• gotcha Network connectivity issues (e.g., 'Connection reset') are frequent if the client application (e.g., Lambda function, EC2 instance) is not in the same VPC as the MSK cluster or if the security groups are improperly configured. Traffic on TCP port 9098 (for SASL_SSL) must be allowed between the client and the MSK brokers.
  Fix: Verify that your client and MSK cluster are in the same VPC or have appropriate VPC peering/connectivity. Ensure security groups allow inbound traffic on port 9098 from your client's security group to the MSK broker security group.
• gotcha Disabling TLS host verification (`ssl_context.check_hostname = False`, `ssl_context.verify_mode = ssl.CERT_NONE`) in some client examples (e.g., `aiokafka`) can introduce security vulnerabilities. This should only be done if you fully understand and accept the risks.
  Fix: Avoid disabling TLS host verification unless absolutely necessary and with a clear understanding of the security implications. If using a custom SSL context, ensure proper certificate validation is in place.
• gotcha Transactional producers may not work correctly with IAM authentication. There's an open issue reporting connection failures when `init_transactions` is called, suggesting a potential timing issue with connection establishment.
  Fix: If using transactional producers, ensure a connection is established by polling or flushing before calling `init_transactions`. Consult the GitHub issues for any updates or workarounds.

Install

• pip install aws-msk-iam-sasl-signer-python Install stable version

Imports

• MSKAuthTokenProvider

  from aws_msk_iam_sasl_signer import MSKAuthTokenProvider

  The `MSKAuthTokenProvider` class is directly available under the `aws_msk_iam_sasl_signer` package, not within a nested module of the same name.

Quickstart

This quickstart demonstrates how to configure a `confluent-kafka-python` consumer to connect to an Amazon MSK cluster using IAM authentication via `aws-msk-iam-sasl-signer-python`. It defines a callback function `oauth_cb` which uses `MSKAuthTokenProvider.generate_auth_token` to retrieve the necessary SASL/OAUTHBEARER token. Ensure that your AWS credentials are configured (e.g., via `~/.aws/credentials`, environment variables, or IAM role for EC2/Lambda) and that `KAFKA_BOOTSTRAP_SERVERS` and `AWS_REGION` environment variables are set.

import os
import socket
import time
from confluent_kafka import Consumer, KafkaException
from aws_msk_iam_sasl_signer import MSKAuthTokenProvider

def oauth_cb(oauth_config):
    # MSKAuthTokenProvider.generate_auth_token returns expiry in milliseconds
    # confluent-kafka-python expects expiry in seconds since epoch
    aws_region = os.environ.get('AWS_REGION', 'us-east-1') # Or specific region for your MSK cluster
    try:
        auth_token, expiry_ms = MSKAuthTokenProvider.generate_auth_token(aws_region)
        return auth_token, expiry_ms / 1000
    except Exception as e:
        print(f"Error generating token: {e}")
        raise KafkaException(f"Failed to get MSK Auth Token: {e}")


# Configure Kafka Consumer
# Ensure KAFKA_BOOTSTRAP_SERVERS environment variable is set (e.g., b-1.yourcluster.abcdef.c5.kafka.us-east-1.amazonaws.com:9098)
# Ensure AWS_REGION environment variable is set
# Ensure your AWS credentials (e.g., via ~/.aws/credentials or environment variables) are configured

consumer_conf = {
    'bootstrap.servers': os.environ.get('KAFKA_BOOTSTRAP_SERVERS', 'localhost:9092'),
    'client.id': socket.gethostname(),
    'group.id': os.environ.get('KAFKA_GROUP_ID', 'my-consumer-group'),
    'security.protocol': 'SASL_SSL',
    'sasl.mechanism': 'OAUTHBEARER',
    'sasl.oauthbearer.token.cb': oauth_cb,
    'debug': 'broker,protocol,security' # Optional: for debugging connection issues
}

topic = os.environ.get('KAFKA_TOPIC', 'my-test-topic')

consumer = None
try:
    consumer = Consumer(consumer_conf)
    consumer.subscribe([topic])

    print(f"Consumer configured for topic: {topic}")
    print(f"Consuming messages. Press Ctrl+C to exit.")

    while True:
        msg = consumer.poll(timeout=1.0) # Poll for messages, 1-second timeout
        if msg is None:
            continue
        if msg.error():
            if msg.error().code() == KafkaException._PARTITION_EOF:
                # End of partition event - not an error
                print(f"%% {msg.topic()} [{msg.partition()}] reached end offset {msg.offset()}")
            elif msg.error():
                raise KafkaException(msg.error())
        else:
            print(f"Received message: key={msg.key().decode('utf-8') if msg.key() else 'None'}, value={msg.value().decode('utf-8')}, topic={msg.topic()}, partition={msg.partition()}, offset={msg.offset()}")

except KafkaException as e:
    print(f"Kafka error: {e}")
except Exception as e:
    print(f"An unexpected error occurred: {e}")
finally:
    if consumer:
        print("Closing consumer...")
        consumer.close()