"""Amazon MSK connector."""

import boto3


def get_kafka_config(arn, credentials):
    """Get Amazon MSK broker string.

    Args:
        arn (str): The MSK Cluster's ARN
    Returns:
        Queue: Instance of Amazon SQS queue.
    """
    kafka = boto3.client(
        'kafka',
        **credentials
    )
    broker_request = kafka.get_bootstrap_brokers(ClusterArn=arn)

    try:
        return {
            'bootstrap_servers': broker_request['BootstrapBrokerStringTls'],
            'security_protocol': 'SSL',
        }
    except KeyError:
        return {
            'bootstrap_servers': broker_request['BootstrapBrokerString'],
        }