From 4f05dcec893495426bdd6d2ce79171d273c6bbf4 Mon Sep 17 00:00:00 2001
From: Jens Langhammer <jens.langhammer@beryju.org>
Date: Mon, 13 Dec 2021 23:45:11 +0100
Subject: [PATCH] sources/oauth: allow oauth types to override their login
 button challenge

Signed-off-by: Jens Langhammer <jens.langhammer@beryju.org>
---
 authentik/sources/oauth/models.py        | 16 +---------------
 authentik/sources/oauth/types/manager.py | 19 +++++++++++++++++++
 2 files changed, 20 insertions(+), 15 deletions(-)

diff --git a/authentik/sources/oauth/models.py b/authentik/sources/oauth/models.py
index f48e2a9c7..8da0f3646 100644
--- a/authentik/sources/oauth/models.py
+++ b/authentik/sources/oauth/models.py
@@ -8,7 +8,6 @@ from rest_framework.serializers import Serializer
 
 from authentik.core.models import Source, UserSourceConnection
 from authentik.core.types import UILoginButton, UserSettingSerializer
-from authentik.flows.challenge import ChallengeTypes, RedirectChallenge
 
 if TYPE_CHECKING:
     from authentik.sources.oauth.types.manager import SourceType
@@ -66,20 +65,7 @@ class OAuthSource(Source):
 
     @property
     def ui_login_button(self) -> UILoginButton:
-        provider_type = self.type
-        return UILoginButton(
-            challenge=RedirectChallenge(
-                instance={
-                    "type": ChallengeTypes.REDIRECT.value,
-                    "to": reverse(
-                        "authentik_sources_oauth:oauth-client-login",
-                        kwargs={"source_slug": self.slug},
-                    ),
-                }
-            ),
-            icon_url=provider_type().icon_url(),
-            name=self.name,
-        )
+        return self.type().ui_login_button()
 
     @property
     def ui_user_settings(self) -> Optional[UserSettingSerializer]:
diff --git a/authentik/sources/oauth/types/manager.py b/authentik/sources/oauth/types/manager.py
index f619148d8..a11f79ff7 100644
--- a/authentik/sources/oauth/types/manager.py
+++ b/authentik/sources/oauth/types/manager.py
@@ -3,8 +3,11 @@ from enum import Enum
 from typing import Callable, Optional, Type
 
 from django.templatetags.static import static
+from django.urls.base import reverse
 from structlog.stdlib import get_logger
 
+from authentik.core.types import UILoginButton
+from authentik.flows.challenge import ChallengeTypes, RedirectChallenge
 from authentik.sources.oauth.views.callback import OAuthCallback
 from authentik.sources.oauth.views.redirect import OAuthRedirect
 
@@ -37,6 +40,22 @@ class SourceType:
         """Get Icon URL for login"""
         return static(f"authentik/sources/{self.slug}.svg")
 
+    def ui_login_button(self) -> UILoginButton:
+        """Allow types to return custom challenges"""
+        return UILoginButton(
+            challenge=RedirectChallenge(
+                instance={
+                    "type": ChallengeTypes.REDIRECT.value,
+                    "to": reverse(
+                        "authentik_sources_oauth:oauth-client-login",
+                        kwargs={"source_slug": self.slug},
+                    ),
+                }
+            ),
+            icon_url=self.icon_url(),
+            name=self.name,
+        )
+
 
 class SourceTypeManager:
     """Manager to hold all Source types."""