Allow using validate_allocations to validate Keycloak group memberships

src/coldfront_plugin_cloud/kc_client.py

@@ -109,3 +109,22 @@ def get_user_groups(self, user_id) -> list[str]:
         r.raise_for_status()
         groups = GroupResponse.model_validate(r.json())
         return [group.name for group in groups.root]
+
+    def get_group_members(self, group_id) -> list[str]:
+        url = f"{self.base_url}/admin/realms/{self.realm}/groups/{group_id}/members"
+        page_size = 100  # Default KeyCloak page size https://www.keycloak.org/docs-api/latest/rest-api/index.html#_query_parameters_32
+        page_offset = 0
+        users = []
+
+        while True:
+            r = self.api_client.get(
+                url, params={"first": page_offset, "max": page_size}
+            )
+            r.raise_for_status()
+            batch = UserResponse.model_validate(r.json())
+            users.extend(user.username for user in batch.root)
+            if len(batch.root) < page_size:
+                break
+            page_offset += page_size
+
+        return users

src/coldfront_plugin_cloud/management/commands/validate_allocations.py

@@ -3,11 +3,13 @@
 from coldfront_plugin_cloud import attributes
 from coldfront_plugin_cloud import utils
 from coldfront_plugin_cloud import tasks
+from coldfront_plugin_cloud import signals
 
 from django.core.management.base import BaseCommand
 from coldfront.core.resource.models import Resource
 from coldfront.core.allocation.models import (
     Allocation,
+    AllocationUser,
 )
 from keystoneauth1.exceptions import http
 from kubernetes.dynamic import exceptions as k8s_exceptions
@@ -46,6 +48,66 @@ def check_institution_specific_code(self, allocation, apply):
                 utils.set_attribute_on_allocation(allocation, attr, "N/A")
                 logger.warning(f'Attribute "{attr}" added to allocation {alloc_str}')
 
+    def validate_keycloak_group_memberships(self, allocation: Allocation, apply: bool):
+        # Fetch or cache Keycloak client
+        kc_client = getattr(self, "_kc_client", None)
+        if kc_client is None:
+            kc_client = tasks.get_kc_client()
+            self._kc_client = kc_client
+
+        resource = allocation.resources.first()
+        group_template = resource.get_attribute(
+            attributes.RESOURCE_KEYCLOAK_GROUP_TEMPLATE
+        )
+        if not group_template:
+            logger.info(
+                f"Keycloak enabled but no group name template specified for resource {resource.name}. Skipping validation"
+            )
+            return
+
+        allocation_users: list[AllocationUser] = (
+            allocation.allocationuser_set.select_related("user").all()
+        )
+        allocation_usernames: set[str] = set(
+            au.user.username for au in allocation_users
+        )
+
+        group_name = tasks._get_keycloak_group_name(allocation, group_template)
+
+        try:
+            group_name = tasks._get_keycloak_group_name(allocation, group_template)
+        except KeyError as e:
+            logger.error(
+                f"Invalid Keycloak group template for allocation {allocation.pk}: missing template variable {e}. Skipping validation"
+            )
+            return
+
+        group_id = kc_client.get_group_id(group_name)
+        if group_id:
+            group_usernames = set(kc_client.get_group_members(group_id))
+        else:
+            group_usernames = set()
+
+        to_add = [
+            au for au in allocation_users if au.user.username not in group_usernames
+        ]
+        to_remove = group_usernames - allocation_usernames
+
+        for au in to_add:
+            username = au.user.username
+            logger.info(f"Adding user {username} to Keycloak group {group_name}")
+            if apply:
+                tasks.add_user_to_keycloak(au.pk)
+        for username in to_remove:
+            logger.info(f"Removing user {username} from Keycloak group {group_name}")
+            if apply:
+                if user_id := kc_client.get_user_id(username):
+                    kc_client.remove_user_from_group(user_id, group_id)
+                else:
+                    logger.warning(
+                        f"User {username} not found in Keycloak, cannot remove from group {group_name}."
+                    )
+
     def handle(self, *args, **options):
         for resource_name in self.PLUGIN_RESOURCE_NAMES:
             resource = Resource.objects.filter(resource_type__name=resource_name)
@@ -70,6 +132,11 @@ def handle(self, *args, **options):
                     )
                     continue
 
+                if signals.is_keycloak_enabled():
+                    self.validate_keycloak_group_memberships(
+                        allocation, options["apply"]
+                    )
+
                 # Check project exists in remote cluster
                 try:
                     allocator.get_project(project_id)

src/coldfront_plugin_cloud/tests/functional/keycloak/test_keycloak.py

@@ -1,4 +1,7 @@
+from unittest import mock
+
 from django.contrib.auth.models import User
+from django.core.management import call_command
 from coldfront.core.resource.models import ResourceAttribute, ResourceAttributeType
 
 from coldfront_plugin_cloud import tasks, kc_client, attributes, utils
@@ -21,6 +24,17 @@ def setUpTestData(cls) -> None:
             value="$resource_name/$allocated_project_id",
         )
 
+    def setUp(self) -> None:
+        mock_allocator = mock.MagicMock()
+        mock_allocator.allocation_str = "Test Allocation of project Test Project"
+        self.patcher = mock.patch(
+            "coldfront_plugin_cloud.tasks.find_allocator", return_value=mock_allocator
+        )
+        self.mock_find_allocator = self.patcher.start()
+
+    def tearDown(self) -> None:
+        self.patcher.stop()
+
     def new_keycloak_user(self, cf_username):
         url = f"{self.kc_admin_client.base_url}/admin/realms/{self.kc_admin_client.realm}/users"
         payload = {
@@ -51,10 +65,10 @@ def test_user_added_to_allocation(self):
         user = self.new_user()
         project = self.new_project(pi=user)
         allocation = self.new_allocation(project, self.resource, 1)
-        allocation_user = self.new_allocation_user(allocation, user)
+        self.new_allocation_user(allocation, user)
 
-        # Simulate triggering the allocation activate signal
-        tasks.add_user_to_keycloak(allocation_user.pk)
+        # Validation should add user to Keycloak group
+        call_command("validate_allocations", apply=True)
 
         # Check that the user exists in Keycloak
         user_id = self.kc_admin_client.get_user_id(user.username)
@@ -72,13 +86,14 @@ def test_user_removed_from_allocation(self):
         allocation = self.new_allocation(project, self.resource, 1)
         allocation_user = self.new_allocation_user(allocation, user)
 
-        tasks.add_user_to_keycloak(allocation_user.pk)
+        call_command("validate_allocations", apply=True)
 
         user_id = self.kc_admin_client.get_user_id(user.username)
         user_groups = self.kc_admin_client.get_user_groups(user_id)
         self.assertIn(f"{self.resource.name}/Test Value", user_groups)
 
-        tasks.remove_user_from_keycloak(allocation_user.pk)
+        allocation_user.delete()
+        call_command("validate_allocations", apply=True)
 
         # Check that the user is no longer in the group
         user_groups = self.kc_admin_client.get_user_groups(user_id)
@@ -91,10 +106,10 @@ def test_user_not_in_keycloak_added_to_allocation(self):
         allocation = self.new_allocation(
             project, self.resource, 1, attr_value="Test Not Created"
         )
-        allocation_user = self.new_allocation_user(allocation, user)
+        self.new_allocation_user(allocation, user)
 
         # Should not raise error
-        tasks.add_user_to_keycloak(allocation_user.pk)
+        call_command("validate_allocations", apply=True)
 
         user_id = self.kc_admin_client.get_user_id(user.username)
         self.assertIsNone(user_id)
@@ -127,12 +142,10 @@ def test_multiple_users_in_same_allocation(self):
 
         # Add multiple users to the allocation
         users = [self.new_user() for _ in range(3)]
-        allocation_users = [
-            self.new_allocation_user(allocation, user) for user in users
-        ]
+        for user in users:
+            self.new_allocation_user(allocation, user)
 
-        for allocation_user in allocation_users:
-            tasks.add_user_to_keycloak(allocation_user.pk)
+        call_command("validate_allocations", apply=True)
 
         # Verify all users are in the group
         for user in users:
@@ -151,8 +164,7 @@ def test_remove_one_user_keeps_others_in_group(self):
             self.new_allocation_user(allocation, user) for user in users
         ]
 
-        for allocation_user in allocation_users:
-            tasks.add_user_to_keycloak(allocation_user.pk)
+        call_command("validate_allocations", apply=True)
 
         tasks.remove_user_from_keycloak(allocation_users[0].pk)
 
@@ -181,10 +193,9 @@ def test_user_in_multiple_allocations_groups(self):
 
         # Add user to both allocations
         allocation_user1 = self.new_allocation_user(allocation1, user)
-        allocation_user2 = self.new_allocation_user(allocation2, user)
+        self.new_allocation_user(allocation2, user)
 
-        tasks.add_user_to_keycloak(allocation_user1.pk)
-        tasks.add_user_to_keycloak(allocation_user2.pk)
+        call_command("validate_allocations", apply=True)
 
         # Verify user is in both groups
         user_id = self.kc_admin_client.get_user_id(user.username)
@@ -229,3 +240,39 @@ def test_user_added_without_keycloak_group_template(self):
         self.assertIsNotNone(user_id)
         user_groups = self.kc_admin_client.get_user_groups(user_id)
         self.assertEqual(user_groups, [])
+
+
+class TestKeyCloakGetGroupMembersPagination(base.TestBase):
+    @classmethod
+    def setUpTestData(cls) -> None:
+        super().setUpTestData()
+        cls.kc_admin_client = kc_client.KeyCloakAPIClient()
+
+    def new_keycloak_user(self, cf_username):
+        url = f"{self.kc_admin_client.base_url}/admin/realms/{self.kc_admin_client.realm}/users"
+        payload = {
+            "username": cf_username,
+            "enabled": True,
+            "email": cf_username,
+        }
+        r = self.kc_admin_client.api_client.post(url, json=payload)
+        r.raise_for_status()
+
+    def test_get_group_members_pagination(self):
+        group_name = "Test Pagination Group"
+        self.kc_admin_client.create_group(group_name)
+        group_id = self.kc_admin_client.get_group_id(group_name)
+
+        # Create 250 users and add them to the group (to ensure pagination is needed, as page size is 100)
+        for i in range(250):
+            username = f"pagination_user_{i}@example.com"
+            self.new_user(username=username)
+            self.new_keycloak_user(username)
+            self.kc_admin_client.add_user_to_group(
+                self.kc_admin_client.get_user_id(username), group_id
+            )
+
+        members = self.kc_admin_client.get_group_members(group_id)
+        self.assertEqual(len(members), 250)
+        expected_usernames = {f"pagination_user_{i}@example.com" for i in range(250)}
+        self.assertEqual(set(members), expected_usernames)