Add ensure_s3_bucket

Author: Octogonapus

axon/client.py

@@ -324,6 +324,65 @@ def ensure_ec2_instance_profile(region, profile_name="axon-ec2-autogenerated-ins
     return instance_profile.arn
 
 
+def ensure_s3_bucket(region):
+    """
+    Ensures that a matching S3 bucket exists.
+
+    :param region: The region, or `None` to pull the region from the environment.
+    :return: The name of the bucket.
+    """
+    client = make_client("s3", region)
+    prefix = "axon-autogenerated-"  # Used to identify the bucket that Axon manages
+
+    def get_axon_bucket():
+        buckets = client.list_buckets()["Buckets"]
+
+        # Return the first matching bucket name, if there is one
+        for bucket in buckets:
+            if bucket["Name"].startswith(prefix):
+                return bucket["Name"]
+
+        return None
+
+    axon_bucket = get_axon_bucket()
+    if axon_bucket is not None:
+        return axon_bucket
+
+    # There is no matching bucket name, so create a new one
+    import random
+    import string
+    while True:
+        bucket_name = prefix + ''.join(
+            random.choice(string.ascii_lowercase + string.digits) for _ in range(30))
+
+        # If the user wants a region in this list, we need to set it
+        if region in ['EU', 'eu-west-1', 'us-west-1', 'us-west-2',
+                      'ap-south-1', 'ap-southeast-1', 'ap-southeast-2',
+                      'ap-northeast-1', 'sa-east-1', 'cn-north-1',
+                      'eu-central-1']:
+            client.create_bucket(ACL='private', Bucket=bucket_name,
+                                 CreateBucketConfiguration={'LocationConstraint': region})
+        else:
+            # Otherwise the region will be us-east-1
+            client.create_bucket(ACL='private', Bucket=bucket_name)
+
+        # Busy loop until the bucket is created. Otherwise, we will set the public access block
+        # too early and its configuration will be lost
+        while True:
+            axon_bucket = get_axon_bucket()
+            if axon_bucket is not None:
+                break
+
+        client.put_public_access_block(Bucket=bucket_name, PublicAccessBlockConfiguration={
+            'BlockPublicAcls': True,
+            'IgnorePublicAcls': True,
+            'BlockPublicPolicy': True,
+            'RestrictPublicBuckets': True
+        })
+
+        return bucket_name
+
+
 def ensure_cluster(ecs_client, cluster_name):
     """
     Ensures that a matching cluster exists. If there are no matching clusters, one is created.
@@ -420,6 +479,7 @@ def impl_ensure_configuration(cluster_name, task_family, region):
     :param region: The region, or `None` to pull the region from the environment.
     """
     client = make_client("ecs", region)
+    ensure_s3_bucket(region)
     ensure_cluster(client, cluster_name)
     ensure_task(client, task_family, region, 2048, 4096)
     ensure_ec2_security_group(region)