From 90151a13aea424186dd48d0a44c8d9bf3053eb53 Mon Sep 17 00:00:00 2001
From: Jens Langhammer <jens.langhammer@beryju.org>
Date: Sat, 6 Nov 2021 14:53:38 +0100
Subject: [PATCH] stages/identification: only allow limited challenges for
 login sources

Signed-off-by: Jens Langhammer <jens.langhammer@beryju.org>
---
 authentik/stages/identification/stage.py | 17 +++++++++++++----
 schema.yml                               | 11 ++++++++++-
 2 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/authentik/stages/identification/stage.py b/authentik/stages/identification/stage.py
index 64a24a0a6..925a936e0 100644
--- a/authentik/stages/identification/stage.py
+++ b/authentik/stages/identification/stage.py
@@ -16,10 +16,16 @@ from structlog.stdlib import get_logger
 
 from authentik.core.api.utils import PassiveSerializer
 from authentik.core.models import Application, Source, User
-from authentik.flows.challenge import Challenge, ChallengeResponse, ChallengeTypes
+from authentik.flows.challenge import (
+    Challenge,
+    ChallengeResponse,
+    ChallengeTypes,
+    RedirectChallenge,
+)
 from authentik.flows.planner import PLAN_CONTEXT_PENDING_USER
 from authentik.flows.stage import PLAN_CONTEXT_PENDING_USER_IDENTIFIER, ChallengeStageView
-from authentik.flows.views.executor import SESSION_KEY_APPLICATION_PRE, challenge_types
+from authentik.flows.views.executor import SESSION_KEY_APPLICATION_PRE
+from authentik.sources.plex.models import PlexAuthenticationChallenge
 from authentik.stages.identification.models import IdentificationStage
 from authentik.stages.identification.signals import identification_failed
 from authentik.stages.password.stage import authenticate
@@ -29,8 +35,11 @@ LOGGER = get_logger()
 
 @extend_schema_field(
     PolymorphicProxySerializer(
-        component_name="ChallengeTypes",
-        serializers=challenge_types(),
+        component_name="LoginChallengeTypes",
+        serializers={
+            RedirectChallenge().fields["component"].default: RedirectChallenge,
+            PlexAuthenticationChallenge().fields["component"].default: PlexAuthenticationChallenge,
+        },
         resource_type_field_name="component",
     )
 )
diff --git a/schema.yml b/schema.yml
index cf57b1451..18128bcb7 100644
--- a/schema.yml
+++ b/schema.yml
@@ -22444,6 +22444,15 @@ components:
           type: string
       required:
       - link
+    LoginChallengeTypes:
+      oneOf:
+      - $ref: '#/components/schemas/RedirectChallenge'
+      - $ref: '#/components/schemas/PlexAuthenticationChallenge'
+      discriminator:
+        propertyName: component
+        mapping:
+          xak-flow-redirect: '#/components/schemas/RedirectChallenge'
+          ak-flow-sources-plex: '#/components/schemas/PlexAuthenticationChallenge'
     LoginMetrics:
       type: object
       description: Login Metrics per 1h
@@ -22471,7 +22480,7 @@ components:
           type: string
           nullable: true
         challenge:
-          $ref: '#/components/schemas/ChallengeTypes'
+          $ref: '#/components/schemas/LoginChallengeTypes'
       required:
       - challenge
       - name