#!/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/__YOUR_USER__@central.tech")
    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()