Iam utils (#53)

setup.py

@@ -120,7 +120,7 @@ def default_setup_args(*, version):
 
 extras_require = dict()
 
-test_requirements = ["tox", "pytest", "pytest-cov"]
+test_requirements = ["tox", "pytest", "pytest-cov", "moto[all]"]
 
 test_requirements = list(set(test_requirements))
 extras_require["tests"] = test_requirements

src/autogluon/cloud/utils/iam.py

@@ -1,3 +1,9 @@
+import json
+from typing import Any, Dict
+
+import boto3
+from botocore.exceptions import ClientError
+
 from .constants import POLICY_ACCOUNT_PLACE_HOLDER, POLICY_BUCKET_PLACE_HOLDER, TRUST_RELATIONSHIP_ACCOUNT_PLACE_HOLDER
 
 
@@ -27,3 +33,43 @@ def replace_iam_policy_place_holder(policy_document, account_id=None, bucket=Non
                     resource.replace(POLICY_BUCKET_PLACE_HOLDER, bucket) for resource in statement["Resource"]
                 ]
     return policy_document
+
+
+def create_iam_role(role_name: str, trust_relationship: Dict[str, Any]) -> str:
+    iam_client = boto3.client("iam")
+    try:
+        response = iam_client.create_role(RoleName=role_name, AssumeRolePolicyDocument=json.dumps(trust_relationship))
+        return response["Role"]["Arn"]
+    except ClientError as error:
+        if error.response["Error"]["Code"] != "EntityAlreadyExists":
+            raise error
+
+
+def create_iam_policy(policy_name: str, policy: Dict[str, Any]) -> str:
+    iam_client = boto3.client("iam")
+    try:
+        response = iam_client.create_policy(PolicyName=policy_name, PolicyDocument=json.dumps(policy))
+        return response["Policy"]["Arn"]
+    except ClientError as error:
+        if error.response["Error"]["Code"] != "EntityAlreadyExists":
+            raise error
+
+
+def attach_iam_policy(role_name: str, policy_arn: str):
+    iam_client = boto3.client("iam")
+    iam_client.attach_role_policy(RoleName=role_name, PolicyArn=policy_arn)
+
+
+def create_instance_profile(instance_profile_name: str) -> str:
+    iam_client = boto3.client("iam")
+    try:
+        response = iam_client.create_instance_profile(InstanceProfileName=instance_profile_name)
+        return response["InstanceProfile"]["Arn"]
+    except ClientError as error:
+        if error.response["Error"]["Code"] != "EntityAlreadyExists":
+            raise error
+
+
+def add_role_to_instance_profile(instance_profile_name: str, role_name: str):
+    iam_client = boto3.client("iam")
+    iam_client.add_role_to_instance_profile(InstanceProfileName=instance_profile_name, RoleName=role_name)

tests/unittests/general/test_iam.py

@@ -0,0 +1,34 @@
+import boto3
+from moto import mock_iam
+
+from autogluon.cloud.utils.iam import (
+    add_role_to_instance_profile,
+    attach_iam_policy,
+    create_iam_policy,
+    create_iam_role,
+    create_instance_profile,
+)
+
+
+@mock_iam
+def test_iam_utils():
+    iam_client = boto3.client("iam")
+    dummy_role = "dummy_role"
+    dummy_trust_relationship = {}
+    dummy_role_arn = create_iam_role(dummy_role, dummy_trust_relationship)
+    assert dummy_role_arn is not None
+    create_iam_role(dummy_role, dummy_trust_relationship)  # check for recreation
+    dummy_policy = {
+        "Version": "2012-10-17",
+        "Statement": [{"Effect": "Allow", "Action": "none:null", "Resource": "*"}],
+    }
+    dummy_policy_arn = create_iam_policy("dommy_policy", dummy_policy)
+    assert dummy_policy_arn is not None
+    attach_iam_policy(dummy_role, dummy_policy_arn)
+    attached_policy = iam_client.list_attached_role_policies(RoleName=dummy_role)["AttachedPolicies"]
+    assert attached_policy[0]["PolicyArn"] == dummy_policy_arn
+    dummy_instance_profile = "dummy_instance_profile"
+    create_instance_profile(dummy_instance_profile)
+    add_role_to_instance_profile(dummy_instance_profile, dummy_role)
+    instance_profile = iam_client.get_instance_profile(InstanceProfileName=dummy_instance_profile)["InstanceProfile"]
+    assert instance_profile["Roles"][0]["Arn"] == dummy_role_arn