Skip to content

Instantly share code, notes, and snippets.

@nutzhub
Last active January 16, 2019 08:08
Show Gist options
  • Save nutzhub/09dd67d8c378fa5db2fc084f415abaa7 to your computer and use it in GitHub Desktop.
Save nutzhub/09dd67d8c378fa5db2fc084f415abaa7 to your computer and use it in GitHub Desktop.
snippet to use sts (temporary credential) to call boto3 api in your machine
#!/usr/bin/env python
import gzip
import io
import boto3
import json
from botocore.exceptions import ClientError
class STS(object):
"""
Sts: Object to manage the persistence of authentication over multiple
runs of an automation script. When testing a script this will
save having to input an MFA token multiple times when using
an account that requires it.
"""
def __init__(self, role_arn, temporary_credentials_path, mfa_arn):
self.temp_creds_path = temporary_credentials_path
self.mfa_arn = mfa_arn
self.role_arn = role_arn
def get_temporary_session(self):
"""
get_temporary_session: checks the temporary credentials stored
on disk, if they fail to authenticate re-attempt to assume
the role. The credentials requested last 15 minutes. For
debugging purposes these can be persisted for up to an hour.
"""
try:
with open(self.temp_creds_path, 'r') as tmp_creds:
credentials = json.loads(tmp_creds.read())
client = boto3.client(
'sts',
aws_access_key_id=credentials['AccessKeyId'],
aws_secret_access_key=credentials['SecretAccessKey'],
aws_session_token=credentials['SessionToken']
)
_ = client.get_caller_identity()['Account']
except (IOError, ClientError, FileNotFoundError):
response = boto3.client('sts').assume_role(
DurationSeconds=900,
RoleArn=self.role_arn,
RoleSessionName='sts-assume-role-script',
SerialNumber=self.mfa_arn,
TokenCode=input('MFA_Token:')
)
credentials = response['Credentials']
with open(self.temp_creds_path, 'w+') as tmp_creds:
tmp_creds.write(json.dumps({
'AccessKeyId': credentials['AccessKeyId'],
'SecretAccessKey': credentials['SecretAccessKey'],
'SessionToken': credentials['SessionToken']}))
return boto3.Session(
aws_access_key_id=credentials['AccessKeyId'],
aws_secret_access_key=credentials['SecretAccessKey'],
aws_session_token=credentials['SessionToken'],
)
# -----------
BUCKET_NAME = "cdl.t1c.dev"
def main():
sts = STS("arn:aws:iam::246898065194:role/r_data-engineer", "/tmp/aws_sts_creds.json",
"arn:aws:iam::627443353872:mfa/[email protected]")
boto3.setup_default_session(profile_name='cg-aws-operation')
session = sts.get_temporary_session()
s3 = session.resource("s3")
for bucket in s3.buckets.all():
print(bucket.name)
if __name__ == "__main__":
main()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment