Python boto3.setup_default_session() Examples

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 run(self, command: str) -> None:
        """Execute one of the available commands.

        Args:
            command: Command in self.commands.
        """
        boto3.setup_default_session(region_name=self._config.aws_region)

        # Validate the configuration.
        try:
            if command not in {'clone_rules', 'compile_rules', 'configure', 'unit_test'}:
                self._config.validate()
            getattr(self, command)()  # Command validation already happened in the ArgumentParser.
        except InvalidConfigError as error:
            sys.exit('ERROR: {}\nPlease run "python3 manage.py configure"'.format(error))
        except TestFailureError as error:
            sys.exit('TEST FAILED: {}'.format(error)) 

def get_client(
    client,
    profile_name,
    aws_access_key_id,
    aws_secret_access_key,
    region=None,
):
    """Shortcut for getting an initialized instance of the boto3 client."""

    boto3.setup_default_session(
        profile_name=profile_name,
        aws_access_key_id=aws_access_key_id,
        aws_secret_access_key=aws_secret_access_key,
        region_name=region,
    )
    return boto3.client(client) 

def set_up_boto3():
    """
    A DRY place to make sure AWS credentials in settings override
    environment based credentials.  Boto3 will fall back to:
    http://boto3.readthedocs.io/en/latest/guide/configuration.html

    Taken from https://github.com/datadesk/django-bakery/blob/
    a2f1f74b03951450d797ec70cc9872d6c694e1e3/bakery/management/commands/__init__.py#L8
    """
    session_kwargs = {}
    if hasattr(settings, "AWS_ACCESS_KEY_ID"):
        session_kwargs["aws_access_key_id"] = settings.AWS_ACCESS_KEY_ID

    if hasattr(settings, "AWS_SECRET_ACCESS_KEY"):
        session_kwargs["aws_secret_access_key"] = settings.AWS_SECRET_ACCESS_KEY
    boto3.setup_default_session(**session_kwargs) 

def test_get_queues(self):
        boto3.setup_default_session()
        sqs_inst = sqs.SQSTalk()

        sqs_inst._get_or_create_queue('LowPriorityTest3')
        sqs_inst._get_or_create_queue('HighPriorityTest4')

        queues = sqs_inst.get_all_queues()
        print(queues)
        self.assertEqual(len(queues), 2)

        queues = sqs_inst.get_all_queues('Low')
        self.assertEqual(len(queues), 1)

        sqs_inst.get_all_queues('High')
        self.assertEqual(len(queues), 1) 

def test_create_queue(self):
        boto3.setup_default_session()

        sqs_inst = sqs.SQSTalk()

        sqs_inst._get_or_create_queue('LowPriorityTest1')
        sqs_inst._get_or_create_queue('HighPriorityTest2')

        expected_low_queue = sqs_inst._sqs.Queue('https://queue.amazonaws.com/123456789012/'
                                                 'LowPriorityTest1')
        expected_hi_queue = sqs_inst._sqs.Queue('https://queue.amazonaws.com/123456789012/'
                                                'HighPriorityTest2')

        self.assertEqual(len(sqs_inst._queues), 2)

        self.assertEqual(expected_low_queue, sqs_inst._queues['LowPriorityTest1'])
        self.assertEqual(expected_hi_queue, sqs_inst._queues['HighPriorityTest2']) 

def __init__(self, type, service=None, region='us-west-2', profile='default'):
        self.region = region
        self.connection_type = type
        self.service = service
        self.client = None
        self.resource = None
        self.profile = profile
        try:
            boto3.setup_default_session(profile_name=self.profile)
        except Exception as e:
            logger.info("Problem setting default boto3 session: {}".format(e)) 

def clean_boto_configuration(monkeypatch):
    monkeypatch.setenv('AWS_ACCESS_KEY_ID', 'AK000000000000000000')
    monkeypatch.setenv('AWS_SECRET_ACCESS_KEY',
                       '0000000000000000000000000000000000000000')
    monkeypatch.setenv('AWS_DEFAULT_REGION', 'us-east-1')
    # Reset previously created default session (for credentials)
    boto3.setup_default_session() 

