[dy] Add OIDC generic provider (mage-ai#4563)

* [dy] Add OIDC generic provider

* [dy] Update environment variables

* [dy] Delete test code

Author: dy46

mage_ai/authentication/oauth/constants.py

@@ -16,6 +16,7 @@ class ProviderName(str, Enum):
     GHE = 'ghe'
     GOOGLE = 'google'
     OKTA = 'okta'
+    OIDC_GENERIC = 'oidc_generic'
 
 
 VALID_OAUTH_PROVIDERS = [e.value for e in ProviderName]

mage_ai/authentication/providers/constants.py

@@ -4,6 +4,7 @@
 from mage_ai.authentication.providers.ghe import GHEProvider
 from mage_ai.authentication.providers.gitlab import GitlabProvider
 from mage_ai.authentication.providers.google import GoogleProvider
+from mage_ai.authentication.providers.oidc import OidcProvider
 from mage_ai.authentication.providers.okta import OktaProvider
 
 NAME_TO_PROVIDER = {
@@ -12,5 +13,6 @@
     ProviderName.GHE: GHEProvider,
     ProviderName.GITLAB: GitlabProvider,
     ProviderName.GOOGLE: GoogleProvider,
+    ProviderName.OIDC_GENERIC: OidcProvider,
     ProviderName.OKTA: OktaProvider,
 }

mage_ai/authentication/providers/gitlab.py

@@ -45,8 +45,8 @@ def get_auth_url_response(self, redirect_uri: str = None, **kwargs) -> Dict:
                 f'{base_url}/oauth',
             ),
             response_type='code',
-            state=uuid.uuid4().hex,
             scope='read_user+write_repository+api',
+            state=uuid.uuid4().hex,
         )
         query_strings = []
         for k, v in query.items():

mage_ai/authentication/providers/oidc.py

@@ -0,0 +1,132 @@
+import urllib.parse
+import uuid
+from typing import Awaitable, Dict
+
+import aiohttp
+import requests
+
+from mage_ai.authentication.oauth.constants import ProviderName
+from mage_ai.authentication.providers.oauth import OauthProvider
+from mage_ai.authentication.providers.sso import SsoProvider
+from mage_ai.authentication.providers.utils import get_base_url
+from mage_ai.server.logger import Logger
+from mage_ai.settings.sso import OIDC_CLIENT_ID, OIDC_CLIENT_SECRET, OIDC_DISCOVERY_URL
+
+logger = Logger().new_server_logger(__name__)
+
+
+class OidcProvider(OauthProvider, SsoProvider):
+    provider = ProviderName.OIDC_GENERIC
+
+    def __init__(self):
+        self.discovery_url = OIDC_DISCOVERY_URL
+        self.client_id = OIDC_CLIENT_ID
+        self.client_secret = OIDC_CLIENT_SECRET
+        self.__validate()
+
+        self.discover()
+
+    def __validate(self):
+        if not self.discovery_url:
+            raise Exception(
+                'OIDC discovery url is empty. '
+                'Make sure the OIDC_DISCOVERY_URL environment variable is set.')
+        if not self.client_id:
+            raise Exception(
+                'OIDC client id is empty. '
+                'Make sure the OIDC_CLIENT_ID environment variable is set.')
+
+    def discover(self) -> Dict:
+        """
+        Call discovery url to get the endpoints for the OIDC server
+        """
+        try:
+            response = requests.get(
+                self.discovery_url,
+                headers={
+                    'Accept': 'application/json',
+                },
+                timeout=10,
+            )
+
+            response.raise_for_status()
+        except Exception:
+            logger.exception('Could not fetch response from OIDC discovery url')
+            raise
+
+        data = response.json()
+
+        self.authorization_endpoint = data.get('authorization_endpoint')
+        self.token_endpoint = data.get('token_endpoint')
+        self.userinfo_endpoint = data.get('userinfo_endpoint')
+
+    def get_auth_url_response(self, redirect_uri: str = None, **kwargs) -> Dict:
+        base_url = get_base_url(redirect_uri)
+        redirect_uri_query = dict(
+            provider=self.provider,
+            redirect_uri=redirect_uri,
+        )
+        query = dict(
+            client_id=self.client_id,
+            redirect_uri=urllib.parse.quote_plus(
+                f'{base_url}/oauth',
+            ),
+            response_type='code',
+            scope='openid profile email',
+            state=uuid.uuid4().hex,
+        )
+        query_strings = []
+        for k, v in query.items():
+            query_strings.append(f'{k}={v}')
+
+        return dict(
+            url=f"{self.authorization_endpoint}?{'&'.join(query_strings)}",
+            redirect_query_params=redirect_uri_query,
+        )
+
+    async def get_access_token_response(self, code: str, **kwargs) -> Awaitable[Dict]:
+        base_url = get_base_url(kwargs.get('redirect_uri'))
+        data = dict()
+
+        payload = dict(
+            client_id=self.client_id,
+            grant_type='authorization_code',
+            code=code,
+            redirect_uri=f'{base_url}/oauth',
+        )
+
+        if self.client_secret:
+            payload['client_secret'] = self.client_secret
+
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                self.token_endpoint,
+                headers={
+                    'Accept': 'application/json',
+                },
+                data=payload,
+                timeout=20,
+            ) as response:
+                data = await response.json()
+
+        return data
+
+    async def get_user_info(self, access_token: str = None, **kwargs) -> Awaitable[Dict]:
+        if access_token is None:
+            raise Exception('Access token is required to fetch user info.')
+        async with aiohttp.ClientSession() as session:
+            async with session.get(
+                self.userinfo_endpoint,
+                headers={
+                    'Authorization': f'Bearer {access_token}',
+                },
+                timeout=10,
+            ) as response:
+                userinfo_resp = await response.json()
+
+        email = userinfo_resp.get('email')
+
+        return dict(
+            email=email,
+            username=userinfo_resp.get('preferred_username', email),
+        )

