fixes #221 - minor rework of PR #222 to retain compatibility with cre…

…ds dicts

CHANGES.rst

@@ -15,6 +15,15 @@ Changelog
   capacity reservation). For the time being, they'll be ignored (i.e. not subtracted
   from On-Demand Running Instances) but a debug-level message will be logged
   for each.
+* `#221 <https://github.com/jantman/awslimitchecker/issues/221>`_ /
+  `PR #222 <https://github.com/jantman/awslimitchecker/pull/222>`_ - Fix bug
+  in handling of STS Credentials where they are cached permanently in
+  ``connectable.Connectable.credentials``, and new AwsLimitChecker instances
+  in the same Python process reuse the first set of STS credentials. This is
+  fixed by storing the Account ID as part of
+  ``connectable.ConnectableCredentials`` and getting new STS creds if the cached
+  account ID does not match the current ``account_id`` on the ``Connectable``
+  object.
 
 0.5.1 (2016-09-25)
 ------------------

awslimitchecker/connectable.py

@@ -50,6 +50,9 @@ class ConnectableCredentials(object):
     `boto3.sts.STSConnection.assume_role <https://boto3.readthedocs.org/en/
     latest/reference/services/sts.html#STS.Client.assume_role>`_ just returns
     a dict. This class provides a compatible interface for boto3.
+
+    We also maintain an ``account_id`` attribute that can be set to the
+    account ID, to ensure that credentials are updated when switching accounts.
     """
 
     def __init__(self, creds_dict):
@@ -59,7 +62,7 @@ def __init__(self, creds_dict):
         self.expiration = creds_dict['Credentials']['Expiration']
         self.assumed_role_id = creds_dict['AssumedRoleUser']['AssumedRoleId']
         self.assumed_role_arn = creds_dict['AssumedRoleUser']['Arn']
-        self.account_id = creds_dict['account_id']
+        self.account_id = None
 
 
 class Connectable(object):
@@ -98,11 +101,10 @@ def _boto3_connection_kwargs(self):
                     logger.debug("Reusing previous STS credentials for "
                                  "account %s", self.account_id)
                 else:
-                    logger.debug("Previous STS credentials are for account %s",
-                                 Connectable.credentials.account_id)
-                    logger.debug("Connecting for account %s role '%s' with STS "
-                                 "(region: %s)", self.account_id,
-                                 self.account_role, self.region)
+                    logger.debug("Previous STS credentials are for account %s; "
+                                 "getting new credentials for current account "
+                                 "(%s)", Connectable.credentials.account_id,
+                                 self.account_id)
                     Connectable.credentials = self._get_sts_token()
             kwargs['aws_access_key_id'] = Connectable.credentials.access_key
             kwargs['aws_secret_access_key'] = Connectable.credentials.secret_key
@@ -177,10 +179,9 @@ def _get_sts_token(self):
             assume_kwargs['TokenCode'] = self.mfa_token
         role = sts.assume_role(**assume_kwargs)
 
-        role['account_id'] = self.account_id
-
         creds = ConnectableCredentials(role)
+        creds.account_id = self.account_id
 
-        logger.debug("Got STS credentials for role; access_key_id=%s",
-                     creds.access_key)
+        logger.debug("Got STS credentials for role; access_key_id=%s "
+                     "(account_id=%s)", creds.access_key, creds.account_id)
         return creds

awslimitchecker/tests/test_connectable.py

@@ -137,6 +137,7 @@ def test_boto3_connection_kwargs_sts_again(self):
         type(mock_creds).access_key = 'sts_ak'
         type(mock_creds).secret_key = 'sts_sk'
         type(mock_creds).session_token = 'sts_token'
+        type(mock_creds).account_id = '123'
 
         with patch('%s._get_sts_token' % pb) as mock_get_sts:
             with patch('%s.logger' % pbm) as mock_logger:
@@ -154,6 +155,33 @@ def test_boto3_connection_kwargs_sts_again(self):
             'aws_session_token': 'sts_token'
         }
 
+    def test_boto3_connection_kwargs_sts_again_other_account(self):
+        cls = ConnectableTester(account_id='123', account_role='myrole',
+                                region='myregion')
+        mock_creds = Mock()
+        type(mock_creds).access_key = 'sts_ak'
+        type(mock_creds).secret_key = 'sts_sk'
+        type(mock_creds).session_token = 'sts_token'
+        type(mock_creds).account_id = '456'
+
+        with patch('%s._get_sts_token' % pb) as mock_get_sts:
+            with patch('%s.logger' % pbm) as mock_logger:
+                mock_get_sts.return_value = mock_creds
+                Connectable.credentials = mock_creds
+                res = cls._boto3_connection_kwargs
+        assert mock_get_sts.mock_calls == [call()]
+        assert mock_logger.mock_calls == [
+            call.debug("Previous STS credentials are for account %s; "
+                       "getting new credentials for current account "
+                       "(%s)", '456', '123')
+        ]
+        assert res == {
+            'region_name': 'myregion',
+            'aws_access_key_id': 'sts_ak',
+            'aws_secret_access_key': 'sts_sk',
+            'aws_session_token': 'sts_token'
+        }
+
     def test_connect(self):
         mock_conn = Mock()
         mock_cc = Mock()