stages/user_write: dynamic groups (#2901)

* stages/user_write: add dynamic groups

Signed-off-by: Jens Langhammer <jens.langhammer@beryju.org>

* simplify functions

Signed-off-by: Jens Langhammer <jens.langhammer@beryju.org>

authentik/stages/user_write/stage.py

@@ -20,6 +20,7 @@ from authentik.stages.prompt.stage import PLAN_CONTEXT_PROMPT
 from authentik.stages.user_write.signals import user_write
 
 LOGGER = get_logger()
+PLAN_CONTEXT_GROUPS = "group"
 
 
 class UserWriteStageView(StageView):
@@ -47,15 +48,8 @@ class UserWriteStageView(StageView):
         """Wrapper for post requests"""
         return self.get(request)
 
-    def get(self, request: HttpRequest) -> HttpResponse:
+    def ensure_user(self) -> tuple[User, bool]:
-        """Save data in the current flow to the currently pending user. If no user is pending,
+        """Ensure a user exists"""
-        a new user is created."""
-        if PLAN_CONTEXT_PROMPT not in self.executor.plan.context:
-            message = _("No Pending data.")
-            messages.error(request, message)
-            LOGGER.debug(message)
-            return self.executor.stage_invalid()
-        data = self.executor.plan.context[PLAN_CONTEXT_PROMPT]
         user_created = False
         if PLAN_CONTEXT_PENDING_USER not in self.executor.plan.context:
             self.executor.plan.context[PLAN_CONTEXT_PENDING_USER] = User(
@@ -68,16 +62,13 @@ class UserWriteStageView(StageView):
             )
             user_created = True
         user: User = self.executor.plan.context[PLAN_CONTEXT_PENDING_USER]
-        # Before we change anything, check if the user is the same as in the request
+        return user, user_created
-        # and we're updating a password. In that case we need to update the session hash
+
-        # Also check that we're not currently impersonating, so we don't update the session
+    def update_user(self, user: User):
-        should_update_seesion = False
+        """Update `user` with data from plan context
-        if (
+
-            any("password" in x for x in data.keys())
+        Only simple attributes are updated, nothing which requires a foreign key or m2m"""
-            and self.request.user.pk == user.pk
+        data = self.executor.plan.context[PLAN_CONTEXT_PROMPT]
-            and SESSION_IMPERSONATE_USER not in self.request.session
-        ):
-            should_update_seesion = True
         for key, value in data.items():
             setter_name = f"set_{key}"
             # Check if user has a setter for this key, like set_password
@@ -98,10 +89,6 @@ class UserWriteStageView(StageView):
                     LOGGER.debug("discarding key", key=key)
                     continue
                 UserWriteStageView.write_attribute(user, key, value)
-        # Extra check to prevent flows from saving a user with a blank username
-        if user.username == "":
-            LOGGER.warning("Aborting write to empty username", user=user)
-            return self.executor.stage_invalid()
         # Check if we're writing from a source, and save the source to the attributes
         if PLAN_CONTEXT_SOURCES_CONNECTION in self.executor.plan.context:
             if USER_ATTRIBUTE_SOURCES not in user.attributes or not isinstance(
@@ -112,17 +99,45 @@ class UserWriteStageView(StageView):
                 PLAN_CONTEXT_SOURCES_CONNECTION
             ]
             user.attributes[USER_ATTRIBUTE_SOURCES].append(connection.source.name)
+
+    def get(self, request: HttpRequest) -> HttpResponse:
+        """Save data in the current flow to the currently pending user. If no user is pending,
+        a new user is created."""
+        if PLAN_CONTEXT_PROMPT not in self.executor.plan.context:
+            message = _("No Pending data.")
+            messages.error(request, message)
+            LOGGER.debug(message)
+            return self.executor.stage_invalid()
+        data = self.executor.plan.context[PLAN_CONTEXT_PROMPT]
+        user, user_created = self.ensure_user()
+        # Before we change anything, check if the user is the same as in the request
+        # and we're updating a password. In that case we need to update the session hash
+        # Also check that we're not currently impersonating, so we don't update the session
+        should_update_session = False
+        if (
+            any("password" in x for x in data.keys())
+            and self.request.user.pk == user.pk
+            and SESSION_IMPERSONATE_USER not in self.request.session
+        ):
+            should_update_session = True
+        self.update_user(user)
+        # Extra check to prevent flows from saving a user with a blank username
+        if user.username == "":
+            LOGGER.warning("Aborting write to empty username", user=user)
+            return self.executor.stage_invalid()
         try:
             with transaction.atomic():
                 user.save()
                 if self.executor.current_stage.create_users_group:
                     user.ak_groups.add(self.executor.current_stage.create_users_group)
