From c4bd8c2ae952b09c4f38e7f47ce1fe9827c66545 Mon Sep 17 00:00:00 2001 From: Jens Langhammer Date: Wed, 29 Jan 2025 17:50:28 +0100 Subject: [PATCH] enterprise/stages/source: fix Source stage not executing authentication/enrollment flow Signed-off-by: Jens Langhammer --- authentik/core/sources/flow_manager.py | 48 +++++++++++---------- authentik/enterprise/stages/source/stage.py | 23 ++++++++-- 2 files changed, 46 insertions(+), 25 deletions(-) diff --git a/authentik/core/sources/flow_manager.py b/authentik/core/sources/flow_manager.py index 3853890bf749..58546f497389 100644 --- a/authentik/core/sources/flow_manager.py +++ b/authentik/core/sources/flow_manager.py @@ -35,8 +35,7 @@ FlowPlanner, ) from authentik.flows.stage import StageView -from authentik.flows.views.executor import NEXT_ARG_NAME, SESSION_KEY_GET, SESSION_KEY_PLAN -from authentik.lib.utils.urls import redirect_with_qs +from authentik.flows.views.executor import NEXT_ARG_NAME, SESSION_KEY_GET from authentik.lib.views import bad_request_message from authentik.policies.denied import AccessDeniedResponse from authentik.policies.utils import delete_none_values @@ -47,8 +46,9 @@ LOGGER = get_logger() -SESSION_KEY_OVERRIDE_FLOW_TOKEN = "authentik/flows/source_override_flow_token" # nosec PLAN_CONTEXT_SOURCE_GROUPS = "source_groups" +SESSION_KEY_SOURCE_FLOW_STAGES = "authentik/flows/source_flow_stages" +SESSION_KEY_OVERRIDE_FLOW_TOKEN = "authentik/flows/source_override_flow_token" # nosec class MessageStage(StageView): @@ -219,28 +219,28 @@ def _prepare_flow( } ) flow_context.update(self.policy_context) - if SESSION_KEY_OVERRIDE_FLOW_TOKEN in self.request.session: - token: FlowToken = self.request.session.get(SESSION_KEY_OVERRIDE_FLOW_TOKEN) - self._logger.info("Replacing source flow with overridden flow", flow=token.flow.slug) - plan = token.plan - plan.context[PLAN_CONTEXT_IS_RESTORED] = token - plan.context.update(flow_context) - for stage in self.get_stages_to_append(flow): - plan.append_stage(stage) - if stages: - for stage in stages: - plan.append_stage(stage) - self.request.session[SESSION_KEY_PLAN] = plan - flow_slug = token.flow.slug - token.delete() - return redirect_with_qs( - "authentik_core:if-flow", - self.request.GET, - flow_slug=flow_slug, - ) flow_context.setdefault(PLAN_CONTEXT_REDIRECT, final_redirect) if not flow: + # We only check for the flow token here if we don't have a flow, otherwise we rely on + # SESSION_KEY_SOURCE_FLOW_STAGES to delegate the usage of this token and dynamically add + # stages that deal with this token to return to another flow + if SESSION_KEY_OVERRIDE_FLOW_TOKEN in self.request.session: + token: FlowToken = self.request.session.get(SESSION_KEY_OVERRIDE_FLOW_TOKEN) + self._logger.info( + "Replacing source flow with overridden flow", flow=token.flow.slug + ) + plan = token.plan + plan.context[PLAN_CONTEXT_IS_RESTORED] = token + plan.context.update(flow_context) + for stage in self.get_stages_to_append(flow): + plan.append_stage(stage) + if stages: + for stage in stages: + plan.append_stage(stage) + redirect = plan.to_redirect(self.request, token.flow) + token.delete() + return redirect return bad_request_message( self.request, _("Configured flow does not exist."), @@ -259,6 +259,8 @@ def _prepare_flow( if stages: for stage in stages: plan.append_stage(stage) + for stage in self.request.session.get(SESSION_KEY_SOURCE_FLOW_STAGES, []): + plan.append_stage(stage) return plan.to_redirect(self.request, flow) def handle_auth( @@ -295,6 +297,8 @@ def handle_existing_link( # When request isn't authenticated we jump straight to auth if not self.request.user.is_authenticated: return self.handle_auth(connection) + # When an override flow token exists we actually still use a flow for link + # to continue the existing flow we came from if SESSION_KEY_OVERRIDE_FLOW_TOKEN in self.request.session: return self._prepare_flow(None, connection) connection.save() diff --git a/authentik/enterprise/stages/source/stage.py b/authentik/enterprise/stages/source/stage.py index 44d405d33a7f..1bd543a173ae 100644 --- a/authentik/enterprise/stages/source/stage.py +++ b/authentik/enterprise/stages/source/stage.py @@ -9,13 +9,13 @@ from guardian.shortcuts import get_anonymous_user from authentik.core.models import Source, User -from authentik.core.sources.flow_manager import SESSION_KEY_OVERRIDE_FLOW_TOKEN +from authentik.core.sources.flow_manager import SESSION_KEY_OVERRIDE_FLOW_TOKEN, SESSION_KEY_SOURCE_FLOW_STAGES from authentik.core.types import UILoginButton from authentik.enterprise.stages.source.models import SourceStage from authentik.flows.challenge import Challenge, ChallengeResponse -from authentik.flows.models import FlowToken +from authentik.flows.models import FlowToken, in_memory_stage from authentik.flows.planner import PLAN_CONTEXT_IS_RESTORED -from authentik.flows.stage import ChallengeStageView +from authentik.flows.stage import ChallengeStageView, StageView from authentik.lib.utils.time import timedelta_from_string PLAN_CONTEXT_RESUME_TOKEN = "resume_token" # nosec @@ -49,6 +49,7 @@ def dispatch(self, request: HttpRequest, *args: Any, **kwargs: Any) -> HttpRespo def get_challenge(self, *args, **kwargs) -> Challenge: resume_token = self.create_flow_token() self.request.session[SESSION_KEY_OVERRIDE_FLOW_TOKEN] = resume_token + self.request.session[SESSION_KEY_SOURCE_FLOW_STAGES] = [in_memory_stage(SourceStageFinal)] return self.login_button.challenge def create_flow_token(self) -> FlowToken: @@ -77,3 +78,19 @@ def create_flow_token(self) -> FlowToken: def challenge_valid(self, response: ChallengeResponse) -> HttpResponse: return self.executor.stage_ok() + + +class SourceStageFinal(StageView): + """Dynamic stage injected in the source flow manager. This is injected in the + flow the source flow manager picks (authentication or enrollment), and will run at the end. + This stage uses the override flow token to resume execution of the initial flow the + source stage is bound to.""" + + def dispatch(self): + token: FlowToken = self.request.session.get(SESSION_KEY_OVERRIDE_FLOW_TOKEN) + self._logger.info("Replacing source flow with overridden flow", flow=token.flow.slug) + plan = token.plan + plan.context[PLAN_CONTEXT_IS_RESTORED] = token + response = plan.to_redirect(self.request, token.flow) + token.delete() + return response