From e7c3bde632928c4143bd942e278bb8201fe299b9 Mon Sep 17 00:00:00 2001 From: Jordan Bradford <36420801+jrdnbradford@users.noreply.github.com> Date: Thu, 26 Sep 2024 12:13:01 -0400 Subject: [PATCH 01/12] Add async group lookup --- oauthenticator/google.py | 77 ++++++++++++----------------- oauthenticator/tests/test_google.py | 3 +- 2 files changed, 34 insertions(+), 46 deletions(-) diff --git a/oauthenticator/google.py b/oauthenticator/google.py index f9d386b1..72ef7bd9 100644 --- a/oauthenticator/google.py +++ b/oauthenticator/google.py @@ -4,6 +4,7 @@ import os +import aiohttp from jupyterhub.auth import LocalAuthenticator from tornado.auth import GoogleOAuth2Mixin from tornado.web import HTTPError @@ -243,7 +244,7 @@ async def update_auth_model(self, auth_model): user_groups = set() if self.allowed_google_groups or self.admin_google_groups: - user_groups = self._fetch_member_groups(user_email, user_domain) + user_groups = await self._fetch_member_groups(user_email, user_domain) # sets are not JSONable, cast to list for auth_state user_info["google_groups"] = list(user_groups) @@ -314,7 +315,7 @@ async def check_allowed(self, username, auth_model): # users should be explicitly allowed via config, otherwise they aren't return False - def _service_client_credentials(self, scopes, user_email_domain): + async def _service_client_credentials(self, scopes, user_email_domain): """ Return a configured service client credentials for the API. """ @@ -338,47 +339,28 @@ def _service_client_credentials(self, scopes, user_email_domain): return credentials - def _service_client(self, service_name, service_version, credentials, http=None): + async def _setup_credentials(self, user_email_domain): """ - Return a configured service client for the API. + Set up the service client for Google API. """ + credentials = await self._service_client_credentials( + scopes=[f"{self.google_api_url}/auth/admin.directory.group.readonly"], + user_email_domain=user_email_domain, + ) + try: - from googleapiclient.discovery import build + from google.auth.transport import requests except: raise ImportError( - "Could not import googleapiclient.discovery's build," + "Could not import google.auth.transport's requests," "you may need to run 'pip install oauthenticator[googlegroups]' or not declare google groups" ) - self.log.debug( - f"service_name is {service_name}, service_version is {service_version}" - ) - - return build( - serviceName=service_name, - version=service_version, - credentials=credentials, - cache_discovery=False, - http=http, - ) - - def _setup_service(self, user_email_domain, http=None): - """ - Set up the service client for Google API. - """ - credentials = self._service_client_credentials( - scopes=[f"{self.google_api_url}/auth/admin.directory.group.readonly"], - user_email_domain=user_email_domain, - ) - service = self._service_client( - service_name='admin', - service_version='directory_v1', - credentials=credentials, - http=http, - ) - return service + request = requests.Request() + credentials.refresh(request) + return credentials - def _fetch_member_groups( + async def _fetch_member_groups( self, member_email, user_email_domain, @@ -389,17 +371,22 @@ def _fetch_member_groups( """ Return a set with the google groups a given user/group is a member of, including nested groups if allowed. """ - # FIXME: When this function is used and waiting for web request - # responses, JupyterHub gets blocked from doing other things. - # Ideally the web requests should be made using an async client - # that can be awaited while JupyterHub handles other things. - # - if not hasattr(self, 'service'): - self.service = self._setup_service(user_email_domain, http) - - resp = self.service.groups().list(userKey=member_email).execute() + if not hasattr(self, 'credentials'): + self.credentials = await self._setup_credentials(user_email_domain) + + headers = {'Authorization': f'Bearer {self.credentials.token}'} + + async with aiohttp.ClientSession() as session: + async with session.get( + f'https://www.googleapis.com/admin/directory/v1/groups?userKey={member_email}', + headers=headers, + ) as resp: + group_data = await resp.json() + member_groups = { - g['email'].split('@')[0] for g in resp.get('groups', []) if g.get('email') + g['email'].split('@')[0] + for g in group_data.get('groups', []) + if g.get('email') } self.log.debug(f"Fetched groups for {member_email}: {member_groups}") @@ -410,7 +397,7 @@ def _fetch_member_groups( for group in member_groups: if group not in processed_groups: processed_groups.add(group) - nested_groups = self._fetch_member_groups( + nested_groups = await self._fetch_member_groups( f"{group}@{user_email_domain}", user_email_domain, http, diff --git a/oauthenticator/tests/test_google.py b/oauthenticator/tests/test_google.py index 4c3a9e0c..070de014 100644 --- a/oauthenticator/tests/test_google.py +++ b/oauthenticator/tests/test_google.py @@ -3,6 +3,7 @@ import logging import re from unittest import mock +from unittest.mock import AsyncMock from pytest import fixture, mark, raises from traitlets.config import Config @@ -211,7 +212,7 @@ async def test_google( handled_user_model = user_model("user1@example.com", "user1") handler = google_client.handler_for_user(handled_user_model) with mock.patch.object( - authenticator, "_fetch_member_groups", lambda *args: {"group1"} + authenticator, "_fetch_member_groups", AsyncMock(return_value={"group1"}) ): auth_model = await authenticator.get_authenticated_user(handler, None) From d32940dc6d10df92261d112a7b5c9e186036599b Mon Sep 17 00:00:00 2001 From: Jordan Bradford <36420801+jrdnbradford@users.noreply.github.com> Date: Thu, 26 Sep 2024 18:44:35 -0400 Subject: [PATCH 02/12] Add `aiohttp` to `requirements` --- requirements.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/requirements.txt b/requirements.txt index 53cc9bc4..47a34eeb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ +aiohttp # jsonschema is used for validating authenticator configurations jsonschema jupyterhub>=2.2 From ebba631fb5746c4a990266e2e0d37d55f4522ed9 Mon Sep 17 00:00:00 2001 From: Jordan Bradford <36420801+jrdnbradford@users.noreply.github.com> Date: Thu, 3 Oct 2024 12:24:56 -0400 Subject: [PATCH 03/12] Use base `httpfetch` --- oauthenticator/google.py | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/oauthenticator/google.py b/oauthenticator/google.py index 9a568926..fdc50230 100644 --- a/oauthenticator/google.py +++ b/oauthenticator/google.py @@ -4,7 +4,6 @@ import os -import aiohttp from jupyterhub.auth import LocalAuthenticator from tornado.auth import GoogleOAuth2Mixin from tornado.web import HTTPError @@ -374,19 +373,15 @@ async def _fetch_member_groups( if not hasattr(self, 'credentials'): self.credentials = await self._setup_credentials(user_email_domain) - headers = {'Authorization': f'Bearer {self.credentials.token}'} - - async with aiohttp.ClientSession() as session: - async with session.get( - f'https://www.googleapis.com/admin/directory/v1/groups?userKey={member_email}', - headers=headers, - ) as resp: - group_data = await resp.json() - checked_groups = checked_groups or set() processed_groups = processed_groups or set() - resp = self.service.groups().list(userKey=member_email).execute() + headers = {'Authorization': f'Bearer {self.credentials.token}'} + url = f'https://www.googleapis.com/admin/directory/v1/groups?userKey={member_email}' + group_data = await self.httpfetch( + url, headers=headers, label="fetching google groups" + ) + member_groups = { g['email'].split('@')[0] for g in group_data.get('groups', []) @@ -399,7 +394,6 @@ async def _fetch_member_groups( if self.include_nested_groups: for group in member_groups: - if group in processed_groups: continue processed_groups.add(group) From 945831389028bc01812e50c34cf0693e301b476d Mon Sep 17 00:00:00 2001 From: Jordan Bradford <36420801+jrdnbradford@users.noreply.github.com> Date: Thu, 3 Oct 2024 12:32:50 -0400 Subject: [PATCH 04/12] Update `_setup_credentials` docs --- oauthenticator/google.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/oauthenticator/google.py b/oauthenticator/google.py index fdc50230..9a358698 100644 --- a/oauthenticator/google.py +++ b/oauthenticator/google.py @@ -340,7 +340,7 @@ async def _service_client_credentials(self, scopes, user_email_domain): async def _setup_credentials(self, user_email_domain): """ - Set up the service client for Google API. + Set up the oauth credentials for Google API. """ credentials = await self._service_client_credentials( scopes=[f"{self.google_api_url}/auth/admin.directory.group.readonly"], From ed38219f8da537c42527c8cea8d51ac9cb5e62ad Mon Sep 17 00:00:00 2001 From: Erik Sundell Date: Thu, 3 Oct 2024 20:35:19 +0200 Subject: [PATCH 05/12] Remove optional dependency on google api client, no longer used --- setup.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/setup.py b/setup.py index 480e04d3..615695be 100644 --- a/setup.py +++ b/setup.py @@ -91,7 +91,6 @@ def run(self): # googlegroups is required for use of GoogleOAuthenticator configured with # either admin_google_groups and/or allowed_google_groups. 'googlegroups': [ - 'google-api-python-client', 'google-auth-oauthlib', ], # mediawiki is required for use of MWOAuthenticator @@ -105,7 +104,6 @@ def run(self): 'pytest-cov', 'requests-mock', # dependencies from googlegroups: - 'google-api-python-client', 'google-auth-oauthlib', # dependencies from mediawiki: 'mwoauth>=0.3.8', From 7884369e7f16330ca1492d725c634f6d8b7b6116 Mon Sep 17 00:00:00 2001 From: Jordan Bradford <36420801+jrdnbradford@users.noreply.github.com> Date: Fri, 4 Oct 2024 09:34:29 -0400 Subject: [PATCH 06/12] Fetch creds per call --- oauthenticator/google.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/oauthenticator/google.py b/oauthenticator/google.py index 9a358698..bd6b9f76 100644 --- a/oauthenticator/google.py +++ b/oauthenticator/google.py @@ -366,17 +366,17 @@ async def _fetch_member_groups( http=None, checked_groups=None, processed_groups=None, + credentials=None, ): """ Return a set with the google groups a given user/group is a member of, including nested groups if allowed. """ - if not hasattr(self, 'credentials'): - self.credentials = await self._setup_credentials(user_email_domain) + credentials = credentials or await self._setup_credentials(user_email_domain) checked_groups = checked_groups or set() processed_groups = processed_groups or set() - headers = {'Authorization': f'Bearer {self.credentials.token}'} + headers = {'Authorization': f'Bearer {credentials.token}'} url = f'https://www.googleapis.com/admin/directory/v1/groups?userKey={member_email}' group_data = await self.httpfetch( url, headers=headers, label="fetching google groups" From 27f1245318a3d57c1b063c86c7a0c73a93137a84 Mon Sep 17 00:00:00 2001 From: Jordan Bradford <36420801+jrdnbradford@users.noreply.github.com> Date: Fri, 4 Oct 2024 09:53:06 -0400 Subject: [PATCH 07/12] Remove irrelevant `async` --- oauthenticator/google.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/oauthenticator/google.py b/oauthenticator/google.py index bd6b9f76..e945aa53 100644 --- a/oauthenticator/google.py +++ b/oauthenticator/google.py @@ -314,7 +314,7 @@ async def check_allowed(self, username, auth_model): # users should be explicitly allowed via config, otherwise they aren't return False - async def _service_client_credentials(self, scopes, user_email_domain): + def _service_client_credentials(self, scopes, user_email_domain): """ Return a configured service client credentials for the API. """ @@ -338,11 +338,11 @@ async def _service_client_credentials(self, scopes, user_email_domain): return credentials - async def _setup_credentials(self, user_email_domain): + def _setup_credentials(self, user_email_domain): """ Set up the oauth credentials for Google API. """ - credentials = await self._service_client_credentials( + credentials = self._service_client_credentials( scopes=[f"{self.google_api_url}/auth/admin.directory.group.readonly"], user_email_domain=user_email_domain, ) @@ -372,7 +372,7 @@ async def _fetch_member_groups( Return a set with the google groups a given user/group is a member of, including nested groups if allowed. """ - credentials = credentials or await self._setup_credentials(user_email_domain) + credentials = credentials or self._setup_credentials(user_email_domain) checked_groups = checked_groups or set() processed_groups = processed_groups or set() From 389d4b2a723fe48ddce215f8c20210124d0fc363 Mon Sep 17 00:00:00 2001 From: Jordan Bradford <36420801+jrdnbradford@users.noreply.github.com> Date: Fri, 4 Oct 2024 13:14:26 -0400 Subject: [PATCH 08/12] Verify token --- oauthenticator/google.py | 29 ++++++++++++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/oauthenticator/google.py b/oauthenticator/google.py index e945aa53..427f4a74 100644 --- a/oauthenticator/google.py +++ b/oauthenticator/google.py @@ -14,6 +14,7 @@ class GoogleOAuthenticator(OAuthenticator, GoogleOAuth2Mixin): user_auth_state_key = "google_user" + _credentials = None @default("login_service") def _login_service_default(self): @@ -314,6 +315,31 @@ async def check_allowed(self, username, auth_model): # users should be explicitly allowed via config, otherwise they aren't return False + def _get_credentials(self, user_email_domain): + """ + Returns the stored credentials or fetches and stores new ones. + + Checks if the credentials are valid before returning them. Refreshes + if necessary and stores the refreshed credentials. + """ + if not self._credentials or not self._is_token_valid(): + self._credentials = self._setup_credentials(user_email_domain) + + return self._credentials + + def _is_token_valid(self): + """ + Checks if the stored token is valid. + """ + if not self._credentials: + return False + if not self._credentials.token: + return False + if self._credentials.expired: + return False + + return True + def _service_client_credentials(self, scopes, user_email_domain): """ Return a configured service client credentials for the API. @@ -357,6 +383,7 @@ def _setup_credentials(self, user_email_domain): request = requests.Request() credentials.refresh(request) + self.log.debug("Credentials refreshed") return credentials async def _fetch_member_groups( @@ -372,7 +399,7 @@ async def _fetch_member_groups( Return a set with the google groups a given user/group is a member of, including nested groups if allowed. """ - credentials = credentials or self._setup_credentials(user_email_domain) + credentials = credentials or self._get_credentials(user_email_domain) checked_groups = checked_groups or set() processed_groups = processed_groups or set() From 876627f4a1484a36c560fc026eb2dd5a80dfeeb8 Mon Sep 17 00:00:00 2001 From: Jordan Bradford <36420801+jrdnbradford@users.noreply.github.com> Date: Fri, 4 Oct 2024 15:48:45 -0400 Subject: [PATCH 09/12] Clarify service creds vars --- oauthenticator/google.py | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/oauthenticator/google.py b/oauthenticator/google.py index 427f4a74..7b41b467 100644 --- a/oauthenticator/google.py +++ b/oauthenticator/google.py @@ -14,7 +14,7 @@ class GoogleOAuthenticator(OAuthenticator, GoogleOAuth2Mixin): user_auth_state_key = "google_user" - _credentials = None + _service_credentials = None @default("login_service") def _login_service_default(self): @@ -315,27 +315,29 @@ async def check_allowed(self, username, auth_model): # users should be explicitly allowed via config, otherwise they aren't return False - def _get_credentials(self, user_email_domain): + def _get_service_credentials(self, user_email_domain): """ Returns the stored credentials or fetches and stores new ones. Checks if the credentials are valid before returning them. Refreshes if necessary and stores the refreshed credentials. """ - if not self._credentials or not self._is_token_valid(): - self._credentials = self._setup_credentials(user_email_domain) + if not self._service_credentials or not self._is_token_valid(): + self._service_credentials = self._setup_service_credentials( + user_email_domain + ) - return self._credentials + return self._service_credentials def _is_token_valid(self): """ Checks if the stored token is valid. """ - if not self._credentials: + if not self._service_credentials: return False - if not self._credentials.token: + if not self._service_credentials.token: return False - if self._credentials.expired: + if self._service_credentials.expired: return False return True @@ -364,7 +366,7 @@ def _service_client_credentials(self, scopes, user_email_domain): return credentials - def _setup_credentials(self, user_email_domain): + def _setup_service_credentials(self, user_email_domain): """ Set up the oauth credentials for Google API. """ @@ -399,7 +401,7 @@ async def _fetch_member_groups( Return a set with the google groups a given user/group is a member of, including nested groups if allowed. """ - credentials = credentials or self._get_credentials(user_email_domain) + credentials = credentials or self._get_service_credentials(user_email_domain) checked_groups = checked_groups or set() processed_groups = processed_groups or set() From 846658cd4d731e784beb2efdab176c369e15341d Mon Sep 17 00:00:00 2001 From: Jordan Bradford <36420801+jrdnbradford@users.noreply.github.com> Date: Fri, 4 Oct 2024 16:27:38 -0400 Subject: [PATCH 10/12] Store multiple domain creds --- oauthenticator/google.py | 24 ++++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/oauthenticator/google.py b/oauthenticator/google.py index 7b41b467..8a7e4024 100644 --- a/oauthenticator/google.py +++ b/oauthenticator/google.py @@ -14,7 +14,7 @@ class GoogleOAuthenticator(OAuthenticator, GoogleOAuth2Mixin): user_auth_state_key = "google_user" - _service_credentials = None + _service_credentials = {} @default("login_service") def _login_service_default(self): @@ -322,22 +322,25 @@ def _get_service_credentials(self, user_email_domain): Checks if the credentials are valid before returning them. Refreshes if necessary and stores the refreshed credentials. """ - if not self._service_credentials or not self._is_token_valid(): - self._service_credentials = self._setup_service_credentials( - user_email_domain + if ( + user_email_domain not in self._service_credentials + or not self._is_token_valid(user_email_domain) + ): + self._service_credentials[user_email_domain] = ( + self._setup_service_credentials(user_email_domain) ) return self._service_credentials - def _is_token_valid(self): + def _is_token_valid(self, user_email_domain): """ Checks if the stored token is valid. """ - if not self._service_credentials: + if not self._service_credentials[user_email_domain]: return False - if not self._service_credentials.token: + if not self._service_credentials[user_email_domain].token: return False - if self._service_credentials.expired: + if self._service_credentials[user_email_domain].expired: return False return True @@ -385,7 +388,7 @@ def _setup_service_credentials(self, user_email_domain): request = requests.Request() credentials.refresh(request) - self.log.debug("Credentials refreshed") + self.log.debug(f"Credentials refreshed for {user_email_domain}") return credentials async def _fetch_member_groups( @@ -402,10 +405,11 @@ async def _fetch_member_groups( """ credentials = credentials or self._get_service_credentials(user_email_domain) + token = credentials[user_email_domain].token checked_groups = checked_groups or set() processed_groups = processed_groups or set() - headers = {'Authorization': f'Bearer {credentials.token}'} + headers = {'Authorization': f'Bearer {token}'} url = f'https://www.googleapis.com/admin/directory/v1/groups?userKey={member_email}' group_data = await self.httpfetch( url, headers=headers, label="fetching google groups" From a829ff7367073d38c66ad3452eed008a62cb1196 Mon Sep 17 00:00:00 2001 From: Jordan Bradford <36420801+jrdnbradford@users.noreply.github.com> Date: Sun, 6 Oct 2024 16:12:56 -0400 Subject: [PATCH 11/12] Update oauthenticator/google.py Co-authored-by: Simon Li --- oauthenticator/google.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/oauthenticator/google.py b/oauthenticator/google.py index 8a7e4024..87a9cf16 100644 --- a/oauthenticator/google.py +++ b/oauthenticator/google.py @@ -403,7 +403,8 @@ async def _fetch_member_groups( """ Return a set with the google groups a given user/group is a member of, including nested groups if allowed. """ - +# WARNING: There's a race condition here if multiple users login at the same time. +# This is currently ignored. credentials = credentials or self._get_service_credentials(user_email_domain) token = credentials[user_email_domain].token checked_groups = checked_groups or set() From 67cbbc6a42564337f4fafb32322f37bcb6fae885 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 6 Oct 2024 20:13:06 +0000 Subject: [PATCH 12/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- oauthenticator/google.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/oauthenticator/google.py b/oauthenticator/google.py index 87a9cf16..28b76d48 100644 --- a/oauthenticator/google.py +++ b/oauthenticator/google.py @@ -403,8 +403,8 @@ async def _fetch_member_groups( """ Return a set with the google groups a given user/group is a member of, including nested groups if allowed. """ -# WARNING: There's a race condition here if multiple users login at the same time. -# This is currently ignored. + # WARNING: There's a race condition here if multiple users login at the same time. + # This is currently ignored. credentials = credentials or self._get_service_credentials(user_email_domain) token = credentials[user_email_domain].token checked_groups = checked_groups or set()