From 241d790e69889b3800fa0d34966b12dd66de1fb7 Mon Sep 17 00:00:00 2001 From: Jens Langhammer Date: Thu, 6 May 2021 22:08:06 +0200 Subject: [PATCH] stages/user_write: if any connection is being sent in the plan context, save it to the user Signed-off-by: Jens Langhammer --- authentik/core/sources/flow_manager.py | 6 ------ authentik/sources/plex/api.py | 1 - authentik/stages/user_write/stage.py | 15 +++++++++++++-- 3 files changed, 13 insertions(+), 9 deletions(-) diff --git a/authentik/core/sources/flow_manager.py b/authentik/core/sources/flow_manager.py index 3426636ca..b3b8bb3aa 100644 --- a/authentik/core/sources/flow_manager.py +++ b/authentik/core/sources/flow_manager.py @@ -11,7 +11,6 @@ from django.utils.translation import gettext as _ from structlog.stdlib import get_logger from authentik.core.models import ( - USER_ATTRIBUTE_SOURCES, Source, SourceUserMatchingModes, User, @@ -271,11 +270,6 @@ class SourceFlowManager: if not self.source.enrollment_flow: self._logger.warning("source has no enrollment flow") return HttpResponseBadRequest() - if USER_ATTRIBUTE_SOURCES not in self.enroll_info or not isinstance( - self.enroll_info[USER_ATTRIBUTE_SOURCES], list - ): - self.enroll_info[USER_ATTRIBUTE_SOURCES] = [] - self.enroll_info[USER_ATTRIBUTE_SOURCES].append(self.source.name) return self._handle_login_flow( self.source.enrollment_flow, **{ diff --git a/authentik/sources/plex/api.py b/authentik/sources/plex/api.py index 38a35c994..800f4112d 100644 --- a/authentik/sources/plex/api.py +++ b/authentik/sources/plex/api.py @@ -1,5 +1,4 @@ """Plex Source Serializer""" -from django.http import Http404 from django.shortcuts import get_object_or_404 from drf_yasg import openapi from drf_yasg.utils import swagger_auto_schema diff --git a/authentik/stages/user_write/stage.py b/authentik/stages/user_write/stage.py index 39bdb4fb2..31fdac412 100644 --- a/authentik/stages/user_write/stage.py +++ b/authentik/stages/user_write/stage.py @@ -7,7 +7,8 @@ from django.utils.translation import gettext as _ from structlog.stdlib import get_logger from authentik.core.middleware import SESSION_IMPERSONATE_USER -from authentik.core.models import User +from authentik.core.models import USER_ATTRIBUTE_SOURCES, User, UserSourceConnection +from authentik.core.sources.stage import PLAN_CONTEXT_SOURCES_CONNECTION from authentik.flows.planner import PLAN_CONTEXT_PENDING_USER from authentik.flows.stage import StageView from authentik.lib.utils.reflection import class_to_path @@ -41,7 +42,7 @@ class UserWriteStageView(StageView): flow_slug=self.executor.flow.slug, ) user_created = True - user = self.executor.plan.context[PLAN_CONTEXT_PENDING_USER] + 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 # 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 @@ -73,6 +74,16 @@ class UserWriteStageView(StageView): 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( + user.attributes.get(USER_ATTRIBUTE_SOURCES), list + ): + user.attributes[USER_ATTRIBUTE_SOURCES] = [] + connection: UserSourceConnection = self.executor.plan.context[ + PLAN_CONTEXT_SOURCES_CONNECTION + ] + user.attributes[USER_ATTRIBUTE_SOURCES].append(connection.source.name) user.save() user_write.send( sender=self, request=request, user=user, data=data, created=user_created