diff --git a/authentik/enterprise/stages/source/stage.py b/authentik/enterprise/stages/source/stage.py index a469db6013..02f2637936 100644 --- a/authentik/enterprise/stages/source/stage.py +++ b/authentik/enterprise/stages/source/stage.py @@ -89,9 +89,9 @@ class SourceStageFinal(StageView): This stage uses the override flow token to resume execution of the initial flow the source stage is bound to.""" - def dispatch(self): + def dispatch(self, *args, **kwargs): token: FlowToken = self.request.session.get(SESSION_KEY_OVERRIDE_FLOW_TOKEN) - self._logger.info("Replacing source flow with overridden flow", flow=token.flow.slug) + 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) diff --git a/authentik/enterprise/stages/source/tests.py b/authentik/enterprise/stages/source/tests.py index 954f2c3206..5dc93a59ff 100644 --- a/authentik/enterprise/stages/source/tests.py +++ b/authentik/enterprise/stages/source/tests.py @@ -4,7 +4,8 @@ from django.urls import reverse from authentik.core.tests.utils import create_test_flow, create_test_user from authentik.enterprise.stages.source.models import SourceStage -from authentik.flows.models import FlowDesignation, FlowStageBinding, FlowToken +from authentik.enterprise.stages.source.stage import SourceStageFinal +from authentik.flows.models import FlowDesignation, FlowStageBinding, FlowToken, in_memory_stage from authentik.flows.planner import PLAN_CONTEXT_IS_RESTORED, FlowPlan from authentik.flows.tests import FlowTestCase from authentik.flows.views.executor import SESSION_KEY_PLAN @@ -87,6 +88,7 @@ class TestSourceStage(FlowTestCase): self.assertIsNotNone(flow_token) session = self.client.session plan: FlowPlan = session[SESSION_KEY_PLAN] + plan.insert_stage(in_memory_stage(SourceStageFinal), index=0) plan.context[PLAN_CONTEXT_IS_RESTORED] = flow_token session[SESSION_KEY_PLAN] = plan session.save() @@ -96,4 +98,6 @@ class TestSourceStage(FlowTestCase): reverse("authentik_api:flow-executor", kwargs={"flow_slug": flow.slug}), follow=True ) self.assertEqual(response.status_code, 200) - self.assertStageRedirects(response, reverse("authentik_core:root-redirect")) + self.assertStageRedirects( + response, reverse("authentik_core:if-flow", kwargs={"flow_slug": flow.slug}) + ) diff --git a/authentik/flows/planner.py b/authentik/flows/planner.py index b5dc77c8ba..eb5d9ccae1 100644 --- a/authentik/flows/planner.py +++ b/authentik/flows/planner.py @@ -76,10 +76,10 @@ class FlowPlan: self.bindings.append(binding) self.markers.append(marker or StageMarker()) - def insert_stage(self, stage: Stage, marker: StageMarker | None = None): + def insert_stage(self, stage: Stage, marker: StageMarker | None = None, index=1): """Insert stage into plan, as immediate next stage""" - self.bindings.insert(1, FlowStageBinding(stage=stage, order=0)) - self.markers.insert(1, marker or StageMarker()) + self.bindings.insert(index, FlowStageBinding(stage=stage, order=0)) + self.markers.insert(index, marker or StageMarker()) def redirect(self, destination: str): """Insert a redirect stage as next stage"""