def download_file(self, bucket: str, key: str, file_obj: IO, wait: bool = False, **kwargs): """ Convenience function for downloading an object out of S3. :param bucket: Name of the bucket. :param key: Name of the key. :param file_obj: File-like object to download the file to. :param wait: If true, will wait up to 100 seconds for the object to exist before reading it. :param kwargs: Any additional arguments to pass to the underlying boto call. """ if wait: logger.info( "Waiting for s3://%s/%s to exist with additional params: %s", bucket, key, kwargs) waiter = self.s3.get_waiter('object_exists') waiter.wait(Bucket=bucket, Key=key, **kwargs) throttled_call( self.s3.download_fileobj, Bucket=bucket, Key=key, Fileobj=file_obj, ExtraArgs=kwargs, )
def test_throttled_call_clienterror_noerr(self, mock_sleep): """Test throttle_call with no error""" mock_func = MagicMock() error_response = {"Error": {"Code": "Throttling"}} client_error = ClientError(error_response, "test") mock_func.side_effect = [client_error, client_error, True] throttled_call(mock_func) self.assertEqual(3, mock_func.call_count)
def _write_to_dynamodb( self, events: List[Dict[str, str]], ): """ Writes the list of records to DynamoDB. :param events: List of records. """ for event in events: throttled_call( self.dynamodb_table.put_item, Item=event, )
def put_object_tags(self, bucket: str, key: str, tags: Dict, **kwargs): """ Convenience function for adding tags to an object. :param bucket: Name of the bucket. :param key: Name of the key. :param tags: A dictionary of tags to set. :param kwargs: Any additional arguments to pass to the underlying boto call. """ throttled_call(self.s3.put_object_tagging, Bucket=bucket, Key=key, Tagging={"TagSet": dict_to_boto3_tags(tags)}, **kwargs)
def write_file(self, bucket: str, key: str, body: str, **kwargs): """ Convenience function for writing an object to S3. :param bucket: Name of the bucket. :param key: Name of the key. :param body: The contents of the object. :param kwargs: Any additional arguments to pass to the underlying boto call. """ throttled_call(self.s3.put_object, Bucket=bucket, Key=key, Body=body, **kwargs)
def put_evaluations(self, result_token: str, evaluations: List[Dict[str, str]]) -> None: """ Convenience function for submitting evaluations to the AWS Config service in chunks of 100. :param result_token: The result token that yielded these evaluations. :param evaluations: A list of dictionaries representing the evaluations. """ for evaluation_group in chunker(evaluations, 100): logger.info("Submitting evaluations for: %s", [(evaluation["ComplianceResourceId"], evaluation["ComplianceType"]) for evaluation in evaluation_group]) throttled_call(self.config_client.put_evaluations, ResultToken=result_token, Evaluations=evaluation_group)
def upsert_item( self, table_name: str, primary_key_name: str, primary_key_value: str, attribute_name: str, attribute_value: str ) -> Dict[str, Any]: """ Insert/update items in a DynamoDB table for a specific primary key value. The value inserted/updated is of type string. :param table_name: DynamoDB table name :param primary_key_name: Table's primary key name :param primary_key_value: Table's primary key value :param attribute_name: Item attribute to be inserted/updated :param attribute_value: Attribute value to be inserted/updated :return: response """ key = {primary_key_name: {'S': primary_key_value}} expression_attribute_values = {':d': {'S': attribute_value}} update_expression = f'SET {attribute_name} = :d' return throttled_call( self.client.update_item, TableName=table_name, Key=key, ExpressionAttributeValues=expression_attribute_values, UpdateExpression=update_expression )
def read_file(self, bucket: str, key: str, wait: bool = False, **kwargs) -> str: """ Convenience function for reading an object out of S3. :param bucket: Name of the bucket. :param key: Name of the key. :param wait: If true, will wait up to 100 seconds for the object to exist before reading it. :param kwargs: Any additional arguments to pass to the underlying boto call. :return: The contents of the object, read and decoded as utf-8. """ if wait: logger.info( "Waiting for s3://%s/%s to exist with additional params: %s", bucket, key, kwargs) waiter = self.s3.get_waiter('object_exists') waiter.wait(Bucket=bucket, Key=key, **kwargs) result = throttled_call(self.s3.get_object, Bucket=bucket, Key=key, **kwargs)["Body"].read().decode("utf-8") return result
def copy_file(self, source_bucket: str, destination_bucket: str, source_key: str, destination_key: str, **kwargs): """ Convenience function for copying an S3 object from one bucket to another. :param source_bucket: Name of the source bucket. :param destination_bucket: Name of the destination bucket. :param source_key: Name of the source object. :param destination_key: Name of the destination object. :param kwargs: Any additional arguments to pass to the underlying boto call. """ throttled_call(self.s3.copy_object, Bucket=destination_bucket, Key=destination_key, CopySource={ "Bucket": source_bucket, "Key": source_key }, **kwargs)
def put_bucket_tags(self, bucket: str, tags: Dict, merge: bool = False): """ Convenience function for tagging a bucket in S3. :param bucket: Name of the bucket. :param tags: Dict of tags. :param merge: True if new tags should be merged with the existing tags on the bucket, otherwise the existing tags will be overwritten entirely. """ if merge: existing_tags = self.get_bucket_tags(bucket=bucket) existing_tags.update(tags) tags = existing_tags throttled_call( self.s3.put_bucket_tagging, Bucket=bucket, Tagging={'TagSet': dict_to_boto3_tags(tags)}, )
def get_object_tags(self, bucket: str, key: str, **kwargs): """ Convenience function for getting tags on an object. :param bucket: Name of the bucket. :param key: Name of the key. :param kwargs: Any additional arguments to pass to the underlying boto call. :return: A dictionary representing the tags on the object. """ boto_tags = throttled_call(self.s3.get_object_tagging, Bucket=bucket, Key=key, **kwargs) return boto3_tags_to_dict(boto_tags["TagSet"])
def get_bucket_tags(self, bucket: str) -> Dict[str, str]: """ Convenience function for getting tags of a bucket in S3. :param bucket: Name of the bucket :return: Dictionary representing the tags of the requested S3 bucket. """ try: result = throttled_call(self.s3.get_bucket_tagging, Bucket=bucket) return boto3_tags_to_dict(result["TagSet"]) except ClientError: logger.warning('Bucket %s has no tags', bucket) return {}
def delete_item( self, table_name: str, primary_key_name: str, primary_key_value: str, ): """ Convenience function for deleting an item from a DynamoDB table with a specific primary key value. :param table_name: DynamoDB table name :param primary_key_name: Table's primary key name :param primary_key_value: Table's primary key value :return: response """ key = {primary_key_name: {'S': primary_key_value}} return throttled_call( self.client.delete_item, TableName=table_name, Key=key, )
def find_amis(self, source_amis: Set[str] = None, newer_than: datetime = None) -> Set[str]: """ Function for finding AMIs that are children of a given set of AMIs. :param source_amis: A set of parent ami-ids to find children of. :param newer_than: Only AMIs newer than this datetime object will be returned. :return: A set of AMI ids. """ filters = {} if source_amis: filters["tag:source_ami"] = list(source_amis) images = throttled_call( self.ec2_client.describe_images, Filters=create_filters(filters), Owners=["self"] )["Images"] image_ids = { image["ImageId"] for image in images if not newer_than or datetime.strptime(image["CreationDate"], DATETIME_FORMAT) > newer_than } logger.info("Found additional images=%s", image_ids) return image_ids
def hash_file(self, bucket: str, key: str, **kwargs) -> str: """ Convenience function for getting the hash of an object from S3 :param bucket: Name of the bucket. :param key: Name of the key. :param kwargs: Any additional arguments to pass to the underlying boto call. :return: SHA256 hash of the object. """ stream = throttled_call(self.s3.get_object, Bucket=bucket, Key=key, **kwargs)["Body"] # 5 megabytes block_size = 5242880 hasher = hashlib.sha256() buffer = stream.read(block_size) while buffer: hasher.update(buffer) buffer = stream.read(block_size) return hasher.hexdigest()
def _find_metadata( self, grouped_events: Dict[str, Dict[str, List[Dict[str, str]]]], ) -> List[Dict[str, str]]: """ Look up metadata on each event. :param grouped_events: List of CloudWatch events grouped by account and region. :return: A list of records with metadata formatted for upload to DynamoDB. """ events_with_metadata = [] for account, events_by_region in grouped_events.items(): for region, events in events_by_region.items(): ec2_client: botostubs.EC2 = self.sts_client.get_boto3_client_for_account( account_id=account, role_name=os.environ["FLAPPY_DETECTOR_ROLE"], client_name="ec2", region_name=region, ) response = throttled_call( ec2_client.describe_instances, InstanceIds=[event["instance_id"] for event in events], ) instance_metadata = { instance["InstanceId"]: boto3_tags_to_dict(instance["Tags"]) for reservation in response["Reservations"] for instance in reservation["Instances"] } for event in events: cur_metadata = instance_metadata[event["instance_id"]] group_name = cur_metadata.get( "spotinst:aws:ec2:group:id", cur_metadata.get("aws:autoscaling:groupName"), ) if not group_name: logger.warning( "Event for instance_id:%s has no associated group, ignoring", event["instance_id"], extra={ "instance_id": event["instance_id"], "event": event, "metadata": cur_metadata, }, ) continue standardized_tags = { "application": cur_metadata.get("application"), "environment": cur_metadata.get("environment") or cur_metadata.get("env"), "team": cur_metadata.get("team"), "group_name": group_name, "region": region, "account": account, } events_with_metadata.append( dict( **standardized_tags, **event, )) return events_with_metadata