From f28e8126af2a57f1f97ef199e986ab9a8989d4ee Mon Sep 17 00:00:00 2001
From: Yogesh Mahajan <yogesh.mahajan@enterprisedb.com>
Date: Thu, 21 Apr 2022 12:48:10 +0530
Subject: [PATCH] Added support for Azure AD OAUTH2 authentication. Fixes #7325

---
 docs/en_US/release_notes_6_9.rst   |  1 +
 web/pgadmin/authenticate/oauth2.py | 21 +++++++++++++--------
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/docs/en_US/release_notes_6_9.rst b/docs/en_US/release_notes_6_9.rst
index 199f27e9c..9a404c1f7 100644
--- a/docs/en_US/release_notes_6_9.rst
+++ b/docs/en_US/release_notes_6_9.rst
@@ -14,6 +14,7 @@ New features
   | `Issue #6830 <https://redmine.postgresql.org/issues/6830>`_ -  Relocate GIS Viewer Button to the Left Side of the Results Table.
   | `Issue #7012 <https://redmine.postgresql.org/issues/7012>`_ -  Disable the master password requirement when using alternative authentication sources.
   | `Issue #7282 <https://redmine.postgresql.org/issues/7282>`_ -  Added options 'Ignore owner' and 'Ignore whitespace' to the schema diff panel.
+  | `Issue #7325 <https://redmine.postgresql.org/issues/7325>`_ -  Added support for Azure AD OAUTH2 authentication.
 
 Housekeeping
 ************
diff --git a/web/pgadmin/authenticate/oauth2.py b/web/pgadmin/authenticate/oauth2.py
index 935d110a7..07d398380 100644
--- a/web/pgadmin/authenticate/oauth2.py
+++ b/web/pgadmin/authenticate/oauth2.py
@@ -88,6 +88,7 @@ class OAuth2Authentication(BaseAuthentication):
     oauth_obj = OAuth(Flask(__name__))
     oauth2_clients = {}
     oauth2_config = {}
+    email_keys = ['mail', 'email']
 
     def __init__(self):
         for oauth2_config in config.OAUTH2_CONFIG:
@@ -119,7 +120,11 @@ class OAuth2Authentication(BaseAuthentication):
 
     def login(self, form):
         profile = self.get_user_profile()
-        if 'email' not in profile or not profile['email']:
+        email_key = \
+            [value for value in self.email_keys if value in profile.keys()]
+        email = profile[email_key[0]] if (len(email_key) > 0) else None
+
+        if not email or email == '':
             current_app.logger.exception(
                 "An email id is required to login into pgAdmin. "
                 "Please update your Oauth2 profile."
@@ -128,10 +133,10 @@ class OAuth2Authentication(BaseAuthentication):
                 "An email id is required to login into pgAdmin. "
                 "Please update your Oauth2 profile.")
 
-        user, msg = self.__auto_create_user(profile)
+        user, msg = self.__auto_create_user(email)
         if user:
             user = db.session.query(User).filter_by(
-                username=profile['email'], auth_source=OAUTH2).first()
+                username=email, auth_source=OAUTH2).first()
             current_app.login_manager.logout_view = \
                 OAuth2Authentication.LOGOUT_VIEW
             return login_user(user), None
@@ -161,17 +166,17 @@ class OAuth2Authentication(BaseAuthentication):
         return False, self.oauth2_clients[
             self.oauth2_current_client].authorize_redirect(redirect_url)
 
-    def __auto_create_user(self, resp):
+    def __auto_create_user(self, email):
         if config.OAUTH2_AUTO_CREATE_USER:
-            user = User.query.filter_by(username=resp['email'],
+            user = User.query.filter_by(username=email,
                                         auth_source=OAUTH2).first()
             if not user:
                 return create_user({
-                    'username': resp['email'],
-                    'email': resp['email'],
+                    'username': email,
+                    'email': email,
                     'role': 2,
                     'active': True,
                     'auth_source': OAUTH2
                 })
 
-        return True, {'username': resp['email']}
+        return True, {'username': email}