mage_ai/frontend/components/Sessions/OidcSignIn.tsx

@@ -0,0 +1,41 @@
+import React, { useMemo } from 'react';
+import { useRouter } from 'next/router';
+
+import KeyboardShortcutButton from '@oracle/elements/Button/KeyboardShortcutButton';
+import { SignInProps } from './constants';
+import { queryFromUrl } from '@utils/url';
+import { set } from '@storage/localStorage';
+
+type OidcSignInProps = {} & SignInProps;
+
+function OidcSignIn({
+  oauthResponse,
+}: OidcSignInProps) {
+  const router = useRouter();
+  const {
+    url: oauthUrl,
+    redirect_query_params: redirectQueryParams = {},
+  } = useMemo(() => oauthResponse || {}, [oauthResponse]);
+
+  return (
+    <>
+      {oauthUrl && (
+        <KeyboardShortcutButton
+          bold
+          inline
+          onClick={() => {
+            const q = queryFromUrl(oauthUrl);
+            const state = q.state;
+            set(state, redirectQueryParams);
+            router.push(oauthUrl);
+          }}
+          uuid="SignForm/oidc_generic"
+        >
+          Sign in with OIDC
+        </KeyboardShortcutButton>
+      )}
+    </>
+  );
+}
+
+export default OidcSignIn;

mage_ai/frontend/interfaces/OauthType.ts

@@ -1,5 +1,6 @@
 import GoogleSignIn from '@components/Sessions/GoogleSignIn';
 import MicrosoftSignIn from '@components/Sessions/MicrosoftSignIn';
+import OidcSignIn from '@components/Sessions/OidcSignIn';
 import OktaSignIn from '@components/Sessions/OktaSignIn';
 
 export enum OauthProviderEnum {
@@ -8,12 +9,14 @@ export enum OauthProviderEnum {
   GITHUB = 'github',
   GITLAB = 'gitlab',
   GOOGLE = 'google',
+  OIDC_GENERIC = 'oidc_generic',
   OKTA = 'okta',
 }
 
 export const OAUTH_PROVIDER_SIGN_IN_MAPPING = {
   [OauthProviderEnum.ACTIVE_DIRECTORY]: MicrosoftSignIn,
   [OauthProviderEnum.GOOGLE]: GoogleSignIn,
+  [OauthProviderEnum.OIDC_GENERIC]: OidcSignIn,
   [OauthProviderEnum.OKTA]: OktaSignIn,
 };
 

mage_ai/settings/__init__.py

@@ -165,21 +165,31 @@ def get_bool_value(value: str) -> bool:
     'SENTRY_DSN',
     'SENTRY_TRACES_SAMPLE_RATE',
     'MAGE_PUBLIC_HOST',
-    'ACTIVE_DIRECTORY_DIRECTORY_ID',
-    'ACTIVE_DIRECTORY_CLIENT_ID',
-    'ACTIVE_DIRECTORY_CLIENT_SECRET',
     'SCHEDULER_TRIGGER_INTERVAL',
     'REQUIRE_USER_PERMISSIONS',
     'ENABLE_PROMETHEUS',
     'OTEL_EXPORTER_OTLP_ENDPOINT',
     'OTEL_EXPORTER_OTLP_HTTP_ENDPOINT',
+    'MAX_FILE_CACHE_SIZE',
+    # Oauth variables
+    'ACTIVE_DIRECTORY_DIRECTORY_ID',
+    'ACTIVE_DIRECTORY_CLIENT_ID',
+    'ACTIVE_DIRECTORY_CLIENT_SECRET',
     'OKTA_DOMAIN_URL',
     'OKTA_CLIENT_ID',
     'OKTA_CLIENT_SECRET',
     'GOOGLE_CLIENT_ID',
     'GOOGLE_CLIENT_SECRET',
+    'OIDC_CLIENT_ID',
+    'OIDC_CLIENT_SECRET',
+    'OIDC_DISCOVERY_URL',
     'GHE_CLIENT_ID',
     'GHE_CLIENT_SECRET',
     'GHE_HOSTNAME',
-    'MAX_FILE_CACHE_SIZE',
+    'BITBUCKET_HOST',
+    'BITBUCKET_OAUTH_KEY',
+    'BITBUCKET_OAUTH_SECRET',
+    'GITLAB_HOST',
+    'GITLAB_CLIENT_ID',
+    'GITLAB_CLIENT_SECRET',
 ]

mage_ai/settings/sso.py

@@ -19,3 +19,9 @@
 ACTIVE_DIRECTORY_CLIENT_ID = os.getenv('ACTIVE_DIRECTORY_CLIENT_ID')
 ACTIVE_DIRECTORY_CLIENT_SECRET = os.getenv('ACTIVE_DIRECTORY_CLIENT_SECRET')
 ACTIVE_DIRECTORY_ROLES_MAPPING = os.getenv('ACTIVE_DIRECTORY_ROLES_MAPPING')
+
+# OIDC
+
+OIDC_CLIENT_ID = os.getenv('OIDC_CLIENT_ID')
+OIDC_CLIENT_SECRET = os.getenv('OIDC_CLIENT_SECRET')
+OIDC_DISCOVERY_URL = os.getenv('OIDC_DISCOVERY_URL')