-        except IntegrityError as exc:
+                if PLAN_CONTEXT_GROUPS in self.executor.plan.context:
+                    user.ak_groups.add(*self.executor.plan.context[PLAN_CONTEXT_GROUPS])
+        except (IntegrityError, ValueError, TypeError) as exc:
             LOGGER.warning("Failed to save user", exc=exc)
             return self.executor.stage_invalid()
         user_write.send(sender=self, request=request, user=user, data=data, created=user_created)
         # Check if the password has been updated, and update the session auth hash
-        if should_update_seesion:
+        if should_update_session:
             update_session_auth_hash(self.request, user)
             LOGGER.debug("Updated session hash", user=user)
         LOGGER.debug(

authentik/stages/user_write/tests.py

@@ -16,7 +16,7 @@ from authentik.flows.tests.test_executor import TO_STAGE_RESPONSE_MOCK
 from authentik.flows.views.executor import SESSION_KEY_PLAN
 from authentik.stages.prompt.stage import PLAN_CONTEXT_PROMPT
 from authentik.stages.user_write.models import UserWriteStage
-from authentik.stages.user_write.stage import UserWriteStageView
+from authentik.stages.user_write.stage import PLAN_CONTEXT_GROUPS, UserWriteStageView
 
 
 class TestUserWriteStage(FlowTestCase):
@@ -30,6 +30,7 @@ class TestUserWriteStage(FlowTestCase):
             designation=FlowDesignation.AUTHENTICATION,
         )
         self.group = Group.objects.create(name="test-group")
+        self.other_group = Group.objects.create(name="other-group")
         self.stage = UserWriteStage.objects.create(
             name="write", create_users_as_inactive=True, create_users_group=self.group
         )
@@ -49,6 +50,7 @@ class TestUserWriteStage(FlowTestCase):
             "email": "test@beryju.org",
             "password": password,
         }
+        plan.context[PLAN_CONTEXT_GROUPS] = [self.other_group]
         plan.context[PLAN_CONTEXT_SOURCES_CONNECTION] = UserSourceConnection(source=self.source)
         session = self.client.session
         session[SESSION_KEY_PLAN] = plan
@@ -63,7 +65,7 @@ class TestUserWriteStage(FlowTestCase):
         user_qs = User.objects.filter(username=plan.context[PLAN_CONTEXT_PROMPT]["username"])
         self.assertTrue(user_qs.exists())
         self.assertTrue(user_qs.first().check_password(password))
-        self.assertEqual(list(user_qs.first().ak_groups.all()), [self.group])
+        self.assertEqual(list(user_qs.first().ak_groups.all()), [self.group, self.other_group])
         self.assertEqual(user_qs.first().attributes, {USER_ATTRIBUTE_SOURCES: [self.source.name]})
 
     def test_user_update(self):

website/docs/flow/stages/user_write.md

@@ -5,3 +5,15 @@ title: User write stage
 This stages writes data from the current context to the current pending user. If no user is pending, a new one is created.
 
 Newly created users can be created as inactive and can be assigned to a selected group.
+
+### Dynamic groups
+
+Starting with authentik 2022.5, users can be added to dynamic groups. To do so, simply set `groups` in the flow plan context before this stage is run, for example
+
+```python
+from authentik.core.models import Group
+group, _ = Group.objects.get_or_create(name="some-group")
+# ["groups"] *must* be set to an array of Group objects, names alone are not enough.
+request.context["groups"] = [group]
+return True
+```