Merge pull request #3932 from allegro/add-auto-region-assignment

Add auto region assignment based on group

Author: hipek8

src/ralph/accounts/ldap.py

@@ -54,7 +54,7 @@ def mirror_groups(self):
     )
     if target_group_names != current_group_names:
         logger.info(
-            "Modifing user groups: current = {}, target = {}".format(
+            "Modifying user groups: current = {}, target = {}".format(
                 ", ".join(current_group_names), ", ".join(target_group_names)
             )
         )

src/ralph/accounts/ldap_helpers.py

@@ -71,13 +71,13 @@ def handle_groups(groups_dns):
         handle_groups(nested_groups_dns)
         return group_map
 
-    def group_name_from_info(self, group_info):
+    def group_name_from_info(self, group_info) -> str | None:
         """Map ldap group names into ralph names if mapping defined."""
         if self.ldap_groups:
             for dn in group_info[1]["distinguishedname"]:
-                mapped = self.ldap_groups.get(dn)
-                if mapped:
+                if mapped := self.ldap_groups.get(dn):
                     return mapped
+            return None
         # return original name if mapping not defined
         else:
             return super(MappedGroupOfNamesType, self).group_name_from_info(group_info)

src/ralph/accounts/management/commands/ldap_sync.py

@@ -11,6 +11,7 @@
 from functools import lru_cache
 from ldap.controls import SimplePagedResultsControl
 
+from ralph.accounts.models import RalphUser, Region
 from ralph.helpers import cache
 
 logger = logging.getLogger(__name__)
@@ -68,20 +69,20 @@ def __exit__(self, type, value, traceback):
 
 
 @cache(seconds=600)
-def get_nested_groups():
+def get_nested_groups() -> tuple[dict[str, set[str]], defaultdict[str, set[str]]]:
     """
     Fetching users in nested group based on custom LDAP filter
     (AUTH_LDAP_NESTED_FILTER) e.g. (memberOf:{}). AUTH_LDAP_NESTED_FILTER
     is a simple dictonary where key is the name of group in DB, the value
     contains DN for nested group.
     """
     # mapping from django group name to set of users (usernames) belonging to it
-    group_users = {}
+    group_name_to_usernames: dict[str, set[str]] = {}
     # mapping from user (username) to set of groups DNs to which he belongs to
-    users_groups = defaultdict(set)
+    username_to_group_names: defaultdict[str, set[str]] = defaultdict(set)
     nested_groups = getattr(settings, "AUTH_LDAP_NESTED_GROUPS", None)
     if not nested_groups:
-        return group_users, users_groups
+        return group_name_to_usernames, username_to_group_names
     nested_filter = getattr(settings, "AUTH_LDAP_NESTED_FILTER", "(memberOf:{})")
     logger.info("Fetching nested groups from LDAP")
     with LDAPConnectionManager() as conn:
@@ -99,7 +100,7 @@ def get_nested_groups():
                 settings.AUTH_LDAP_QUERY_PAGE_SIZE,
             )
             logger.info("{} fetched".format(ralph_group_name))
