Update region resolution logic and refactor / add more tests

Author: danielpops

kafka/msk.py

@@ -34,6 +34,11 @@ def __init__(self, host, boto_session):
         self.host = host
         self.boto_session = boto_session
 
+        # This will raise if the region can't be determined
+        # Do this during init instead of waiting for failures downstream
+        if self.region:
+            pass
+
     @property
     def access_key(self):
         return self.boto_session.get_credentials().access_key
@@ -48,11 +53,21 @@ def token(self):
 
     @property
     def region(self):
-        # TODO: This logic is not perfect and should be revisited
+        # Try to get the region information from the broker hostname
         for host in self.host.split(','):
             if 'amazonaws.com' in host:
                 return host.split('.')[-3]
-        return 'us-west-2'
+
+        # If the region can't be determined from hostname, try the boto session
+        # This will only have a value if:
+        #  - `AWS_DEFAULT_REGION` environment variable is set
+        #  - `~/.aws/config` region variable is set
+        region = self.boto_session.get_config_variable('region')
+        if region:
+            return region
+
+        # Otherwise give up
+        raise Exception('Could not determine region from broker host(s) or aws configuration')
 
     @property
     def _credential(self):

test/test_msk.py

@@ -13,21 +13,92 @@
     import mock
 
 
-@pytest.fixture(params=[{'session_token': 'session_token', 'host': 'localhost'}, {'session_token': None, 'host': 'localhost.us-east-1.amazonaws.com'}])
-def msk_client(request):
+@pytest.fixture
+def boto_session():
     # To avoid a package dependency on the optional botocore library, we mock the module out
     sys.modules['botocore.session'] = mock.MagicMock()
     from botocore.session import Session  # pylint: disable=import-error
 
-    session = Session()
-    session.get_credentials = mock.MagicMock(return_value=mock.MagicMock(id='the_actual_credentials', access_key='akia', secret_key='secret', token=request.param['session_token']))
-    yield AwsMskIamClient(
-        host=request.param["host"],
-        boto_session = session,
+    boto_session = Session()
+    boto_session.get_credentials = mock.MagicMock(return_value=mock.MagicMock(id='the_actual_credentials', access_key='akia', secret_key='secret', token=None))
+    yield boto_session
+
+
+def test_aws_msk_iam_region_from_config(boto_session):
+    # Region determined by configuration
+    boto_session.get_config_variable = mock.MagicMock(return_value='us-west-2')
+    msk_client = AwsMskIamClient(
+        host='localhost',
+        boto_session = boto_session,
     )
+    msg = msk_client.first_message()
+    assert msg
+    assert isinstance(msg, bytes)
+    actual = json.loads(msg.decode('utf-8'))
+
+    expected = {
+        'version': '2020_10_22',
+        'host': msk_client.host,
+        'user-agent': 'kafka-python',
+        'action': 'kafka-cluster:Connect',
+        'x-amz-algorithm': 'AWS4-HMAC-SHA256',
+        'x-amz-credential': '{}/{}/us-west-2/kafka-cluster/aws4_request'.format(msk_client.access_key, datetime.datetime.utcnow().strftime('%Y%m%d')),
+        'x-amz-date': mock.ANY,
+        'x-amz-signedheaders': 'host',
+        'x-amz-expires': '900',
+        'x-amz-signature': mock.ANY,
+    }
+    TestCase().assertEqual(actual, expected)
 
 
-def test_aws_msk_iam(msk_client):
+def test_aws_msk_iam_region_from_hostname(boto_session):
+    # Region determined by hostname
+    msk_client = AwsMskIamClient(
+        host='localhost.us-east-1.amazonaws.com',
+        boto_session = boto_session,
+    )
+    msg = msk_client.first_message()
+    assert msg
+    assert isinstance(msg, bytes)
+    actual = json.loads(msg.decode('utf-8'))
+
+    expected = {
+        'version': '2020_10_22',
+        'host': msk_client.host,
+        'user-agent': 'kafka-python',
+        'action': 'kafka-cluster:Connect',
+        'x-amz-algorithm': 'AWS4-HMAC-SHA256',
+        'x-amz-credential': '{}/{}/us-east-1/kafka-cluster/aws4_request'.format(msk_client.access_key, datetime.datetime.utcnow().strftime('%Y%m%d')),
+        'x-amz-date': mock.ANY,
+        'x-amz-signedheaders': 'host',
+        'x-amz-expires': '900',
+        'x-amz-signature': mock.ANY,
+    }
+    TestCase().assertEqual(actual, expected)
+
+
+def test_aws_msk_iam_no_region(boto_session):
+    # No region from config
+    boto_session.get_config_variable = mock.MagicMock(return_value=None)
+
+    with TestCase().assertRaises(Exception) as e:
+        # No region from hostname
+        msk_client = AwsMskIamClient(
+            host='localhost',
+            boto_session = boto_session,
+        )
+    assert 'Could not determine region from broker host(s) or aws configuration' == str(e.exception)
+
+
+@pytest.mark.parametrize('session_token', [(None), ('the_token')])
+def test_aws_msk_iam_permanent_and_temporary_credentials(session_token, request):
+    boto_session = request.getfixturevalue('boto_session')
+    if session_token:
+        boto_session.get_credentials.return_value.token = session_token
+    msk_client = AwsMskIamClient(
+        host='localhost.us-east-1.amazonaws.com',
+        boto_session = boto_session,
+    )
     msg = msk_client.first_message()
     assert msg
     assert isinstance(msg, bytes)
@@ -39,12 +110,12 @@ def test_aws_msk_iam(msk_client):
         'user-agent': 'kafka-python',
         'action': 'kafka-cluster:Connect',
         'x-amz-algorithm': 'AWS4-HMAC-SHA256',
-        'x-amz-credential': '{}/{}/{}/kafka-cluster/aws4_request'.format(msk_client.access_key, datetime.datetime.utcnow().strftime('%Y%m%d'), 'us-west-2' if msk_client.host == 'localhost' else 'us-east-1'),
+        'x-amz-credential': '{}/{}/us-east-1/kafka-cluster/aws4_request'.format(msk_client.access_key, datetime.datetime.utcnow().strftime('%Y%m%d')),
         'x-amz-date': mock.ANY,
         'x-amz-signedheaders': 'host',
         'x-amz-expires': '900',
         'x-amz-signature': mock.ANY,
     }
-    if msk_client.token:
-        expected['x-amz-security-token'] = msk_client.token
+    if session_token:
+        expected['x-amz-security-token'] = session_token
     TestCase().assertEqual(actual, expected)