Separate ensure_ec2_instance_profile from ensure_ec2_role

Author: Octogonapus

axon/client.py

@@ -31,6 +31,10 @@ def make_client(name, region):
         return boto3.client(name, region_name=region)
 
 
+def make_resource(name):
+    return boto3.resource(name)
+
+
 def ensure_log_group(group_name, region):
     """
     Ensures that a log group is present. If there is a matching log group, nothing is created.
@@ -208,11 +212,12 @@ def ensure_role(client, role_name):
 def ensure_task_role(region):
     """
     Ensures a task role exists. If there is one matching role, its Arn is returned. If there are
-    multiple matching roles, a RuntimeError is raised. If there are no matching roles, a new one
+    multiple matching roles, a `RuntimeError` is raised. If there are no matching roles, a new one
     is created.
 
-    TODO: Fix this.
+    TODO: Fix this:
     This method does not check that a matching role has the correct policies.
+
     :param region: The region, or `None` to pull the region from the environment.
     :return: The role Arn.
     """
@@ -253,9 +258,14 @@ def ensure_task_role(region):
     return role_arn
 
 
-def ensure_ec2_role(region):
-    role_name = "axon-ec2-autogenerated-role"
-    profile_name = "axon-ec2-autogenerated-instance-profile"
+def ensure_ec2_role(region, role_name="axon-ec2-autogenerated-role"):
+    """
+    Ensures the EC2 role exists. Creates the role if it does not exist.
+
+    :param region: The region, or `None` to pull the region from the environment.
+    :param role_name: The name of the role to ensure.
+    :return: The role Arn.
+    """
     client = make_client("iam", region)
     role_arn = ensure_role(client, role_name)
     if role_arn is None:
@@ -282,15 +292,36 @@ def ensure_ec2_role(region):
         client.attach_role_policy(RoleName=role_name,
                                   PolicyArn="arn:aws:iam::aws:policy/AmazonS3FullAccess")
 
+    return role_arn
+
+
+def ensure_ec2_instance_profile(region, profile_name="axon-ec2-autogenerated-instance-profile",
+                                role_name="axon-ec2-autogenerated-role"):
+    """
+    Ensures the EC2 instance profile exists and has the EC2 role attached.
+
+    :param region: The region, or `None` to pull the region from the environment.
+    :param profile_name: The name of the instance profile to ensure.
+    :param role_name: The name of the role to ensure.
+    :return: The instance profile Arn.
+    """
+    client = make_client("iam", region)
+    iam_resource = make_resource('iam')
+
+    # Get or create the instance profile
     try:
-        client.get_instance_profile(InstanceProfileName=profile_name)
+        local_profile = client.get_instance_profile(InstanceProfileName=profile_name)
     except:
-        client.create_instance_profile(InstanceProfileName=profile_name)
-        iam_resource = boto3.resource('iam')
-        instance_profile = iam_resource.InstanceProfile(profile_name)
+        local_profile = client.create_instance_profile(InstanceProfileName=profile_name)
+
+    instance_profile = iam_resource.InstanceProfile(
+        local_profile['InstanceProfile']['InstanceProfileName'])
+
+    if role_name not in [role.name for role in instance_profile.roles]:
+        # Add the role if it does not exist
         instance_profile.add_role(RoleName=role_name)
 
-    return role_arn
+    return instance_profile.arn
 
 
 def ensure_cluster(ecs_client, cluster_name):
@@ -395,6 +426,7 @@ def impl_ensure_configuration(cluster_name, task_family, region):
     ensure_ecs_security_group(region)
     ensure_task_role(region)
     ensure_ec2_role(region)
+    ensure_ec2_instance_profile(region)
 
 
 def impl_start_task(cluster_name, task_family, revision, region):