stages/identification: dynamically find login challenges (#11571)

Signed-off-by: Jens Langhammer <jens@goauthentik.io>

authentik/sources/oauth/types/apple.py

@@ -15,12 +15,13 @@ from authentik.sources.oauth.models import OAuthSource
 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()

authentik/sources/plex/models.py

@@ -19,9 +19,10 @@ from authentik.core.models import (
 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()

authentik/stages/identification/stage.py

@@ -26,23 +26,31 @@ from authentik.flows.models import FlowDesignation
 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",
     )
 )

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