Last active
April 4, 2019 20:17
-
-
Save bkruger99/6bbaacf1e7fa49891d421d6a1a7ba9c9 to your computer and use it in GitHub Desktop.
aws connection helper for python. Allows you to use a single call to get either a client or resource and also has sts support built in. If you use this, I'd appreciate feedback.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/usr/bin/env python | |
import boto3 | |
import sys | |
""" | |
Generic AWS Helper call class to setup a boto3 Session. Now with assume role support. | |
Pass in 'type=' to do either 'client' or 'resource' | |
usage: | |
ec2 = aws(type='client', service_name='ec2') | |
sqs_resource = aws(type='resource', service_name='sqs', RoleArn='arn:aws:iam::012345678901:role/example-role', | |
RoleSessionName='SomeSessionName') | |
This will allow for either using your ~/.aws credentials or allow you to override in the function calls. | |
Python 2 and 3 compatible without six. | |
""" | |
def aws(type='client', **kwargs): | |
""" | |
This makes boto3 connection. Client is the default. | |
:param: type (str) - client type. Either "resource" or "client" right now | |
:param: **kwargs - anything else passed in. | |
:returns: Your aws object type you requested. | |
""" | |
myargs = {} | |
if 'service_name' not in kwargs: | |
print("You need to specify a service_name") | |
raise | |
myargs.update(**kwargs) | |
if 'RoleArn' in kwargs and 'RoleSessionName' in myargs: | |
stscreds = __role_arn_to_session(**myargs) | |
myargs.update(stscreds) | |
myargs = __stripargs(**myargs) | |
session = boto3.Session() | |
client = eval("session." + type)(**myargs) | |
return client | |
# sts assume role | |
# originally from: https://gist.github.com/gene1wood/938ff578fbe57cf894a105b4107702de | |
# slightly modified. | |
def __role_arn_to_session(**args): | |
""" | |
Pass in at least "RoleArn" and "RoleSessionName" with your args in the 'aws' function above. | |
""" | |
clientargs = __stripargs(**args) | |
stsargs = __stripargs(sts=True, **args) | |
clientargs['service_name'] = 'sts' | |
client = boto3.client(**clientargs) | |
response = client.assume_role(**stsargs) | |
return { | |
'aws_access_key_id': response['Credentials']['AccessKeyId'], | |
'aws_secret_access_key': response['Credentials']['SecretAccessKey'], | |
'aws_session_token': response['Credentials']['SessionToken']} | |
# Used to strip out STS arguments. | |
def __stripargs(sts=False, **args): | |
stsTuple = ('RoleArn', 'RoleSessionName', 'Policy', 'DurationSeconds', 'ExternalId', 'SerialNumber', 'TokenCode') | |
clientargs = dict(args) | |
stsargs = {} | |
# Check if python 3 or newer. If not, then it's probably 2. | |
if sys.version_info.major >= 3: | |
for k,v in args.items(): | |
if k in stsTuple: | |
stsargs[k] = v | |
del clientargs[k] | |
else: | |
for k, v in args.iteritems(): | |
if k in stsTuple: | |
stsargs[k] = v | |
del clientargs[k] | |
if sts is not True: | |
return clientargs | |
else: | |
return stsargs |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
I have updated this to be python2&3 compatible. A little repetition, but doesn't require an additional library installation.