Added functionality to attach/detach a policy to a role for use with … (

#1266)

* Added functionality to attach/detach a policy to a role for use with Bedrock

* added tests for the bedrock view

* added migration file

* Updated policy name to match what is in terraform

* fixed error in name

controlpanel/api/aws.py

@@ -499,6 +499,34 @@ def delete_role(self, name):
 
         role.delete()
 
+    def attach_policy(self, iam_role_name, attach_policies):
+
+        try:
+            role = self.boto3_session.resource("iam").Role(iam_role_name)
+            role.load()
+        except botocore.exceptions.ClientError as e:
+            if e.response["Error"]["Code"] == "NoSuchEntity":
+                log.warning(f"Role {iam_role_name}: Does not exist")
+                raise e
+
+        for attach_policy in attach_policies or []:
+            role.attach_policy(
+                PolicyArn=iam_arn(f"policy/{attach_policy}"),
+            )
+
+    def remove_policy(self, iam_role_name, remove_policies):
+
+        try:
+            role = self.boto3_session.resource("iam").Role(iam_role_name)
+            role.load()
+        except botocore.exceptions.ClientError as e:
+            if e.response["Error"]["Code"] == "NoSuchEntity":
+                log.warning(f"Role {iam_role_name}: Does not exist")
+                raise e
+
+        for policy in remove_policies or []:
+            role.detach_policy(PolicyArn=iam_arn(f"policy/{policy}"))
+
     def list_role_names(self, prefix="/"):
         roles = self.boto3_session.resource("iam").roles.filter(PathPrefix=prefix).all()
         return [role.name for role in list(roles)]

controlpanel/api/cluster.py

@@ -199,6 +199,7 @@ class User(EntityResource):
     }
 
     READ_INLINE_POLICIES = f"{settings.ENV}-read-user-roles-inline-policies"
