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
