from . import cloudformation
from . import servicecatalog
from . import organizations
from . import guardduty
from . import codebuild
from . import codecommit
from boto3.session import Session
import logging
logger = logging.getLogger(__file__)
def make_better(service_name, client):
if service_name == 'cloudformation':
return cloudformation.make_better(client)
elif service_name == 'servicecatalog':
return servicecatalog.make_better(client)
elif service_name == 'organizations':
return organizations.make_better(client)
elif service_name == 'guardduty':
return guardduty.make_better(client)
elif service_name == 'codebuild':
return codebuild.make_better(client)
elif service_name == 'codecommit':
return codecommit.make_better(client)
return client
[docs]class ClientContextManager(object):
"""
ClientContextManager allows you to use boto3 client as a python context manager.
This allows you to perform the following::
with ClientContextManager('cloudformation') as cloudformation_client:
cloudformation_client.create_stack(**args)
"""
def __init__(self, service_name, **kwargs):
super().__init__()
self.service_name = service_name
self.kwargs = kwargs
def __enter__(self):
self.client = Session().client(
self.service_name,
**self.kwargs
)
self.client = make_better(self.service_name, self.client)
return self.client
def __exit__(self, *args, **kwargs):
self.client = None
[docs]class MultiRegionClientContextManager(object):
"""
MultiRegionClientContextManager allows you to use boto3 client as a python context manager for multiple regions.
This allows you to perform the following::
with MultiRegionClientContextManager('cloudformation', ['us-east-1','eu-west-1']) as cloudformation_clients:
for region_name, cloudformation_client in cloudformation_clients.items():
cloudformation_client.create_stack(**args)
If you want to deploy to multiple regions at the same time then you should use Python Threads
"""
def __init__(self, service_name, regions, **kwargs):
super().__init__()
self.service_name = service_name
self.regions = regions
self.clients = {}
self.kwargs = kwargs
def __enter__(self):
for region in self.regions:
c = make_better(
self.service_name,
Session().client(
self.service_name,
region_name=region,
**self.kwargs
)
)
self.clients[region] = c
return self.clients
def __exit__(self, *args, **kwargs):
self.clients = None
[docs]class CrossAccountClientContextManager(object):
"""
CrossAccountClientContextManager allows you to use boto3 client as a python context manager for another account.
This allows you to perform the following::
with CrossAccountClientContextManager(
'cloudformation',
'arn:aws:iam::0123456789010:role/deployer',
'deployment_account_session',
) as deployment_account_cloudformation:
deployment_account_cloudformation.create_stack(**args)
"""
def __init__(self, service_name, role_arn, role_session_name, **kwargs):
super().__init__()
self.service_name = service_name
self.role_arn = role_arn
self.role_session_name = role_session_name
self.kwargs = kwargs
def __enter__(self):
sts = Session().client('sts')
assumed_role_object = sts.assume_role(
RoleArn=self.role_arn,
RoleSessionName=self.role_session_name,
)
self.credentials = assumed_role_object['Credentials']
kwargs = {
"service_name": self.service_name,
"aws_access_key_id": self.credentials['AccessKeyId'],
"aws_secret_access_key": self.credentials['SecretAccessKey'],
"aws_session_token": self.credentials['SessionToken'],
}
if self.kwargs is not None:
kwargs.update(self.kwargs)
self.client = Session().client(**kwargs)
self.client = make_better(self.service_name, self.client)
return self.client
def __exit__(self, *args, **kwargs):
self.client = None
[docs]class CrossMultipleAccountsClientContextManager(object):
"""
CrossMultipleAccountsClientContextManager allows you to use boto3 client as a python context manager for another account.
This allows you to perform the following::
with CrossMultipleAccountsClientContextManager(
'cloudformation',
[
('arn:aws:iam::0123456789010:role/deployer', 'deployment_account_session'),
('arn:aws:iam::097167856333:role/deployer', 'deployment_account_session_nested'),
],
) as deployment_account_cloudformation:
deployment_account_cloudformation.create_stack(**args)
"""
def __init__(self, service_name, assumable_details, **kwargs):
super().__init__()
self.service_name = service_name
self.assumable_details = assumable_details
self.kwargs = kwargs
def __enter__(self):
credentials = {}
for assumable_detail in self.assumable_details:
role_arn, role_session_name = assumable_detail
logger.info('About to assume: {} with session name: {}'.format(role_arn, role_session_name))
sts = Session().client('sts', **credentials)
assumed_role_object = sts.assume_role(
RoleArn=role_arn,
RoleSessionName=role_session_name,
)
new_credentials = assumed_role_object['Credentials']
credentials = {
"aws_access_key_id": new_credentials['AccessKeyId'],
"aws_secret_access_key": new_credentials['SecretAccessKey'],
"aws_session_token": new_credentials['SessionToken'],
}
self.client = Session().client(self.service_name, **credentials)
self.client = make_better(self.service_name, self.client)
return self.client
def __exit__(self, *args, **kwargs):
self.client = None