core: add FlowToken which saves the pickled flow plan, replace standard token in email stage to allow finishing flows in different sessions

closes #1801

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

authentik/core/models.py

@@ -25,7 +25,6 @@ from structlog.stdlib import get_logger
 from authentik.core.exceptions import PropertyMappingExpressionException
 from authentik.core.signals import password_changed
 from authentik.core.types import UILoginButton, UserSettingSerializer
-from authentik.flows.models import Flow
 from authentik.lib.config import CONFIG
 from authentik.lib.generators import generate_id
 from authentik.lib.models import CreatedUpdatedModel, DomainlessURLValidator, SerializerModel
@@ -203,7 +202,7 @@ class Provider(SerializerModel):
     name = models.TextField()
 
     authorization_flow = models.ForeignKey(
-        Flow,
+        "authentik_flows.Flow",
         on_delete=models.CASCADE,
         help_text=_("Flow used when authorizing this provider."),
         related_name="provider_authorization",
@@ -324,7 +323,7 @@ class Source(ManagedModel, SerializerModel, PolicyBindingModel):
     property_mappings = models.ManyToManyField("PropertyMapping", default=None, blank=True)
 
     authentication_flow = models.ForeignKey(
-        Flow,
+        "authentik_flows.Flow",
         blank=True,
         null=True,
         default=None,
@@ -333,7 +332,7 @@ class Source(ManagedModel, SerializerModel, PolicyBindingModel):
         related_name="source_authentication",
     )
     enrollment_flow = models.ForeignKey(
-        Flow,
+        "authentik_flows.Flow",
         blank=True,
         null=True,
         default=None,

authentik/flows/migrations/0020_flowtoken.py

@@ -0,0 +1,46 @@
+# Generated by Django 3.2.9 on 2021-12-05 13:50
+
+import django.db.models.deletion
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("authentik_core", "0018_auto_20210330_1345_squashed_0028_alter_token_intent"),
+        (
+            "authentik_flows",
+            "0019_alter_flow_background_squashed_0024_alter_flow_compatibility_mode",
+        ),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name="FlowToken",
+            fields=[
+                (
+                    "token_ptr",
+                    models.OneToOneField(
+                        auto_created=True,
+                        on_delete=django.db.models.deletion.CASCADE,
+                        parent_link=True,
+                        primary_key=True,
+                        serialize=False,
+                        to="authentik_core.token",
+                    ),
+                ),
+                ("_plan", models.TextField()),
+                (
+                    "flow",
+                    models.ForeignKey(
+                        on_delete=django.db.models.deletion.CASCADE, to="authentik_flows.flow"
+                    ),
+                ),
+            ],
+            options={
+                "verbose_name": "Flow Token",
+                "verbose_name_plural": "Flow Tokens",
+            },
+            bases=("authentik_core.token",),
+        ),
+    ]

authentik/flows/models.py

@@ -1,4 +1,6 @@
 """Flow models"""
+from base64 import b64decode, b64encode
+from pickle import dumps, loads  # nosec
 from typing import TYPE_CHECKING, Optional, Type
 from uuid import uuid4
 
@@ -9,11 +11,13 @@ from model_utils.managers import InheritanceManager
 from rest_framework.serializers import BaseSerializer
 from structlog.stdlib import get_logger
 
+from authentik.core.models import Token
 from authentik.core.types import UserSettingSerializer
 from authentik.lib.models import InheritanceForeignKey, SerializerModel
 from authentik.policies.models import PolicyBindingModel
 
 if TYPE_CHECKING:
+    from authentik.flows.planner import FlowPlan
     from authentik.flows.stage import StageView
 
 LOGGER = get_logger()
@@ -260,3 +264,30 @@ class ConfigurableStage(models.Model):
     class Meta:
 
         abstract = True
