8000 feat: Introduce the functionality to override token_uri in credential… · TJB-1/google-auth-library-python@73bc7e9 · GitHub
[go: up one dir, main page]
Skip to content

Commit 73bc7e9

Browse files
sai-sunder-ssai-sunder-s
authored
feat: Introduce the functionality to override token_uri in credentials (googleapis#1159)
* feat: Introduce the functionality to override token_uri in credentials * update rt
1 parent 75326e3 commit 73bc7e9

File tree

10 files changed

+216
-4
lines changed

10 files changed

+216
-4
lines changed

‎google/auth/compute_engine/credentials.py

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -154,7 +154,11 @@ def with_scopes(self, scopes, default_scopes=None):
154154
_DEFAULT_TOKEN_URI = "https://www.googleapis.com/oauth2/v4/token"
155155

156156

157-
class IDTokenCredentials(credentials.CredentialsWithQuotaProject, credentials.Signing):
157+
class IDTokenCredentials(
158+
credentials.CredentialsWithQuotaProject,
159+
credentials.Signing,
160+
credentials.CredentialsWithTokenUri,
161+
):
158162
"""Open ID Connect ID Token-based service account credentials.
159163
160164
These credentials relies on the default service account of a GCE instance.
@@ -302,6 +306,27 @@ def with_quota_project(self, quota_project_id):
302306
quota_project_id=quota_project_id,
303307
)
304308

309+
@_helpers.copy_docstring(credentials.CredentialsWithTokenUri)
310+
def with_token_uri(self, token_uri):
311+
312+
# since the signer is already instantiated,
313+
# the request is not needed
314+
if self._use_metadata_identity_endpoint:
315+
raise ValueError(
316+
"If use_metadata_identity_endpoint is set, token_uri" " must not be set"
317+
)
318+
else:
319+
return self.__class__(
320+
None,
321+
service_account_email=self._service_account_email,
322+
token_uri=token_uri,
323+
target_audience=self._target_audience,
324+
additional_claims=self._additional_claims.copy(),
325+
signer=self.signer,
326+
use_metadata_identity_endpoint=False,
327+
quota_project_id=self.quota_project_id,
328+
)
329+
305330
def _make_authorization_grant_assertion(self):
306331
"""Create the OAuth 2.0 assertion.
307332
This assertion is used during the OAuth 2.0 grant to acquire an

‎google/auth/credentials.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,21 @@ def with_quota_project(self, quota_project_id):
150150
raise NotImplementedError("This credential does not support quota project.")
151151

152152

153+
class CredentialsWithTokenUri(Credentials):
154+
"""Abstract base for credentials supporting ``with_token_uri`` factory"""
155+
156+
def with_token_uri(self, token_uri):
157+
"""Returns a copy of these credentials with a modified token uri.
158+
159+
Args:
160+
token_uri (str): The uri to use for fetching/exchanging tokens
161+
162+
Returns:
163+
google.oauth2.credentials.Credentials: A new credentials instance.
164+
"""
165+
raise NotImplementedError("This credential does not use token uri.")
166+
167+
153168
class AnonymousCredentials(Credentials):
154169
"""Credentials that do not provide any authentication information.
155170

‎google/auth/external_account.py

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,11 @@
5555

5656

5757
@six.add_metaclass(abc.ABCMeta)
58-
class Credentials(credentials.Scoped, credentials.CredentialsWithQuotaProject):
58+
class Credentials(
59+
credentials.Scoped,
60+
credentials.CredentialsWithQuotaProject,
61+
credentials.CredentialsWithTokenUri,
62+
):
5963
"""Base class for all external account credentials.
6064
6165
This is used to instantiate Credentials for exchanging external account
@@ -382,6 +386,26 @@ def with_quota_project(self, quota_project_id):
382386
d.pop("workforce_pool_user_project")
383387
return self.__class__(**d)
384388

389+
@_helpers.copy_docstring(credentials.CredentialsWithTokenUri)
390+
def with_token_uri(self, token_uri):
391+
d = dict(
392+
audience=self._audience,
393+
subject_token_type=self._subject_token_type,
394+
token_url=token_uri,
395+
credential_source=self._credential_source,
396+
service_account_impersonation_url=self._service_account_impersonation_url,
397+
service_account_impersonation_options=self._service_account_impersonation_options,
398+
client_id=self._client_id,
399+
client_secret=self._client_secret,
400+
quota_project_id=self._quota_project_id,
401+
scopes=self._scopes,
402+
default_scopes=self._default_scopes,
403+
workforce_pool_user_project=self._workforce_pool_user_project,
404+
)
405+
if not self.is_workforce_pool:
406+
d.pop("workforce_pool_user_project")
407+
return self.__class__(**d)
408+
385409
def _initialize_impersonated_credentials(self):
386410
"""Generates an impersonated credentials.
387411

‎google/oauth2/credentials.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -254,6 +254,23 @@ def with_quota_project(self, quota_project_id):
254254
enable_reauth_refresh=self._enable_reauth_refresh,
255255
)
256256

257+
@_helpers.copy_docstring(credentials.CredentialsWithTokenUri)
258+
def with_token_uri(self, token_uri):
259+
260+
return self.__class__(
261+
self.token,
262+
refresh_token=self.refresh_token,
263+
id_token=self.id_token,
264+
token_uri=token_uri,
265+
client_id=self.client_id,
266+
client_secret=self.client_secret,
267+
scopes=self.scopes,
268+
default_scopes=self.default_scopes,
269+
quota_project_id=self.quota_project_id,
270+
rapt_token=self.rapt_token,
271+
enable_reauth_refresh=self._enable_reauth_refresh,
272+
)
273+
257274
@_helpers.copy_docstring(credentials.Credentials)
258275
def refresh(self, request):
259276
scopes = self._scopes if self._scopes is not None else self._default_scopes

‎google/oauth2/service_account.py

Lines changed: 36 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,10 @@
8484

8585

8686
class Credentials(
87-
credentials.Signing, credentials.Scoped, credentials.CredentialsWithQuotaProject
87+
credentials.Signing,
88+
credentials.Scoped,
89+
credentials.CredentialsWithQuotaProject,
90+
credentials.CredentialsWithTokenUri,
8891
):
8992
"""Service account credentials
9093
@@ -364,6 +367,22 @@ def with_quota_project(self, quota_project_id):
364367
always_use_jwt_access=self._always_use_jwt_access,
365368
)
366369

370+
@_helpers.copy_docstring(credentials.CredentialsWithTokenUri)
371+
def with_token_uri(self, token_uri):
372+
373+
return self.__class__(
374+
self._signer,
375+
service_account_email=self._service_account_email,
376+
default_scopes=self._default_scopes,
377+
scopes=self._scopes,
378+
token_uri=token_uri,
379+
subject=self._subject,
380+
project_id=self._project_id,
381+
quota_project_id=self._quota_project_id,
382+
additional_claims=self._additional_claims.copy(),
383+
always_use_jwt_access=self._always_use_jwt_access,
384+
)
385+
367386
def _make_authorization_grant_assertion(self):
368387
"""Create the OAuth 2.0 assertion.
369388
@@ -455,7 +474,11 @@ def signer_email(self):
455474
return self._service_account_email
456475

457476

458-
class IDTokenCredentials(credentials.Signing, credentials.CredentialsWithQuotaProject):
477+
class IDTokenCredentials(
478+
credentials.Signing,
479+
credentials.CredentialsWithQuotaProject,
480+
credentials.CredentialsWithTokenUri,
481+
):
459482
"""Open ID Connect ID Token-based service account credentials.
460483
461484
These credentials are largely similar to :class:`.Credentials`, but instead
@@ -627,6 +650,17 @@ def with_quota_project(self, quota_project_id):
627650
quota_project_id=quota_project_id,
628651
)
629652

653+
@_helpers.copy_docstring(credentials.CredentialsWithTokenUri)
654+
def with_token_uri(self, token_uri):
655+
return self.__class__(
656+
self._signer,
657+
service_account_email=self._service_account_email,
658+
token_uri=token_uri,
659+
target_audience=self._target_audience,
660+
additional_claims=self._additional_claims.copy(),
661+
quota_project_id=self._quota_project_id,
662+
)
663+
630664
def _make_authorization_grant_assertion(self):
631665
"""Create the OAuth 2.0 assertion.
632666

‎system_tests/secrets.tar.enc

0 Bytes
Binary file not shown.

‎tests/compute_engine/test_credentials.py

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -483,6 +483,50 @@ def test_with_quota_project(self, sign, get, utcnow):
483483
# Check that the signer have been initialized with a Request object
484484
assert isinstance(self.credentials._signer._request, transport.Request)
485485

486+
@mock.patch(
487+
"google.auth._helpers.utcnow",
488+
return_value=datetime.datetime.utcfromtimestamp(0),
489+
)
490+
@mock.patch("google.auth.compute_engine._metadata.get", autospec=True)
491+
@mock.patch("google.auth.iam.Signer.sign", autospec=True)
492+
def test_with_token_uri(self, sign, get, utcnow):
493+
get.side_effect = [
494+
{"email": "service-account@example.com", "scopes": ["one", "two"]}
495+
]
496+
sign.side_effect = [b"signature"]
497+
498+
request = mock.create_autospec(transport.Request, instance=True)
499+
self.credentials = credentials.IDTokenCredentials(
500+
request=request,
501+
target_audience="https://audience.com",
502+
token_uri="http://xyz.com",
503+
)
504+
assert self.credentials._token_uri == "http://xyz.com"
505+
creds_with_token_uri = self.credentials.with_token_uri("http://abc.com")
506+
assert creds_with_token_uri._token_uri == "http://abc.com"
507+
508+
@mock.patch(
509+
"google.auth._helpers.utcnow",
510+
return_value=datetime.datetime.utcfromtimestamp(0),
511+
)
512+
@mock.patch("google.auth.compute_engine._metadata.get", autospec=True)
513+
@mock.patch("google.auth.iam.Signer.sign", autospec=True)
514+
def test_with_token_uri_exception(self, sign, get, utcnow):
515+
get.side_effect = [
516+
{"email": "service-account@example.com", "scopes": ["one", "two"]}
517+
]
518+
sign.side_effect = [b"signature"]
519+
520+
request = mock.create_autospec(transport.Request, instance=True)
521+
self.credentials = credentials.IDTokenCredentials(
522+
request=request,
523+
target_audience="https://audience.com",
524+
use_metadata_identity_endpoint=True,
525+
)
526+
assert self.credentials._token_uri is None
527+
with pytest.raises(ValueError):
528+
self.credentials.with_token_uri("http://abc.com")
529+
486530
@responses.activate
487531
def test_with_quota_project_integration(self):
488532
""" Test that it is possible to refresh credentials

