def test_put(self):
in_five_min_epoch = int(time.time()) + 5 * 60
aws_credentials = AWSCredentials(
access_key_id="test_aki",
secret_access_key="test_sak",
session_token="test_st",
expiration=in_five_min_epoch,
)
cache = AWSCredentialsCache()
cache.put(
credentials=aws_credentials,
account_id="123456789012",
role_name="test_rn",
role_session_name="test_rsn",
)
keys_values = cache.cache.items()
self.assertEqual(len(keys_values), 1)
for key, value in keys_values:
self.assertEqual(
key,
AWSCredentialsCache.build_cache_key(
account_id="123456789012",
role_name="test_rn",
role_session_name="test_rsn"),
)
self.assertEqual(
value,
AWSCredentials(
access_key_id="test_aki",
secret_access_key="test_sak",
session_token="test_st",
expiration=in_five_min_epoch,
),
)
def test_is_expired_when_not_expired(self):
in_five_min_epoch = int(time.time()) + 5 * 60
aws_credentials = AWSCredentials(
access_key_id="test",
secret_access_key="test",
session_token="test",
expiration=in_five_min_epoch,
)
self.assertFalse(aws_credentials.is_expired())
def test_get_session_without_region(self):
in_five_min_epoch = int(time.time()) + 5 * 60
aws_credentials = AWSCredentials(
access_key_id="test",
secret_access_key="test",
session_token="test",
expiration=in_five_min_epoch,
)
session = aws_credentials.get_session()
self.assertIsNone(session.region_name)
def test_get_session_with_region(self):
in_five_min_epoch = int(time.time()) + 5 * 60
aws_credentials = AWSCredentials(
access_key_id="test",
secret_access_key="test",
session_token="test",
expiration=in_five_min_epoch,
)
session = aws_credentials.get_session(region_name="us-east-4")
self.assertEqual(session.region_name, "us-east-4")
def test_from_dict(self):
in_five_min_epoch = int(time.time()) + 5 * 60
aws_credentials = AWSCredentials(
access_key_id="testaki",
secret_access_key="testsak",
session_token="testst",
expiration=in_five_min_epoch,
)
data = {
"access_key_id": "testaki",
"secret_access_key": "testsak",
"session_token": "testst",
"expiration": in_five_min_epoch,
}
aws_credentials_from_dict = AWSCredentials(**data)
self.assertEqual(aws_credentials, aws_credentials_from_dict)
def test_get_with_region(self):
in_five_min_epoch = int(time.time()) + 5 * 60
aws_credentials = AWSCredentials(
access_key_id="test_aki",
secret_access_key="test_sak",
session_token="test_st",
expiration=in_five_min_epoch,
)
cache = AWSCredentialsCache()
cache.put(
credentials=aws_credentials,
account_id="123456789012",
role_name="test_rn",
role_session_name="test_rsn",
)
session = cache.get(
account_id="123456789012",
role_name="test_rn",
role_session_name="test_rsn",
region_name="us-east-4",
)
self.assertEqual(session.region_name, "us-east-4")
creds = session.get_credentials()
self.assertEqual(creds.access_key, "test_aki")
self.assertEqual(
creds.secret_key,
"test_sak",
)
self.assertEqual(creds.token, "test_st")
def test_from_dict(self):
cache = AWSCredentialsCache()
in_five_min_epoch = int(time.time()) + 5 * 60
aws_credentials = AWSCredentials(
access_key_id="test_aki",
secret_access_key="test_sak",
session_token="test_st",
expiration=in_five_min_epoch,
)
cache.put(
credentials=aws_credentials,
account_id="123456789012",
role_name="test_rn",
role_session_name="test_rsn",
)
aws_credentials = AWSCredentials(
access_key_id="test_aki2",
secret_access_key="test_sak2",
session_token="test_st2",
expiration=in_five_min_epoch,
)
cache.put(
credentials=aws_credentials,
account_id="123456789012",
role_name="test_rn2",
role_session_name="test_rsn2",
)
data = {
"cache": {
"123456789012:test_rn:test_rsn": {
"access_key_id": "test_aki",
"secret_access_key": "test_sak",
"session_token": "test_st",
"expiration": in_five_min_epoch,
},
"123456789012:test_rn2:test_rsn2": {
"access_key_id": "test_aki2",
"secret_access_key": "test_sak2",
"session_token": "test_st2",
"expiration": in_five_min_epoch,
},
}
}
from_data_cache = AWSCredentialsCache(**data)
self.assertEqual(from_data_cache, cache)
def test_put_wrong_account_id(self):
in_five_min_epoch = int(time.time()) + 5 * 60
aws_credentials = AWSCredentials(
access_key_id="test_aki",
secret_access_key="test_sak",
session_token="test_st",
expiration=in_five_min_epoch,
)
cache = AWSCredentialsCache()
with self.assertRaises(ValueError):
cache.put(
credentials=aws_credentials,
account_id="987654321098",
role_name="test_rn",
role_session_name="test_rsn",
)
def test_get_expired_returns_none(self):
five_min_ago_epoch = int(time.time()) - 5 * 60
aws_credentials = AWSCredentials(
access_key_id="test_aki",
secret_access_key="test_sak",
session_token="test_st",
expiration=five_min_ago_epoch,
)
cache = AWSCredentialsCache()
cache.put(
credentials=aws_credentials,
account_id="123456789012",
role_name="test_rn",
role_session_name="test_rsn",
)
session = cache.get(account_id="123456789012",
role_name="test_rn",
role_session_name="test_rsn")
self.assertIsNone(session)
def get_session(
self,
account_id: str,
region_name: Optional[str] = None,
credentials_cache: Optional[AWSCredentialsCache] = None,
) -> boto3.Session:
"""Get a session for an account_id by iterating through the :class:`.AccessStep`s
of this :class:`.MultiHopAccessor`.
Args:
account_id: account to access
region_name: region to use during session creation.
Returns:
boto3 Session for accessing account_id
"""
logger = Logger()
cws = boto3.Session(region_name=region_name)
for access_step in self.access_steps:
access_account_id = access_step.account_id if access_step.account_id else account_id
role_name = access_step.role_name
external_id = access_step.external_id
session = None
if credentials_cache is not None:
session = credentials_cache.get(
account_id=access_account_id,
role_name=role_name,
role_session_name=self.role_session_name,
region_name=region_name,
)
if session is None:
logger.debug(event=LogEvent.AuthToAccountStart)
sts_client = cws.client("sts")
role_arn = f"arn:aws:iam::{access_account_id}:role/{role_name}"
assume_args = {
"RoleArn": role_arn,
"RoleSessionName": self.role_session_name
}
if external_id:
assume_args["ExternalId"] = external_id
assume_resp = sts_client.assume_role(**assume_args)
creds = assume_resp["Credentials"]
expiration_datetime = creds["Expiration"]
credentials = AWSCredentials(
access_key_id=creds["AccessKeyId"],
secret_access_key=creds["SecretAccessKey"],
session_token=creds["SessionToken"],
expiration=int(expiration_datetime.timestamp()),
)
session = boto3.Session(
aws_access_key_id=creds["AccessKeyId"],
aws_secret_access_key=creds["SecretAccessKey"],
aws_session_token=creds["SessionToken"],
region_name=region_name,
)
if credentials_cache is not None:
credentials_cache.put(
credentials=credentials,
account_id=access_account_id,
role_name=role_name,
role_session_name=self.role_session_name,
)
logger.debug(event=LogEvent.AuthToAccountEnd)
cws = session
return session