def aws_upload(args: argparse.Namespace, zip_data: BinaryIO):
    if args.profile:
        session = boto3.session.Session()
        log.debug("Available profiles: %s", session.available_profiles)
        log.info("Using profile: %s", args.profile)
        boto3.setup_default_session(profile_name=args.profile)
    role_arn = set_up_role()
    if lambda_exists():
        upload(zip_data)
    else:
        arn = create_lambda(role_arn, zip_data, args.skill)
        print("Lambda ARN:", arn) 

def __init__(self, profile):
        if profile:
            boto3.setup_default_session(profile_name=profile)
        self.ssm = boto3.client('ssm') 

def test_find_instance_ids(mock_resource):
    """
    Test the find_instance_ids function
    """
    instance_id = 'abcdef-12345'
    instance = [MagicMock(id=instance_id)]
    boto3.setup_default_session(region_name='us-east-1')
    mock_resource.return_value = instance
    assert find_instance_ids('blah') == [instance_id] 

def client_with_default_region(service, default_region = DEFAULT_REGION):
    """
    Return a boto3 client for the named *service* in the *default_region* if no
    region is specified in the default session (via the environment, AWS
    config, or ``boto3.setup_default_session``).
    """
    try:
        return boto3.client(service)
    except NoRegionError:
        return boto3.client(service, region_name = default_region) 

def main():
    """Main function
    """
    args = parse_args()

    if args.profile:
        boto3.setup_default_session(profile_name=args.profile)

    name_filter = "*"
    for a in args.keywords:
        name_filter += a + "*"

    hosts = get_hosts(name_filter, args.region)

    if not hosts:
        sys.exit("No hosts found matching those keywords.")

    host_name_key = 'PrivateIpAddress'
    if args.public_dns:
        host_name_key = 'PublicDnsName'

    if len(hosts) == 1:
        print("Logging in to %s..." % hosts[0][1])
        ssh(hosts[0][0][host_name_key], args.user, args.identity)
    else:
        choice = None

        for i, host in enumerate(hosts, 1):
            print("{0}) {1} - {2}".format(i, host[1], host[0][host_name_key]))
        print("Enter) exit")

        try:
            choice = int(input("Choose wisely: "))
        except (SyntaxError, NameError, ValueError):
            sys.exit("You chose... poorly")

        if choice == 0 or choice > len(hosts):
            sys.exit("You chose... poorly")
        else:
            print("Logging in to %s..." % hosts[choice - 1][1])
            ssh(hosts[choice - 1][0][host_name_key], args.user, args.identity) 

def mock_public_bucket():
    """A pytest fixture that creates a temp working dir, mocks the Configuration class to resolve the mocked config"""
    os.environ['AWS_ACCESS_KEY_ID'] = "fake-access-key"
    os.environ['AWS_SECRET_ACCESS_KEY'] = "fake-secret-key"

    with mock_s3():
        boto3.setup_default_session()
        bucket = "gigantum-test-bucket"
        conn = boto3.resource('s3', region_name='us-east-1')
        conn.create_bucket(Bucket=bucket)
        conn.meta.client.put_bucket_acl(ACL='public-read', Bucket=bucket)

        with tempfile.NamedTemporaryFile('wt') as tf:
            tf.write("asdf" * 1000)
            tf.seek(0)
            conn.meta.client.upload_file(tf.name, bucket, 'test-file-1.bin')
            conn.meta.client.upload_file(tf.name, bucket, 'test-file-2.bin')

        with tempfile.NamedTemporaryFile('wt') as tf:
            tf.write("1234" * 100)
            tf.seek(0)
            conn.meta.client.upload_file(tf.name, bucket, 'metadata/test-file-3.bin')
            conn.meta.client.upload_file(tf.name, bucket, 'metadata/test-file-4.bin')

        with tempfile.NamedTemporaryFile('wt') as tf:
            tf.write("5656" * 100)
            tf.seek(0)
            conn.meta.client.upload_file(tf.name, bucket, 'metadata/sub/test-file-5.bin')

        yield bucket 

def setup_aws_client(cls, config):
        aws_access_key_id = config['aws_access_key_id']
        aws_secret_access_key = config['aws_secret_access_key']

        LOGGER.info('Attempting to create AWS session')
        boto3.setup_default_session(aws_access_key_id=aws_access_key_id,
                                    aws_secret_access_key=aws_secret_access_key) 

def __init__(
        self, bucket, secret=None, region=None, bare=False, private=False, profile=None
    ):
        if profile:
            boto3.setup_default_session(profile_name=profile)
        self.s3 = boto3.resource("s3", region_name=region)
        self.bucket = bucket
        self.secret = secret
        self.index = "" if bare else "index.html"
        self.acl = "private" if private else "public-read" 

def main():

    args = parse_args()

    try:
        boto3.setup_default_session(profile_name=args.profile, region_name=args.region)
        conn = boto3.client('ec2')
    except ProfileNotFound as e:
        print(e)
        sys.exit(1)

    # Retrieve a list of instances that match the filters
    reservations = conn.describe_instances(Filters=get_filters(args))
    if len(reservations['Reservations']) == 0:
        print('No instances matching criteria')
        sys.exit(1)

    instance_dns_names = [
            instance['PublicDnsName']
            for reservation
            in reservations['Reservations']
            for instance
            in reservation['Instances']]
    if args.all_matching_instances:
        pass
    else:
        # Pick a random instance from the results
        instance_dns_names = [instance_dns_names[random.randrange(0, len(instance_dns_names))]]

    if args.command:
        remote_command = ' '.join(args.command)
    else:
        remote_command = ''

    for dns_name in instance_dns_names:
        if args.ssh_user:
            dns_name = '%s@%s' % (args.ssh_user, dns_name)

        ssh_cmd = 'ssh %s %s %s' % (args.ssh_args, dns_name, remote_command)
        os.system(ssh_cmd) 

def __init__(self):
        self.default_wait_secs = 5
        self.verbose = bool(int(os.getenv('VERBOSE', 0)))
        self.profile = os.getenv('AWS_PROFILE')

        try:
            self.region = os.getenv('AWS_DEFAULT_REGION')
            assert self.region
        except:
            try:
                self.region = os.getenv('AWS_REGION')
                assert self.region
            except:
                self.region = 'us-east-1'

        self.session = boto3.session.Session()
        boto3.setup_default_session()
        if self.profile:
            print(
                'profile={} region={}'.format(
                    self.profile,
                    self.region
                ),
                file=sys.stderr
            )
            self.session = boto3.session.Session(profile_name=self.profile)
            boto3.setup_default_session(profile_name=self.profile) 

def __init__(self):
        self.default_wait_secs = 5
        self.verbose = bool(int(os.getenv('VERBOSE', 0)))
        self.profile = os.getenv('AWS_PROFILE')

        try:
            self.region = os.getenv('AWS_DEFAULT_REGION')
            assert self.region
        except:
            try:
                self.region = os.getenv('AWS_REGION')
                assert self.region
            except:
                self.region = 'us-east-1'

        self.session = boto3.session.Session()
        boto3.setup_default_session()
        if self.profile:
            print(
                'profile={} region={}'.format(
                    self.profile,
                    self.region
                ),
                file=sys.stderr
            )
            self.session = boto3.session.Session(profile_name=self.profile)
            boto3.setup_default_session(profile_name=self.profile) 

def __init__(self):
        self.default_wait_secs = 5
        self.verbose = bool(int(os.getenv('VERBOSE', 0)))
        self.profile = os.getenv('AWS_PROFILE')

        try:
            self.region = os.getenv('AWS_DEFAULT_REGION')
            assert self.region
        except:
            try:
                self.region = os.getenv('AWS_REGION')
                assert self.region
            except:
                self.region = 'us-east-1'

        self.session = boto3.session.Session()
        boto3.setup_default_session()
        if self.profile:
            print(
                'profile={} region={}'.format(
                    self.profile,
                    self.region
                ),
                file=sys.stderr
            )
            self.session = boto3.session.Session(profile_name=self.profile)
            boto3.setup_default_session(profile_name=self.profile) 

def __init__(self):
        self.default_wait_secs = 5
        self.verbose = bool(int(os.getenv('VERBOSE', 0)))
        self.profile = os.getenv('AWS_PROFILE')

        try:
            self.region = os.getenv('AWS_DEFAULT_REGION')
            assert self.region
        except:
            try:
                self.region = os.getenv('AWS_REGION')
                assert self.region
            except:
                self.region = 'us-east-1'

        self.session = boto3.session.Session()
        boto3.setup_default_session()
        if self.profile:
            print(
                'profile={} region={}'.format(
                    self.profile,
                    self.region
                ),
                file=sys.stderr
            )
            self.session = boto3.session.Session(profile_name=self.profile)
            boto3.setup_default_session(profile_name=self.profile) 

def __init__(self):
        self.default_wait_secs = 5
        self.verbose = bool(int(os.getenv('VERBOSE', 0)))
        self.profile = os.getenv('AWS_PROFILE')

        try:
            self.region = os.getenv('AWS_DEFAULT_REGION')
            assert self.region
        except:
            try:
                self.region = os.getenv('AWS_REGION')
                assert self.region
            except:
                self.region = 'us-east-1'

        self.session = boto3.session.Session()
        boto3.setup_default_session()
        if self.profile:
            print(
                'profile={} region={}'.format(
                    self.profile,
                    self.region
                ),
                file=sys.stderr
            )
            self.session = boto3.session.Session(profile_name=self.profile)
            boto3.setup_default_session(profile_name=self.profile) 

def test_default_session():
    boto3.setup_default_session(region_name="us-east-1")
    assert wr._utils.ensure_session().region_name == "us-east-1"
    boto3.setup_default_session(region_name="us-east-2")
    assert wr._utils.ensure_session().region_name == "us-east-2"
    boto3.setup_default_session(region_name="us-west-1")
    assert wr._utils.ensure_session().region_name == "us-west-1"
    boto3.setup_default_session(region_name="us-west-2")
    assert wr._utils.ensure_session().region_name == "us-west-2" 

def test_iterateCore_exists(mocker):
  # mock the get regions part
  mockreturn = lambda service: ['us-east-1']
  mockee = 'boto3.session.Session.get_available_regions'
  mocker.patch(mockee, side_effect=mockreturn)

  # undo some region settings from before
  import boto3
  boto3.setup_default_session(region_name='us-east-1')

  # create mock redshift
  import boto3
  redshift_client = boto3.client('redshift')
  redshift_client.create_cluster(
    ClusterIdentifier='abc',
    NodeType='abc',
    MasterUsername='abc',
    MasterUserPassword='abc'
  )

  # test
  rpi = RedshiftPerformanceIterator()
  rpi.region_include = ['us-east-1']
  x = list(rpi.iterate_core())
  assert len(x) == 1


# cannot name function "test_iterator" because the filename is as such
# pytest .../test_iterator.py -k 'test_iterator' would run all tests, not just this one 

def test_iterateCore_exists(mocker):
  # mock the get regions part
  mockreturn = lambda service: ['us-east-1']
  mockee = 'boto3.session.Session.get_available_regions'
  mocker.patch(mockee, side_effect=mockreturn)

  # undo some region settings from before
  import boto3
  boto3.setup_default_session(region_name='us-east-1')

  # create mock redshift
  import boto3
  ec2_client = boto3.resource('ec2')
  ec2_client.create_instances(
    MinCount = 1,
    MaxCount = 1,
    InstanceType='t2.medium'
  )

  # test
  rpi = Ec2Iterator()
  rpi.region_include=['us-east-1']
  x = list(rpi.iterate_core())
  assert len(x) == 1


# cannot name function "test_iterator" because the filename is as such
# pytest .../test_iterator.py -k 'test_iterator' would run all tests, not just this one 

def get(self):
        # get cloudtrail ec2 type changes for all instances
        logger.debug("Downloading cloudtrail data (from %i regions)"%len(self.region_include))
        df_2 = []
        import boto3

        # add some spaces for aligning the progress bars
        desc="Cloudtrail events in all regions"
        desc = "%-50s"%desc

        iter_wrap = self.region_include
        iter_wrap = self.tqdmman(iter_wrap, desc=desc, total=len(self.region_include))
        for region_name in iter_wrap:
          boto3.setup_default_session(region_name = region_name)
          df_1 = super().get()
          df_1['Region'] = region_name # bugfix, field name was "region" (lower-case)
          df_2.append(df_1.reset_index())

        # concatenate
        df_3 = pd.concat(df_2, axis=0, sort=False)

        # check if empty
        if df_3.shape[0]==0:
          return df_3

        # sort again
        df_3 = df_3.set_index(["Region", "ServiceName", "ResourceName", "EventTime"]).sort_index()

        return df_3 

def set_resource(self, region_name):
        """
        Easy-to-mock function since moto mock of cloudwatch is giving pagination error
        """
        boto3.setup_default_session(region_name = region_name)
        self.cloudwatch_resource = boto3.resource('cloudwatch') 

def __iter__(self):
    # over-ride the __iter__ to get the ec2 resource object for the current code (backwards compatibility)

    # method 1 for ec2
    # ec2_it = self.ec2_resource.instances.all()
    # return ec2_it

    # boto3 ec2 and cloudwatch data
    ec2_resource_all = {}
    import boto3

    # TODO cannot use directly use the iterator exposed in "ec2_it"
    # because it would return the dataframes from Cloudwatch,
    # whereas in the cloudwatch data fetch here, the data gets cached to redis.
    # Once the redshift.iterator can cache to redis, then the cloudwatch part here
    # can also be dropped, as well as using the "ec2_it" iterator directly
    # for ec2_dict in self.ec2_it:
    for ec2_dict, ec2_id, ec2_launchtime, _ in super().__iter__():
      if ec2_dict['Region'] not in ec2_resource_all.keys():
        boto3.setup_default_session(region_name = ec2_dict['Region'])
        ec2_resource_all[ec2_dict['Region']] = boto3.resource('ec2')

      ec2_resource_single = ec2_resource_all[ec2_dict['Region']]
      ec2_l = ec2_resource_single.instances.filter(InstanceIds=[ec2_dict['InstanceId']])
      ec2_l = list(ec2_l)
      if len(ec2_l)==0:
        continue # not found

      # yield first entry
      ec2_obj = ec2_l[0]
      ec2_obj.region_name = ec2_dict['Region']

      yield ec2_dict, ec2_id, ec2_launchtime, ec2_obj 

def _get_rules_from_aws(self):
        """
        Load the EC2 security rules off AWS into a list of dict.

        Returns:
            list
        """
        list_of_rules = list()

        if self.profile:
            boto3.setup_default_session(profile_name=self.profile)

        if self.region:
            ec2 = boto3.client('ec2', region_name=self.region)
        else:
            ec2 = boto3.client('ec2')

        security_groups = ec2.describe_security_groups(Filters=self.filters)

        for group in security_groups['SecurityGroups']:
            group_dict = dict()
            group_dict['id'] = group['GroupId']
            group_dict['name'] = group['GroupName']
            group_dict['description'] = group.get('Description', None)

            if (group.get('IpPermissions', None) or
                    group.get('IpPermissionsEgress', None)):
                group_dict['rules'] = list()

            for rule in group.get('IpPermissions', None):
                rule_dict = self._build_rule(rule)
                rule_dict['direction'] = "INGRESS"

                group_dict['rules'].append(rule_dict)

            for rule in group.get('IpPermissionsEgress', None):
                rule_dict = self._build_rule(rule)
                rule_dict['direction'] = "EGRESS"

                group_dict['rules'].append(rule_dict)

            list_of_rules.append(group_dict)

        return list_of_rules