‎tests/oauth2/test_credentials.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -701,6 +701,18 @@ def test_with_quota_project(self):
701701
creds.apply(headers)
702702
assert "x-goog-user-project" in headers
703703

704+
def test_with_token_uri(self):
705+
info = AUTH_USER_INFO.copy()
706+
707+
creds = credentials.Credentials.from_authorized_user_info(info)
708+
new_token_uri = "https://oauth2-eu.googleapis.com/token"
709+
710+
assert creds._token_uri == credentials._GOOGLE_OAUTH2_TOKEN_ENDPOINT
711+
712+
creds_with_new_token_uri = creds.with_token_uri(new_token_uri)
713+
714+
assert creds_with_new_token_uri._token_uri == new_token_uri
715+
704716
def test_from_authorized_user_info(self):
705717
info = AUTH_USER_INFO.copy()
706718

‎tests/oauth2/test_service_account.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -155,6 +155,13 @@ def test_with_quota_project(self):
155155
new_credentials.apply(hdrs, token="tok")
156156
assert "x-goog-user-project" in hdrs
157157

158+
def test_with_token_uri(self):
159+
credentials = self.make_credentials()
160+
new_token_uri = "https://example2.com/oauth2/token"
161+
assert credentials._token_uri == self.TOKEN_URI
162+
creds_with_new_token_uri = credentials.with_token_uri(new_token_uri)
163+
assert creds_with_new_token_uri._token_uri == new_token_uri
164+
158165
def test__with_always_use_jwt_access(self):
159166
credentials = self.make_credentials()
160167
assert not credentials._always_use_jwt_access
@@ -464,6 +471,13 @@ def test_with_quota_project(self):
464471
new_credentials = credentials.with_quota_project("project-foo")
465472
assert new_credentials._quota_project_id == "project-foo"
466473

474+
def test_with_token_uri(self):
475+
credentials = self.make_credentials()
476+
new_token_uri = "https://example2.com/oauth2/token"
477+
assert credentials._token_uri == self.TOKEN_URI
478+
creds_with_new_token_uri = credentials.with_token_uri(new_token_uri)
479+
assert creds_with_new_token_uri._token_uri == new_token_uri
480+
467481
def test__make_authorization_grant_assertion(self):
468482
credentials = self.make_credentials()
469483
token = credentials._make_authorization_grant_assertion()

‎tests/test_external_account.py

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -542,6 +542,33 @@ def test_with_scopes_full_options_propagated(self):
542542
workforce_pool_user_project=None,
543543
)
544544

545+
def test_with_token_uri(self):
546+
credentials = self.make_credentials()
547+
new_token_uri = "https://eu-sts.googleapis.com/v1/token"
548+
549+
assert credentials._token_url == self.TOKEN_URL
550+
551+
creds_with_new_token_uri = credentials.with_token_uri(new_token_uri)
552+
553+
assert creds_with_new_token_uri._token_url == new_token_uri
554+
555+
def test_with_token_uri_workforce_pool(self):
556+
credentials = self.make_workforce_pool_credentials(
557+
workforce_pool_user_project=self.WORKFORCE_POOL_USER_PROJECT
558+
)
559+
560+
new_token_uri = "https://eu-sts.googleapis.com/v1/token"
561+
562+
assert credentials._token_url == self.TOKEN_URL
563+
564+
creds_with_new_token_uri = credentials.with_token_uri(new_token_uri)
565+
566+
assert creds_with_new_token_uri._token_url == new_token_uri
567+
assert (
568+
creds_with_new_token_uri.info.get("workforce_pool_user_project")
569+
== self.WORKFORCE_POOL_USER_PROJECT
570+
)
571+
545572
def test_with_quota_project(self):
546573
credentials = self.make_credentials()
547574

0 commit comments

Comments
 (0)
0