-            group_users[ralph_group_name] = set(
+            group_name_to_usernames[ralph_group_name] = set(
                 [
                     u[1][settings.AUTH_LDAP_USER_USERNAME_ATTR][0]
                     .decode("utf-8")
@@ -109,13 +110,13 @@ def get_nested_groups():
             )
             logger.info(
                 "Users in nested group {}: {}".format(
-                    ralph_group_name, group_users[ralph_group_name]
+                    ralph_group_name, group_name_to_usernames[ralph_group_name]
                 )
             )
-            for username in group_users[ralph_group_name]:
+            for username in group_name_to_usernames[ralph_group_name]:
                 # notice group DN here, not Django group name!
-                users_groups[username].add(ldap_group_name)
-    return group_users, users_groups
+                username_to_group_names[username].add(ldap_group_name)
+    return group_name_to_usernames, username_to_group_names
 
 
 def _make_paged_query(conn, search_base, search_scope, ad_query, attr_list, page_size):
@@ -156,7 +157,29 @@ def _make_paged_query(conn, search_base, search_scope, ad_query, attr_list, page
     return result
 
 
-class NestedGroups(object):
+def _add_regions(user: RalphUser, region_names: list[str]):
+    for region_name in region_names:
+        try:
+            user.regions.add(Region.objects.get(name=region_name))
+            logger.info("Assigned {} to region {}".format(user.username, region_name))
+        except Region.DoesNotExist:
+            logger.warning(
+                "Region {} does not exist, cannot assign to user {}".format(
+                    region_name,
+                    user.username,
+                )
+            )
+
+
+def assign_user_to_group(user: RalphUser, group: Group):
+    user.groups.add(group)
+    default_regions = settings.DEFAULT_REGIONS_FOR_GROUP.get(group.name, [])
+    _add_regions(user, default_regions)
+
+    logger.info("Added {} to {}".format(user.username, group.name))
+
+
+class NestedGroups:
     """
     Class fetch nested groups and mapping them to standard Django's
     group (get or create). django_auth_ldap and their class for nested
@@ -170,19 +193,17 @@ def __init__(self):
     def get_group_from_db(self, name):
         return Group.objects.get_or_create(name=name)[0]
 
-    def handle(self, user):
+    def handle(self, user: RalphUser):
         """
         Match user to group in fetched groups from LDAP and assign user
         to Django's group.
         """
-
         if not self.group_users:
             return
         for group_name, users in self.group_users.items():
             if user.username in users:
                 group = self.get_group_from_db(group_name)
-                user.groups.add(group)
-                logger.info("Added {} to {}".format(user.username, group_name))
+                assign_user_to_group(user, group)
 
 
 class Command(BaseCommand):

src/ralph/accounts/tests/tests.py

@@ -4,13 +4,17 @@
 
 from django.conf import settings
 from django.contrib.auth.hashers import check_password
-from django.contrib.auth.models import Permission
+from django.contrib.auth.models import Permission, Group
 from django.test import TestCase
 from django.urls import reverse
 from rest_framework import status
 
 from ralph.accounts.ldap import manager_country_attribute_populate
-from ralph.accounts.management.commands.ldap_sync import _truncate, ldap_module_exists
+from ralph.accounts.management.commands.ldap_sync import (
+    _truncate,
+    ldap_module_exists,
+    assign_user_to_group,
+)
 from ralph.accounts.models import RalphUser, Region
 from ralph.api.tests._base import RalphAPITestCase
 from ralph.assets.tests.factories import (
@@ -22,6 +26,7 @@
 from ralph.licences.tests.factories import LicenceFactory
 from ralph.tests import factories
 from ralph.tests.mixins import ClientMixin
+from django.test.utils import override_settings
 
 NO_LDAP_MODULE = not ldap_module_exists
 
@@ -234,3 +239,42 @@ def make_request():
         # Check if password is actually changed
         self.admin.refresh_from_db()
         check_password(new_password, self.admin.password)
+
+
+class RalphUserRegionTests(TestCase):
+    def test_user_region_str(self):
+        region = Region.objects.create(name="PL")
+        user = factories.UserFactory()
+        user.regions.add(region)
+        self.assertEqual(str(user.regions.first()), "PL")
+
+    def test_automatic_region_assignment(self):
+        region_pl = Region.objects.create(name="PL")
+        region_cz = Region.objects.create(name="CZ")
+        multi_region_group = Group.objects.create(name="Multi region group")
+        user = factories.UserFactory()
+
+        with override_settings(
+            DEFAULT_REGIONS_FOR_GROUP={"Multi region group": ["PL", "CZ"]}
+        ):
+            assign_user_to_group(user, multi_region_group)  # noqa
+
+        self.assertIn(region_cz, user.regions.all())
+        self.assertIn(region_pl, user.regions.all())
+
+    def test_automatic_region_assignment_when_user_already_in_group(self):
+        region_pl = Region.objects.create(name="PL")
+        region_cz = Region.objects.create(name="CZ")
+        multi_region_group = Group.objects.create(name="Multi region group")
+        user = factories.UserFactory()
+        assign_user_to_group(user, multi_region_group)  # noqa
+
+        self.assertEqual(user.regions.all().count(), 0)
+
+        with override_settings(
+            DEFAULT_REGIONS_FOR_GROUP={"Multi region group": ["PL", "CZ"]}
+        ):
+            assign_user_to_group(user, multi_region_group)  # noqa
+
+        self.assertIn(region_cz, user.regions.all())
+        self.assertIn(region_pl, user.regions.all())