def setUp(self):
self.environ = {
'AWS_DATA_PATH': os.environ['AWS_DATA_PATH'],
'AWS_DEFAULT_REGION': 'us-east-1',
'AWS_ACCESS_KEY_ID': 'access_key',
'AWS_SECRET_ACCESS_KEY': 'secret_key',
'AWS_CONFIG_FILE': '',
}
self.environ_patch = mock.patch('os.environ', self.environ)
self.environ_patch.start()
emitter = HierarchicalEmitter()
session = Session(EnvironmentVariables, emitter)
session.register_component('data_loader', _LOADER)
load_plugins({}, event_hooks=emitter)
driver = CLIDriver(session=session)
self.session = session
self.driver = driver
def assume_role(session: Session,
role_arn: str,
duration: int = 3600,
session_name: str = None) -> Session:
# noinspection PyTypeChecker
fetcher = AssumeRoleCredentialFetcher(session.create_client,
session.get_credentials(),
role_arn,
extra_args={
'DurationSeconds': duration,
'RoleSessionName': session_name
})
role_session = Session()
role_session.register_component(
'credential_provider',
CredentialResolver([AssumeRoleProvider(fetcher)]))
return role_session
def setUp(self):
self.environ = {
'AWS_DATA_PATH': os.environ['AWS_DATA_PATH'],
'AWS_DEFAULT_REGION': 'us-east-1',
'AWS_ACCESS_KEY_ID': 'access_key',
'AWS_SECRET_ACCESS_KEY': 'secret_key',
'AWS_CONFIG_FILE': '',
}
self.environ_patch = mock.patch('os.environ', self.environ)
self.environ_patch.start()
emitter = HierarchicalEmitter()
session = Session(EnvironmentVariables, emitter)
session.register_component('data_loader', _LOADER)
load_plugins({}, event_hooks=emitter)
driver = CLIDriver(session=session)
self.session = session
self.driver = driver
def create_session(self, profile=None):
session = Session(profile=profile)
# We have to set bogus credentials here or otherwise we'll trigger
# an early credential chain resolution.
sts = session.create_client(
'sts',
aws_access_key_id='spam',
aws_secret_access_key='eggs',
)
stubber = Stubber(sts)
stubber.activate()
assume_role_provider = AssumeRoleProvider(
load_config=lambda: session.full_config,
client_creator=lambda *args, **kwargs: sts,
cache={},
profile_name=profile,
credential_sourcer=CanonicalNameCredentialSourcer([
self.env_provider, self.container_provider,
self.metadata_provider
])
)
component_name = 'credential_provider'
resolver = session.get_component(component_name)
available_methods = [p.METHOD for p in resolver.providers]
replacements = {
'env': self.env_provider,
'iam-role': self.metadata_provider,
'container-role': self.container_provider,
'assume-role': assume_role_provider
}
for name, provider in replacements.items():
try:
index = available_methods.index(name)
except ValueError:
# The provider isn't in the session
continue
resolver.providers[index] = provider
session.register_component(
'credential_provider', resolver
)
return session, stubber
def create_session(self, profile=None):
session = Session(profile=profile)
# We have to set bogus credentials here or otherwise we'll trigger
# an early credential chain resolution.
sts = session.create_client(
'sts',
aws_access_key_id='spam',
aws_secret_access_key='eggs',
)
self.mock_client_creator.return_value = sts
stubber = Stubber(sts)
stubber.activate()
assume_role_provider = AssumeRoleProvider(
load_config=lambda: session.full_config,
client_creator=self.mock_client_creator,
cache={},
profile_name=profile,
credential_sourcer=CanonicalNameCredentialSourcer([
self.env_provider, self.container_provider,
self.metadata_provider
]),
profile_provider_builder=ProfileProviderBuilder(session),
)
component_name = 'credential_provider'
resolver = session.get_component(component_name)
available_methods = [p.METHOD for p in resolver.providers]
replacements = {
'env': self.env_provider,
'iam-role': self.metadata_provider,
'container-role': self.container_provider,
'assume-role': assume_role_provider
}
for name, provider in replacements.items():
try:
index = available_methods.index(name)
except ValueError:
# The provider isn't in the session
continue
resolver.providers[index] = provider
session.register_component('credential_provider', resolver)
return session, stubber
def setup_aws_client(config):
role_arn = "arn:aws:iam::{}:role/{}".format(
config['account_id'].replace('-', ''), config['role_name'])
session = Session()
fetcher = AssumeRoleCredentialFetcher(session.create_client,
session.get_credentials(),
role_arn,
extra_args={
'DurationSeconds': 3600,
'RoleSessionName': 'TapS3CSV',
'ExternalId':
config['external_id']
},
cache=JSONFileCache())
refreshable_session = Session()
refreshable_session.register_component(
'credential_provider',
CredentialResolver([AssumeRoleProvider(fetcher)]))
LOGGER.info("Attempting to assume_role on RoleArn: %s", role_arn)
boto3.setup_default_session(botocore_session=refreshable_session)
def _get_boto3_session(region: str,
role_arn: str = None,
assume_duration: int = 3600) -> Session:
"""Creates a boto3 session, optionally assuming a role.
Args:
region: The AWS region for the session.
role_arn: The ARN to assume for the session.
assume_duration: The duration (in seconds) to assume the role.
Returns:
object: A boto3 Session.
"""
# By default return a basic session
if not role_arn:
return Session(region_name=region)
# The following assume role example was taken from
# https://github.com/boto/botocore/issues/761#issuecomment-426037853
# Create a session used to assume role
assume_session = BotocoreSession()
fetcher = AssumeRoleCredentialFetcher(
assume_session.create_client,
assume_session.get_credentials(),
role_arn,
extra_args={
"DurationSeconds": assume_duration,
},
cache=JSONFileCache(),
)
role_session = BotocoreSession()
role_session.register_component(
"credential_provider",
CredentialResolver([Boto3Manager.AssumeRoleProvider(fetcher)]),
)
return Session(region_name=region, botocore_session=role_session)
def setup_aws_client(config):
role_arn = "arn:aws:iam::{}:role/{}".format(
config["account_id"].replace("-", ""), config["role_name"])
session = Session()
fetcher = AssumeRoleCredentialFetcher(
session.create_client,
session.get_credentials(),
role_arn,
extra_args={
"DurationSeconds": 3600,
"RoleSessionName": "TapS3CSV",
"ExternalId": config["external_id"],
},
cache=JSONFileCache(),
)
refreshable_session = Session()
refreshable_session.register_component(
"credential_provider",
CredentialResolver([AssumeRoleProvider(fetcher)]))
LOGGER.info("Attempting to assume_role on RoleArn: %s", role_arn)
boto3.setup_default_session(botocore_session=refreshable_session)
def setup_aws_client(config):
if 'role_name' in config:
role_arn = "arn:aws:iam::{}:role/{}".format(
config['account_id'].replace('-', ''), config['role_name'])
session = Session()
fetcher = AssumeRoleCredentialFetcher(session.create_client,
session.get_credentials(),
role_arn,
extra_args={
'DurationSeconds':
3600,
'RoleSessionName':
'TapDynamodDB',
'ExternalId':
config['external_id']
},
cache=JSONFileCache())
refreshable_session = Session()
refreshable_session.register_component(
'credential_provider',
CredentialResolver([AssumeRoleProvider(fetcher)]))
LOGGER.info("Attempting to assume_role on RoleArn: %s", role_arn)
boto3.setup_default_session(botocore_session=refreshable_session)
elif 'aws_access_key_id' in config and 'aws_secret_access_key' in config:
LOGGER.info(
"Attempting to pass AWS credentials from 'aws_access_key_id' and 'aws_secret_access_key' config values"
)
boto3.setup_default_session(
aws_access_key_id=config['aws_access_key_id'],
aws_secret_access_key=config['aws_secret_access_key'],
aws_session_token=config.get('aws_session_token', None))
session = Session()
