authentik.enterprise.stages.source.stage

  1"""Source stage logic"""
  2
  3from typing import Any
  4from uuid import uuid4
  5
  6from django.http import HttpRequest, HttpResponse
  7from django.utils.text import slugify
  8from django.utils.timezone import now
  9from guardian.shortcuts import get_anonymous_user
 10
 11from authentik.core.models import Source, User
 12from authentik.core.sources.flow_manager import (
 13    SESSION_KEY_OVERRIDE_FLOW_TOKEN,
 14    SESSION_KEY_SOURCE_FLOW_CONTEXT,
 15    SESSION_KEY_SOURCE_FLOW_STAGES,
 16)
 17from authentik.core.types import UILoginButton
 18from authentik.enterprise.stages.source.models import SourceStage
 19from authentik.flows.challenge import Challenge, ChallengeResponse
 20from authentik.flows.models import FlowToken, in_memory_stage
 21from authentik.flows.planner import PLAN_CONTEXT_IS_REDIRECTED, PLAN_CONTEXT_IS_RESTORED
 22from authentik.flows.stage import ChallengeStageView, StageView
 23from authentik.lib.utils.time import timedelta_from_string
 24
 25PLAN_CONTEXT_RESUME_TOKEN = "resume_token"  # nosec
 26
 27
 28class SourceStageView(ChallengeStageView):
 29    """Suspend the current flow execution and send the user to a source,
 30    after which this flow execution is resumed."""
 31
 32    login_button: UILoginButton
 33
 34    def dispatch(self, request: HttpRequest, *args: Any, **kwargs: Any) -> HttpResponse:
 35        current_stage: SourceStage = self.executor.current_stage
 36        source: Source = (
 37            Source.objects.filter(pk=current_stage.source_id).select_subclasses().first()
 38        )
 39        if not source:
 40            self.logger.warning("Source does not exist")
 41            return self.executor.stage_invalid("Source does not exist")
 42        self.login_button = source.ui_login_button(self.request)
 43        if not self.login_button:
 44            self.logger.warning("Source does not have a UI login button")
 45            return self.executor.stage_invalid("Invalid source")
 46        restore_token = self.executor.plan.context.get(PLAN_CONTEXT_IS_RESTORED)
 47        override_token = self.request.session.get(SESSION_KEY_OVERRIDE_FLOW_TOKEN)
 48        if restore_token and override_token and restore_token.pk == override_token.pk:
 49            del self.request.session[SESSION_KEY_OVERRIDE_FLOW_TOKEN]
 50            return self.executor.stage_ok()
 51        return super().dispatch(request, *args, **kwargs)
 52
 53    def get_challenge(self, *args, **kwargs) -> Challenge:
 54        resume_token = self.create_flow_token()
 55        self.request.session[SESSION_KEY_OVERRIDE_FLOW_TOKEN] = resume_token
 56        self.request.session[SESSION_KEY_SOURCE_FLOW_STAGES] = [in_memory_stage(SourceStageFinal)]
 57        self.request.session[SESSION_KEY_SOURCE_FLOW_CONTEXT] = {
 58            PLAN_CONTEXT_IS_REDIRECTED: self.executor.flow,
 59        }
 60        return self.login_button.challenge
 61
 62    def create_flow_token(self) -> FlowToken:
 63        """Save the current flow state in a token that can be used to resume this flow"""
 64        pending_user: User = self.get_pending_user()
 65        if pending_user.is_anonymous or not pending_user.pk:
 66            pending_user = get_anonymous_user()
 67        current_stage: SourceStage = self.executor.current_stage
 68        identifier = slugify(f"ak-source-stage-{current_stage.name}-{str(uuid4())}")
 69        # Don't check for validity here, we only care if the token exists
 70        tokens = FlowToken.objects.filter(identifier=identifier)
 71        valid_delta = timedelta_from_string(current_stage.resume_timeout)
 72        if not tokens.exists():
 73            return FlowToken.objects.create(
 74                expires=now() + valid_delta,
 75                user=pending_user,
 76                identifier=identifier,
 77                flow=self.executor.flow,
 78                _plan=FlowToken.pickle(self.executor.plan),
 79            )
 80        token = tokens.first()
 81        # Check if token is expired and rotate key if so
 82        if token.is_expired:
 83            token.expire_action()
 84        return token
 85
 86    def challenge_valid(self, response: ChallengeResponse) -> HttpResponse:
 87        return self.executor.stage_ok()
 88
 89
 90class SourceStageFinal(StageView):
 91    """Dynamic stage injected in the source flow manager. This is injected in the
 92    flow the source flow manager picks (authentication or enrollment), and will run at the end.
 93    This stage uses the override flow token to resume execution of the initial flow the
 94    source stage is bound to."""
 95
 96    def dispatch(self, *args, **kwargs):
 97        token: FlowToken = self.request.session.get(SESSION_KEY_OVERRIDE_FLOW_TOKEN)
 98        self.logger.info("Replacing source flow with overridden flow", flow=token.flow.slug)
 99        plan = token.plan
100        plan.context.update(self.executor.plan.context)
101        plan.context[PLAN_CONTEXT_IS_RESTORED] = token
102        response = plan.to_redirect(self.request, token.flow)
103        token.delete()
104        return response