From dc4f341399d76fff6509a2cb0185855a2b294b9a Mon Sep 17 00:00:00 2001
From: "Jens L." <jens@goauthentik.io>
Date: Mon, 30 Sep 2024 23:27:17 +0200
Subject: [PATCH] stages/identification: dynamically find login challenges
 (#11571)

Signed-off-by: Jens Langhammer <jens@goauthentik.io>
---
 authentik/sources/oauth/types/apple.py   |  3 ++-
 authentik/sources/plex/models.py         |  3 ++-
 authentik/stages/identification/stage.py | 22 +++++++++++++++-------
 schema.yml                               |  4 ++--
 4 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/authentik/sources/oauth/types/apple.py b/authentik/sources/oauth/types/apple.py
index 1e3b8a2099d7..5c7645077f85 100644
--- a/authentik/sources/oauth/types/apple.py
+++ b/authentik/sources/oauth/types/apple.py
@@ -15,12 +15,13 @@
 from authentik.sources.oauth.types.registry import SourceType, registry
 from authentik.sources.oauth.views.callback import OAuthCallback
 from authentik.sources.oauth.views.redirect import OAuthRedirect
+from authentik.stages.identification.stage import LoginChallengeMixin
 
 LOGGER = get_logger()
 APPLE_CLIENT_ID_PARTS = 3
 
 
-class AppleLoginChallenge(Challenge):
+class AppleLoginChallenge(LoginChallengeMixin, Challenge):
     """Special challenge for apple-native authentication flow, which happens on the client."""
 
     client_id = CharField()
diff --git a/authentik/sources/plex/models.py b/authentik/sources/plex/models.py
index 86cda34ab43e..7ce05e598535 100644
--- a/authentik/sources/plex/models.py
+++ b/authentik/sources/plex/models.py
@@ -19,9 +19,10 @@
 from authentik.core.types import UILoginButton, UserSettingSerializer
 from authentik.flows.challenge import Challenge, ChallengeResponse
 from authentik.lib.generators import generate_id
+from authentik.stages.identification.stage import LoginChallengeMixin
 
 
-class PlexAuthenticationChallenge(Challenge):
+class PlexAuthenticationChallenge(LoginChallengeMixin, Challenge):
     """Challenge shown to the user in identification stage"""
 
     client_id = CharField()
diff --git a/authentik/stages/identification/stage.py b/authentik/stages/identification/stage.py
index 3cdccf3c0085..1a040d8cc1bf 100644
--- a/authentik/stages/identification/stage.py
+++ b/authentik/stages/identification/stage.py
@@ -26,23 +26,31 @@
 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, SESSION_KEY_GET
+from authentik.lib.utils.reflection import all_subclasses
 from authentik.lib.utils.urls import reverse_with_qs
 from authentik.root.middleware import ClientIPMiddleware
-from authentik.sources.oauth.types.apple import AppleLoginChallenge
-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
 
 
+class LoginChallengeMixin:
+    """Base login challenge for Identification stage"""
+
+
+def get_login_serializers():
+    mapping = {
+        RedirectChallenge().fields["component"].default: RedirectChallenge,
+    }
+    for cls in all_subclasses(LoginChallengeMixin):
+        mapping[cls().fields["component"].default] = cls
+    return mapping
+
+
 @extend_schema_field(
     PolymorphicProxySerializer(
         component_name="LoginChallengeTypes",
-        serializers={
-            RedirectChallenge().fields["component"].default: RedirectChallenge,
-            PlexAuthenticationChallenge().fields["component"].default: PlexAuthenticationChallenge,
-            AppleLoginChallenge().fields["component"].default: AppleLoginChallenge,
-        },
+        serializers=get_login_serializers,
         resource_type_field_name="component",
     )
 )
diff --git a/schema.yml b/schema.yml
index 646f7233678d..8beb39dd208e 100644
--- a/schema.yml
+++ b/schema.yml
@@ -41620,14 +41620,14 @@ components:
     LoginChallengeTypes:
       oneOf:
       - $ref: '#/components/schemas/RedirectChallenge'
-      - $ref: '#/components/schemas/PlexAuthenticationChallenge'
       - $ref: '#/components/schemas/AppleLoginChallenge'
+      - $ref: '#/components/schemas/PlexAuthenticationChallenge'
       discriminator:
         propertyName: component
         mapping:
           xak-flow-redirect: '#/components/schemas/RedirectChallenge'
-          ak-source-plex: '#/components/schemas/PlexAuthenticationChallenge'
           ak-source-oauth-apple: '#/components/schemas/AppleLoginChallenge'
+          ak-source-plex: '#/components/schemas/PlexAuthenticationChallenge'
     LoginMetrics:
       type: object
       description: Login Metrics per 1h