+    BEDROCK_POLICY_NAME = "analytical-platform-bedrock-integration"
 
     ATTACH_POLICIES = [
         READ_INLINE_POLICIES,
@@ -375,6 +376,13 @@ def has_required_installation_charts(self):
                 return False
         return True
 
+    def set_bedrock_access(self):
+        bedrock_policy = [self.BEDROCK_POLICY_NAME]
+        if self.user.is_bedrock_enabled:
+            self.aws_role_service.attach_policy(self.iam_role_name, bedrock_policy)
+        else:
+            self.aws_role_service.remove_policy(self.iam_role_name, bedrock_policy)
+
 
 class App(EntityResource):
     """

controlpanel/api/migrations/0035_user_is_bedrock_enabled.py

@@ -0,0 +1,17 @@
+# Generated by Django 4.2.7 on 2024-03-08 11:59
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("api", "0034_remove_null_blank_from_namespace"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="user",
+            name="is_bedrock_enabled",
+            field=models.BooleanField(default=False),
+        ),
+    ]

controlpanel/api/models/user.py

@@ -37,6 +37,7 @@ class User(AbstractUser):
         choices=MIGRATION_STATES,
         default=VOID,
     )
+    is_bedrock_enabled = models.BooleanField(default=False)
 
     REQUIRED_FIELDS = ["email", "auth0_id"]
 
@@ -134,6 +135,9 @@ def is_bucket_admin(self, bucket_id):
     def reset_mfa(self):
         auth0.ExtendedAuth0().users.reset_mfa(self.auth0_id)
 
+    def set_bedrock_access(self):
+        cluster.User(self).set_bedrock_access()
+
     def save(self, *args, **kwargs):
         existing = User.objects.filter(pk=self.pk).first()
         if not existing:

controlpanel/frontend/jinja2/user-detail.html

@@ -63,6 +63,36 @@ <h2 class="govuk-heading-m govuk-error-summary-heading">
 </section>
 {% endif %}
 
+{% if request.user.has_perm('api.add_superuser') %}
+<section class="cpanel-section">
+  <form action="{{ url('set-bedrock', kwargs={ "pk": user.auth0_id }) }}" method="post">
+    {{ csrf_input }}
+    {{ govukCheckboxes({
+      "name": "is_bedrock_enabled",
+      "fieldset": {
+        "legend": {
+          "text": "Enable Bedrock",
+          "classes": "govuk-fieldset__legend--m",
+        },
+      },
+      "hint": {
+        "text": "Toggle access to Bedrock for a user."
+      },
+      "items": [
+        {
+          "value": "True",
+          "text": "Bedrock Enabled",
+          "checked": user.is_bedrock_enabled
+        },
+      ]
+    }) }}
+    <button class="govuk-button govuk-button--secondary">
+      Save changes
+    </button>
+  </form>
+</section>
+{% endif %}
+
 {% if request.user.has_perm('api.reset_mfa') %}
 <section class="cpanel-section">
   <form action="{{ url('reset-mfa', kwargs={ "pk": user.auth0_id }) }}" method="post">

controlpanel/frontend/urls.py

@@ -79,6 +79,7 @@
     path("users/", views.UserList.as_view(), name="list-users"),
     path("users/<str:pk>/", views.UserDetail.as_view(), name="manage-user"),
     path("users/<str:pk>/delete/", views.UserDelete.as_view(), name="delete-user"),
+    path("users/<str:pk>/bedrock/", views.EnableBedrockUser.as_view(), name="set-bedrock"),
     path("users/<str:pk>/edit/", views.SetSuperadmin.as_view(), name="set-superadmin"),
     path("users/<str:pk>/reset-mfa/", views.ResetMFA.as_view(), name="reset-mfa"),
     path(

controlpanel/frontend/views/__init__.py

@@ -73,6 +73,7 @@
 from controlpanel.frontend.views.user import (
     ResetMFA,
     SetSuperadmin,
+    EnableBedrockUser,
     UserDelete,
     UserDetail,
     UserList,

controlpanel/frontend/views/user.py

@@ -3,6 +3,8 @@
 
 # Third-party
 from django.contrib import messages
+from django.forms import BaseModelForm
+from django.http import HttpResponse
 from django.urls import reverse_lazy
 from django.views.generic.base import RedirectView
 from django.views.generic.detail import DetailView, SingleObjectMixin
@@ -94,6 +96,21 @@ def get_success_url(self):
         return reverse_lazy("manage-user", kwargs={"pk": self.object.auth0_id})
 
 
+class EnableBedrockUser(OIDCLoginRequiredMixin, PermissionRequiredMixin, UpdateView):
+    fields = ["is_bedrock_enabled"]
+    http_method_names = ["post"]
+    model = User
+    permission_required = "api.add_superuser"
+
+    def form_valid(self, form):
+        self.object.set_bedrock_access()
+        return super().form_valid(form)
+
+    def get_success_url(self):
+        messages.success(self.request, "Successfully updated bedrock status")
+        return reverse_lazy("manage-user", kwargs={"pk": self.object.auth0_id})
+
+
 class ResetMFA(
     OIDCLoginRequiredMixin,
     PermissionRequiredMixin,

tests/api/fixtures/aws.py

@@ -128,6 +128,52 @@ def airflow_prod_policy(iam):
     return result["Policy"]
 
 
+@pytest.fixture(autouse=True)
+def bedrock_policy(iam):
+    result = iam.meta.client.create_policy(
+        PolicyName="analytical-platform-bedrock-integration",
+        PolicyDocument=json.dumps(
+            {
+                "Version": "2012-10-17",
+                "Statement": [
+                    {
+                        "Sid": "BedrockEnable",
+                        "Effect": "Allow",
+                        "Action": ["iam:AllowBedrock"],
+                        "Resource": [
+                            "arn:aws:iam::{settings.AWS_DATA_ACCOUNT_ID}:role/{settings.ENV}_user_*"  # noqa: E501
+                        ],
+                    },
+                ],
+            }
+        ),
+    )
+    return result["Policy"]
+
+
+@pytest.fixture(autouse=True)
+def test_policy(iam):
+    result = iam.meta.client.create_policy(
+        PolicyName="a-test-policy",
+        PolicyDocument=json.dumps(
+            {
+                "Version": "2012-10-17",
+                "Statement": [
+                    {
+                        "Sid": "ThisIsATestPolicy",
+                        "Effect": "Allow",
+                        "Action": ["iam:TestPolicy"],
+                        "Resource": [
+                            "arn:aws:iam::{settings.AWS_DATA_ACCOUNT_ID}:role/{settings.ENV}_user_*"  # noqa: E501
+                        ],
+                    },
+                ],
+            }
+        ),
+    )
+    return result["Policy"]
+
+
 @pytest.fixture(autouse=True)
 def logs_bucket(s3):
     bucket = s3.Bucket(settings.LOGS_BUCKET_NAME)

tests/api/test_aws.py

@@ -1111,3 +1111,71 @@ def test_revoke_access_doesnt_remove_prefixes(s3_access_policy):
     assert s3_access_policy.statements[f"listSubFolders{bucket_hash}"].get("Resource", None) is not None
     assert s3_access_policy.statements[f"listFolder{bucket_hash}"]["Condition"]["StringEquals"]["s3:prefix"] != [""]  # noqa
     assert s3_access_policy.statements[f"listSubFolders{bucket_hash}"]["Condition"]["StringLike"]["s3:prefix"] != []  # noqa
+
+
+def test_attach_policy(iam, managed_policy, airflow_dev_policy, airflow_prod_policy, test_policy):
+    """
+    Ensure EKS settngs are in the policy document when running on that
+    infrastructure.
+    """
+    user = {
+        "auth0_id": "normal_user",
+        "user_name": "testing-bob",
+        "iam_role_name": "testing-bob",
+    }
+
+    aws.AWSRole().create_role(
+        user["iam_role_name"],
+        User.aws_user_policy(user["auth0_id"], user["user_name"]),
+        User.ATTACH_POLICIES,
+    )
+
+    aws.AWSRole().attach_policy(
+        user["iam_role_name"],
+        ["a-test-policy"]
+    )
+    role = iam.Role(user["iam_role_name"])
+
+
+    attached_policies = list(role.attached_policies.all())
+    assert len(attached_policies) == 4
+    arns = [policy.arn for policy in attached_policies]
+    assert managed_policy["Arn"] in arns
+    assert airflow_dev_policy["Arn"] in arns
+    assert airflow_prod_policy["Arn"] in arns
+    assert test_policy["Arn"] in arns
+
+
+def test_remove_policy(iam, managed_policy, airflow_dev_policy, airflow_prod_policy, test_policy):
+    """
+    Ensure EKS settngs are in the policy document when running on that
+    infrastructure.
+    """
+    user = {
+        "auth0_id": "normal_user",
+        "user_name": "testing-bob",
+        "iam_role_name": "testing-bob",
+    }
+
+    policies = User.ATTACH_POLICIES + ["a-test-policy"]
+
+    aws.AWSRole().create_role(
+        user["iam_role_name"],
+        User.aws_user_policy(user["auth0_id"], user["user_name"]),
+        policies,
+    )
+
+    aws.AWSRole().remove_policy(
+        user["iam_role_name"],
+        ["a-test-policy"]
+    )
+    role = iam.Role(user["iam_role_name"])
+
+
+    attached_policies = list(role.attached_policies.all())
+    assert len(attached_policies) == 3
+    arns = [policy.arn for policy in attached_policies]
+    assert managed_policy["Arn"] in arns
+    assert airflow_dev_policy["Arn"] in arns
+    assert airflow_prod_policy["Arn"] in arns
+    assert test_policy["Arn"] not in arns

tests/frontend/views/test_user.py

@@ -41,6 +41,14 @@ def reset_mfa(client, users, *args):
     return client.post(reverse("reset-mfa", kwargs=kwargs))
 
 
+def set_bedrock(client, users, *args):
+    data = {
+        "is_bedrock_enabled": True,
+    }
+    kwargs = {"pk": users["other_user"].auth0_id}
+    return client.post(reverse("set-bedrock", kwargs=kwargs), data)
+
+
 @pytest.mark.parametrize(
     "view,user,expected_status",
     [
@@ -58,6 +66,9 @@ def reset_mfa(client, users, *args):
         (reset_mfa, "superuser", status.HTTP_302_FOUND),
         (reset_mfa, "normal_user", status.HTTP_403_FORBIDDEN),
         (reset_mfa, "other_user", status.HTTP_403_FORBIDDEN),
+        (set_bedrock, "superuser", status.HTTP_302_FOUND),
+        (set_bedrock, "normal_user", status.HTTP_403_FORBIDDEN),
+        (set_bedrock, "other_user", status.HTTP_403_FORBIDDEN),
     ],
 )
 def test_permission(client, users, view, user, expected_status):