+
+
+class FlowToken(Token):
+    """Subclass of a standard Token, stores the currently active flow plan upon creation.
+    Can be used to later resume a flow."""
+
+    flow = models.ForeignKey(Flow, on_delete=models.CASCADE)
+    _plan = models.TextField()
+
+    @staticmethod
+    def pickle(plan) -> str:
+        """Pickle into string"""
+        data = dumps(plan)
+        return b64encode(data).decode()
+
+    @property
+    def plan(self) -> "FlowPlan":
+        """Load Flow plan from pickled version"""
+        return loads(b64decode(self._plan.encode()))  # nosec
+
+    def __str__(self) -> str:
+        return f"Flow Token {super.__str__()}"
+
+    class Meta:
+
+        verbose_name = _("Flow Token")
+        verbose_name_plural = _("Flow Tokens")

authentik/flows/planner.py

@@ -24,6 +24,9 @@ PLAN_CONTEXT_SSO = "is_sso"
 PLAN_CONTEXT_REDIRECT = "redirect"
 PLAN_CONTEXT_APPLICATION = "application"
 PLAN_CONTEXT_SOURCE = "source"
+# Is set by the Flow Planner when a FlowToken was used, and the currently active flow plan
+# was restored.
+PLAN_CONTEXT_IS_RESTORED = "is_restored"
 GAUGE_FLOWS_CACHED = UpdatingGauge(
     "authentik_flows_cached",
     "Cached flows",

authentik/flows/views/executor.py

@@ -34,8 +34,16 @@ from authentik.flows.challenge import (
     WithUserInfoChallenge,
 )
 from authentik.flows.exceptions import EmptyFlowException, FlowNonApplicableException
-from authentik.flows.models import ConfigurableStage, Flow, FlowDesignation, FlowStageBinding, Stage
+from authentik.flows.models import (
+    ConfigurableStage,
+    Flow,
+    FlowDesignation,
+    FlowStageBinding,
+    FlowToken,
+    Stage,
+)
 from authentik.flows.planner import (
+    PLAN_CONTEXT_IS_RESTORED,
     PLAN_CONTEXT_PENDING_USER,
     PLAN_CONTEXT_REDIRECT,
     FlowPlan,
@@ -55,6 +63,7 @@ SESSION_KEY_APPLICATION_PRE = "authentik_flows_application_pre"
 SESSION_KEY_GET = "authentik_flows_get"
 SESSION_KEY_POST = "authentik_flows_post"
 SESSION_KEY_HISTORY = "authentik_flows_history"
+QS_KEY_TOKEN = "flow_token"  # nosec
 
 
 def challenge_types():
@@ -127,8 +136,31 @@ class FlowExecutorView(APIView):
         message = exc.__doc__ if exc.__doc__ else str(exc)
         return self.stage_invalid(error_message=message)
 
+    def _check_flow_token(self, get_params: QueryDict):
+        """Check if the user is using a flow token to restore a plan"""
+        tokens = FlowToken.filter_not_expired(key=get_params[QS_KEY_TOKEN])
+        if not tokens.exists():
+            return False
+        token: FlowToken = tokens.first()
+        try:
+            plan = token.plan
+        except (AttributeError, EOFError, ImportError, IndexError) as exc:
+            LOGGER.warning("f(exec): Failed to restore token plan", exc=exc)
+        finally:
+            token.delete()
+        if not isinstance(plan, FlowPlan):
+            return None
+        plan.context[PLAN_CONTEXT_IS_RESTORED] = True
+        self._logger.debug("f(exec): restored flow plan from token", plan=plan)
+        return plan
+
     # pylint: disable=unused-argument, too-many-return-statements
     def dispatch(self, request: HttpRequest, flow_slug: str) -> HttpResponse:
+        get_params = QueryDict(request.GET.get("query", ""))
+        if QS_KEY_TOKEN in get_params:
+            plan = self._check_flow_token(get_params)
+            if plan:
+                self.request.session[SESSION_KEY_PLAN] = plan
         # Early check if there's an active Plan for the current session
         if SESSION_KEY_PLAN in self.request.session:
             self.plan = self.request.session[SESSION_KEY_PLAN]
@@ -156,7 +188,7 @@ class FlowExecutorView(APIView):
                 # we don't show an error message here, but rather call _flow_done()
                 return self._flow_done()
         # Initial flow request, check if we have an upstream query string passed in
-        request.session[SESSION_KEY_GET] = QueryDict(request.GET.get("query", ""))
+        request.session[SESSION_KEY_GET] = get_params
         # We don't save the Plan after getting the next stage
         # as it hasn't been successfully passed yet
         try:

authentik/stages/email/stage.py

@@ -12,17 +12,16 @@ from rest_framework.fields import CharField
 from rest_framework.serializers import ValidationError
 from structlog.stdlib import get_logger
 
-from authentik.core.models import Token
 from authentik.flows.challenge import Challenge, ChallengeResponse, ChallengeTypes
-from authentik.flows.planner import PLAN_CONTEXT_PENDING_USER
+from authentik.flows.models import FlowToken
+from authentik.flows.planner import PLAN_CONTEXT_IS_RESTORED, PLAN_CONTEXT_PENDING_USER
 from authentik.flows.stage import ChallengeStageView
-from authentik.flows.views.executor import SESSION_KEY_GET
+from authentik.flows.views.executor import QS_KEY_TOKEN, SESSION_KEY_GET
 from authentik.stages.email.models import EmailStage
 from authentik.stages.email.tasks import send_mails
 from authentik.stages.email.utils import TemplateEmailMessage
 
 LOGGER = get_logger()
-QS_KEY_TOKEN = "etoken"  # nosec
 PLAN_CONTEXT_EMAIL_SENT = "email_sent"
 
 
@@ -56,7 +55,7 @@ class EmailStageView(ChallengeStageView):
         relative_url = f"{base_url}?{urlencode(kwargs)}"
         return self.request.build_absolute_uri(relative_url)
 
-    def get_token(self) -> Token:
+    def get_token(self) -> FlowToken:
         """Get token"""
         pending_user = self.executor.plan.context[PLAN_CONTEXT_PENDING_USER]
         current_stage: EmailStage = self.executor.current_stage
@@ -65,10 +64,14 @@ class EmailStageView(ChallengeStageView):
         )  # + 1 because django timesince always rounds down
         identifier = slugify(f"ak-email-stage-{current_stage.name}-{pending_user}")
         # Don't check for validity here, we only care if the token exists
-        tokens = Token.objects.filter(identifier=identifier)
+        tokens = FlowToken.objects.filter(identifier=identifier)
         if not tokens.exists():
-            return Token.objects.create(
+            return FlowToken.objects.create(
-                expires=now() + valid_delta, user=pending_user, identifier=identifier
+                expires=now() + valid_delta,
+                user=pending_user,
+                identifier=identifier,
+                flow=self.executor.flow,
+                _plan=FlowToken.pickle(self.executor.plan),
             )
         token = tokens.first()
         # Check if token is expired and rotate key if so
@@ -97,13 +100,9 @@ class EmailStageView(ChallengeStageView):
 
     def get(self, request: HttpRequest, *args, **kwargs) -> HttpResponse:
         # Check if the user came back from the email link to verify
-        if QS_KEY_TOKEN in request.session.get(SESSION_KEY_GET, {}):
+        if QS_KEY_TOKEN in request.session.get(
-            tokens = Token.filter_not_expired(key=request.session[SESSION_KEY_GET][QS_KEY_TOKEN])
+            SESSION_KEY_GET, {}
-            if not tokens.exists():
+        ) and self.executor.plan.context.get(PLAN_CONTEXT_IS_RESTORED, False):
-                return self.executor.stage_invalid(_("Invalid token"))
-            token = tokens.first()
-            self.executor.plan.context[PLAN_CONTEXT_PENDING_USER] = token.user
-            token.delete()
             messages.success(request, _("Successfully verified Email."))
             if self.executor.current_stage.activate_user_on_success:
                 self.executor.plan.context[PLAN_CONTEXT_PENDING_USER].is_active = True