Google Groups whitelist and admin group

oauthenticator/google.py

@@ -9,11 +9,11 @@
 import urllib.parse
 
 from tornado import gen
-from tornado.httpclient import AsyncHTTPClient
+from tornado.httpclient import HTTPRequest, AsyncHTTPClient
 from tornado.auth import GoogleOAuth2Mixin
 from tornado.web import HTTPError
 
-from traitlets import Unicode, List, default, validate
+from traitlets import Set, Unicode, List, default, validate
 
 from jupyterhub.auth import LocalAuthenticator
 from jupyterhub.utils import url_path_join
@@ -22,6 +22,21 @@
 
 
 class GoogleOAuthenticator(OAuthenticator, GoogleOAuth2Mixin):
+    google_api_url = Unicode("https://www.googleapis.com", config=True)
+
+    @default('google_api_url')
+    def _google_api_url(self):
+        """get default google apis url from env"""
+        google_api_url = os.getenv('GOOGLE_API_URL')
+
+        # default to googleapis.com
+        if not google_api_url:
+            google_api_url = 'https://www.googleapis.com'
+
+        return google_api_url
+
+    # add the following to your jupyterhub_config.py to check groups
+    # c.GoogleOAuthenticator.scope = ['openid', 'email', 'https://www.googleapis.com/auth/admin.directory.group.readonly']
     @default('scope')
     def _scope_default(self):
         return ['openid', 'email']
@@ -32,7 +47,15 @@ def _authorize_url_default(self):
 
     @default("token_url")
     def _token_url_default(self):
-        return "https://www.googleapis.com/oauth2/v4/token"
+        return "%s/oauth2/v4/token" % (self.google_api_url)
+
+    google_group_whitelist = Set(
+        config=True, help="Automatically whitelist members of selected groups"
+    )
+
+    admin_google_groups = Set(
+        config=True, help="Groups whose members should have Jupyterhub admin privileges"
+    )
 
     user_info_url = Unicode(
         "https://www.googleapis.com/oauth2/v1/userinfo", config=True
@@ -127,10 +150,68 @@ async def authenticate(self, handler, data=None):
                 # unambiguous domain, use only base name
                 username = user_email.split('@')[0]
 
-        return {
-            'name': username,
-            'auth_state': {'access_token': access_token, 'google_user': bodyjs},
-        }
+        # Check if user is a member of any admin groups.
+        # These checks are performed here, as it requires `access_token`.
+        is_admin = False
+        is_admin_group_specified = False
+        if self.admin_google_groups:
+            is_admin_group_specified = True
+            is_admin = await self._check_user_in_groups(self.admin_google_groups , user_email, user_email_domain, access_token)
+
+        # Check if user is a member of any whitelisted groups.
+        # These checks are performed here, as it requires `access_token`.
+        user_in_group = False
+        is_group_specified = False
+
+        if self.google_group_whitelist:
+            is_group_specified = True
+            user_in_group = await self._check_user_in_groups(self.google_group_whitelist , user_email, user_email_domain, access_token)
+
+        no_config_specified = not is_group_specified
+
+        if is_admin_group_specified and is_admin:
+            self.log.debug("%s is in the admin group", username)
+            return {
+                'name': username,
+                'auth_state': {'access_token': access_token, 'google_user': bodyjs},
+                'admin': is_admin,
+            }
+        elif is_admin_group_specified and is_group_specified and user_in_group:
+            self.log.debug("%s can login on this server", username)
+            return {
+                'name': username,
+                'auth_state': {'access_token': access_token, 'google_user': bodyjs},
+                'admin': is_admin,
+            }
+        elif (
+            (is_group_specified and user_in_group)
+            or no_config_specified
+        ):
+            self.log.debug("%s can login on this server", username)
+            return {
+                'name': username,
+                'auth_state': {'access_token': access_token, 'google_user': bodyjs},
+            }
+        else:
+            self.log.warning("%s not in group whitelist", username)
+            return None
+
+
+    async def _check_user_in_groups(self, groups, user_email, user_email_domain, access_token):
+        http_client = AsyncHTTPClient()
+        # Check if user is a member of any group in the whitelist
+        for group in groups:
+            url = "%s/admin/directory/v1/groups/%s/members/%s?access_token=%s" % (
+                self.google_api_url,
+                "%s@%s" % (group, user_email_domain),
+                user_email,
+                access_token,
+            )
+            req = HTTPRequest(url, method="GET")
+            resp = await http_client.fetch(req, raise_error=False)
+            if resp.code == 200:
+                return True  # user _is_ in group
+        return False
 
 
 class LocalGoogleOAuthenticator(LocalAuthenticator, GoogleOAuthenticator):