From f3f6fe048d44b4463f17af8b3658a8819eb87bb8 Mon Sep 17 00:00:00 2001
From: Philipp Rudiger <prudiger@anaconda.com>
Date: Wed, 8 Jul 2020 19:39:14 +0200
Subject: [PATCH] Add oauth-jwt-user config variable (#1470)

---
 panel/auth.py          |  9 +++++++--
 panel/command/serve.py | 13 +++++++++++++
 panel/config.py        | 15 +++++++++++++++
 3 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/panel/auth.py b/panel/auth.py
index 4addb4c9af..132a321c92 100644
--- a/panel/auth.py
+++ b/panel/auth.py
@@ -202,7 +202,8 @@ async def get(self):
         await self.get_authenticated_user(**params)
 
     def _on_auth(self, user_info, access_token):
-        self.set_secure_cookie('user', user_info[self._USER_KEY])
+        user_key = config.oauth_jwt_user or self._USER_KEY
+        self.set_secure_cookie('user', user_info[user_key])
         id_token = base64url_encode(json.dumps(user_info))
         if state.encryption:
             access_token = state.encryption.encrypt(access_token.encode('utf-8'))
@@ -434,7 +435,8 @@ def _on_auth(self, id_token, access_token):
         signing_input, _ = id_token.encode('utf-8').rsplit(b".", 1)
         _, payload_segment = signing_input.split(b".", 1)
         decoded = json.loads(base64url_decode(payload_segment).decode('utf-8'))
-        self.set_secure_cookie('user', decoded['email'])
+        user_key = config.oauth_jwt_user or self._USER_KEY
+        self.set_secure_cookie('user', decoded[user_key])
         if state.encryption:
             access_token = state.encryption.encrypt(access_token.encode('utf-8'))
             id_token = state.encryption.encrypt(id_token.encode('utf-8'))
@@ -454,6 +456,8 @@ class AzureAdLoginHandler(OAuthIDTokenLoginHandler, OAuth2Mixin):
     _OAUTH_AUTHORIZE_URL_ = 'https://login.microsoftonline.com/{tenant}/oauth2/authorize'
     _OAUTH_USER_URL_ = ''
 
+    _USER_KEY = 'unique_name'
+
     @property
     def _OAUTH_ACCESS_TOKEN_URL(self):
         return self._OAUTH_ACCESS_TOKEN_URL_.format(**config.oauth_extra_params)
@@ -478,6 +482,7 @@ class GoogleLoginHandler(OAuthIDTokenLoginHandler, OAuth2Mixin):
 
     _SCOPE = ['profile', 'email']
 
+    _USER_KEY = 'email'
 
 
 class LogoutHandler(tornado.web.RequestHandler):
diff --git a/panel/command/serve.py b/panel/command/serve.py
index 702a9090cf..169a338c5d 100644
--- a/panel/command/serve.py
+++ b/panel/command/serve.py
@@ -73,6 +73,11 @@ class Serve(_BkServe):
             type    = str,
             help    = "Additional parameters to use.",
         )),
+        ('--oauth-jwt-user', dict(
+            action  = 'store',
+            type    = str,
+            help    = "The key in the ID JWT token to consider the user.",
+        )),
         ('--oauth-encryption-key', dict(
             action = 'store',
             type    = str,
@@ -190,6 +195,14 @@ def customize_kwargs(self, args, server_kwargs):
             elif args.oauth_redirect_uri:
                 config.oauth_redirect_uri = args.oauth_redirect_uri
 
+            if args.oauth_jwt_user and config.oauth_jwt_user:
+                raise ValueError(
+                    "Supply OAuth JWT user either using environment "
+                    "variable or via explicit argument, not both."
+                )
+            elif args.oauth_jwt_user:
+                config.oauth_jwt_user = args.oauth_jwt_user
+
         if config.cookie_secret:
             kwargs['cookie_secret'] = config.cookie_secret
 
diff --git a/panel/config.py b/panel/config.py
index 1b12390abd..1210d3a7aa 100644
--- a/panel/config.py
+++ b/panel/config.py
@@ -119,6 +119,9 @@ class _config(param.Parameterized):
     _oauth_secret = param.String(default=None, doc="""
         A client secret to provide to the OAuth provider.""")
 
+    _oauth_jwt_user = param.String(default=None, doc="""
+        The key in the ID JWT token to consider the user.""")
+
     _oauth_redirect_uri = param.String(default=None, doc="""
         A redirect URI to provide to the OAuth provider.""")
 
@@ -319,6 +322,18 @@ def oauth_redirect_uri(self, value):
         validate_config(self, '_oauth_redirect_uri', value)
         self._oauth_redirect_uri_ = value
 
+    @property
+    def oauth_jwt_user(self):
+        if self._oauth_jwt_user_ is not None:
+            return self._oauth_jwt_user_
+        else:
+            return os.environ.get('PANEL_OAUTH_JWT_USER', _config._oauth_jwt_user)
+
+    @oauth_jwt_user.setter
+    def oauth_secret(self, value):
+        validate_config(self, '_oauth_jwt_user', value)
+        self._oauth_jwt_user_ = value
+
     @property
     def oauth_encryption_key(self):
         if self._oauth_encryption_key_ is not None: