iaptoolkit 0.3.0a2__tar.gz → 0.3.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (24) hide show
  1. {iaptoolkit-0.3.0a2 → iaptoolkit-0.3.2}/PKG-INFO +5 -9
  2. {iaptoolkit-0.3.0a2 → iaptoolkit-0.3.2}/README.md +2 -2
  3. {iaptoolkit-0.3.0a2 → iaptoolkit-0.3.2}/pyproject.toml +3 -11
  4. iaptoolkit-0.3.2/src/iaptoolkit/__init__.py +247 -0
  5. {iaptoolkit-0.3.0a2 → iaptoolkit-0.3.2}/src/iaptoolkit/exceptions.py +4 -23
  6. {iaptoolkit-0.3.0a2 → iaptoolkit-0.3.2}/src/iaptoolkit/headers.py +2 -6
  7. iaptoolkit-0.3.2/src/iaptoolkit/tokens/__init__.py +24 -0
  8. iaptoolkit-0.3.2/src/iaptoolkit/tokens/service_account.py +179 -0
  9. {iaptoolkit-0.3.0a2 → iaptoolkit-0.3.2}/src/iaptoolkit/tokens/structs.py +23 -7
  10. iaptoolkit-0.3.2/src/iaptoolkit/tokens/token_datastore.py +72 -0
  11. {iaptoolkit-0.3.0a2 → iaptoolkit-0.3.2}/src/iaptoolkit/utils/urls.py +3 -3
  12. iaptoolkit-0.3.0a2/src/iaptoolkit/__init__.py +0 -223
  13. iaptoolkit-0.3.0a2/src/iaptoolkit/tokens/base.py +0 -68
  14. iaptoolkit-0.3.0a2/src/iaptoolkit/tokens/oauth2/__init__.py +0 -110
  15. iaptoolkit-0.3.0a2/src/iaptoolkit/tokens/oauth2/datastore_oauth2.py +0 -39
  16. iaptoolkit-0.3.0a2/src/iaptoolkit/tokens/oauth2/gua.py +0 -11
  17. iaptoolkit-0.3.0a2/src/iaptoolkit/tokens/oidc/__init__.py +0 -149
  18. iaptoolkit-0.3.0a2/src/iaptoolkit/tokens/oidc/datastore_oidc.py +0 -35
  19. iaptoolkit-0.3.0a2/src/iaptoolkit/tokens/oidc/gsa.py +0 -11
  20. iaptoolkit-0.3.0a2/src/iaptoolkit/tokens/token_datastore.py +0 -63
  21. iaptoolkit-0.3.0a2/src/iaptoolkit/utils/__init__.py +0 -0
  22. {iaptoolkit-0.3.0a2 → iaptoolkit-0.3.2}/LICENSE +0 -0
  23. {iaptoolkit-0.3.0a2 → iaptoolkit-0.3.2}/src/iaptoolkit/constants.py +0 -0
  24. {iaptoolkit-0.3.0a2/src/iaptoolkit/tokens → iaptoolkit-0.3.2/src/iaptoolkit/utils}/__init__.py +0 -0
@@ -1,22 +1,18 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: iaptoolkit
3
- Version: 0.3.0a2
3
+ Version: 0.3.2
4
4
  Summary: Library of common utils for interacting with Identity-Aware Proxies
5
- Home-page: https://github.com/RAVoigt/iaptoolkit
6
- License: MIT
7
5
  Author: Rob Voigt
8
6
  Author-email: code@ravoigt.com
9
7
  Requires-Python: >=3.11,<4.0
10
- Classifier: License :: OSI Approved :: MIT License
11
8
  Classifier: Programming Language :: Python :: 3
12
9
  Classifier: Programming Language :: Python :: 3.11
13
10
  Classifier: Programming Language :: Python :: 3.12
11
+ Classifier: Programming Language :: Python :: 3.13
14
12
  Requires-Dist: google-auth (>=2.29.0,<3.0.0)
15
- Requires-Dist: kvcommon (>=0.1.4,<0.2.0)
16
- Requires-Dist: pytest (>=7.4.4,<8.0.0)
13
+ Requires-Dist: kvcommon[k8s] (>=0.2.3,<0.3.0)
17
14
  Requires-Dist: requests (>=2.31.0,<3.0.0)
18
15
  Requires-Dist: toml (>=0.10.2,<0.11.0)
19
- Project-URL: Repository, https://github.com/RAVoigt/iaptoolkit
20
16
  Description-Content-Type: text/markdown
21
17
 
22
18
  # IAP Toolkit
@@ -38,9 +34,9 @@ https://pypi.org/project/iaptoolkit/
38
34
  ```python
39
35
  import requests
40
36
 
41
- from iaptoolkit import IAPToolkit_OIDC
37
+ from iaptoolkit import IAPToolkit
42
38
 
43
- iaptk_oidc = IAPToolkit_OIDC(google_iap_client_id="EXAMPLE_ID_123456789ABCDEF")
39
+ iaptk = IAPToolkit(google_iap_client_id="EXAMPLE_ID_123456789ABCDEF")
44
40
  allowed_domains = ["example.com", ]
45
41
 
46
42
 
@@ -17,9 +17,9 @@ https://pypi.org/project/iaptoolkit/
17
17
  ```python
18
18
  import requests
19
19
 
20
- from iaptoolkit import IAPToolkit_OIDC
20
+ from iaptoolkit import IAPToolkit
21
21
 
22
- iaptk_oidc = IAPToolkit_OIDC(google_iap_client_id="EXAMPLE_ID_123456789ABCDEF")
22
+ iaptk = IAPToolkit(google_iap_client_id="EXAMPLE_ID_123456789ABCDEF")
23
23
  allowed_domains = ["example.com", ]
24
24
 
25
25
 
@@ -1,12 +1,9 @@
1
1
  [tool.poetry]
2
2
  name = "iaptoolkit"
3
- version = "0.3.0a2"
3
+ version = "0.3.2"
4
4
  description = "Library of common utils for interacting with Identity-Aware Proxies"
5
5
  authors = ["Rob Voigt <code@ravoigt.com>"]
6
6
  readme = "README.md"
7
- license = "MIT"
8
- repository = "https://github.com/RAVoigt/iaptoolkit"
9
- homepage = "https://github.com/RAVoigt/iaptoolkit"
10
7
 
11
8
  [build-system]
12
9
  requires = ["poetry-core>=1.0.0"]
@@ -19,7 +16,7 @@ Repository = "https://github.com/RAVoigt/iaptoolkit"
19
16
  # ================================
20
17
  # Tools etc.
21
18
  [tool.black]
22
- line-length = 100
19
+ line-length = 120
23
20
  target-version = ['py311']
24
21
  include = '\.pyi?$'
25
22
 
@@ -29,9 +26,8 @@ include = '\.pyi?$'
29
26
  python = "^3.11"
30
27
  google-auth = "^2.29.0"
31
28
  requests = "^2.31.0"
32
- pytest = "^7.4.4"
33
29
  toml = "^0.10.2"
34
- kvcommon = "^0.1.4"
30
+ kvcommon = {extras = ["k8s"], version = "^0.2.3"}
35
31
 
36
32
  [tool.poetry.dev-dependencies]
37
33
  black = "*"
@@ -43,7 +39,3 @@ pytest = "*"
43
39
  pytest-cov = "*"
44
40
  pytest-socket = "*"
45
41
  pyfakefs = "^5.3.2"
46
-
47
- # [tool.poetry.extras]
48
- # flask = ["flask"]
49
- # k8s = ["kubernetes"]
@@ -0,0 +1,247 @@
1
+ from __future__ import annotations
2
+
3
+ import logging
4
+
5
+ logging.getLogger(__name__).addHandler(logging.NullHandler())
6
+
7
+ import typing as t
8
+ from urllib.parse import ParseResult
9
+ from urllib.parse import urlparse
10
+
11
+ from kvcommon import logger
12
+
13
+ from iaptoolkit import headers
14
+ from iaptoolkit.exceptions import ServiceAccountTokenException
15
+ from iaptoolkit.tokens.service_account import ServiceAccount
16
+ from iaptoolkit.tokens.structs import ResultAddTokenHeader
17
+
18
+ from iaptoolkit.tokens.structs import TokenRefreshStruct
19
+ from iaptoolkit.tokens.structs import TokenStruct
20
+ from iaptoolkit.utils.urls import is_url_safe_for_token
21
+
22
+ LOG = logger.get_logger("iaptk")
23
+
24
+
25
+ class IAPToolkit:
26
+ """
27
+ Class to encapsulate client-specific vars and forward them to static functions
28
+ """
29
+
30
+ _GOOGLE_IAP_CLIENT_ID: str
31
+
32
+ def __init__(
33
+ self,
34
+ google_iap_client_id: str,
35
+ ) -> None:
36
+ self._GOOGLE_IAP_CLIENT_ID = google_iap_client_id
37
+
38
+ @staticmethod
39
+ def sanitize_request_headers(request_headers: dict) -> dict:
40
+ return headers.sanitize_request_headers(request_headers)
41
+
42
+ def get_token_oidc(self, bypass_cached: bool = False) -> TokenStruct:
43
+ try:
44
+ return ServiceAccount.get_token(iap_client_id=self._GOOGLE_IAP_CLIENT_ID, bypass_cached=bypass_cached)
45
+ except ServiceAccountTokenException as ex:
46
+ LOG.debug(ex)
47
+ raise
48
+
49
+ def get_token_oidc_str(self, bypass_cached: bool = False) -> str:
50
+ struct = self.get_token_oidc(bypass_cached=bypass_cached)
51
+ return struct.id_token
52
+
53
+ def get_token_oauth2(self, bypass_cached: bool = False) -> TokenRefreshStruct:
54
+ # TODO
55
+ raise NotImplementedError()
56
+
57
+ def get_token_oauth2_str(self, bypass_cached: bool = False) -> str:
58
+ struct = self.get_token_oauth2(bypass_cached=bypass_cached)
59
+ return struct.id_token
60
+
61
+ def get_token_and_add_to_headers(
62
+ self,
63
+ request_headers: dict,
64
+ use_oauth2: bool = False,
65
+ use_auth_header: bool = False,
66
+ bypass_cached: bool = False,
67
+ ) -> bool:
68
+ """
69
+ Retrieves an auth token and inserts it into the supplied request_headers dict.
70
+ request_headers is modified in-place
71
+
72
+ Params:
73
+ request_headers: dict of headers to insert into
74
+ use_oauth2: Use OAuth2.0 credentials and respective token, else use OIDC (default)
75
+ As a general guideline, OIDC is the assumed default approach for ServiceAccounts.
76
+ use_auth_header: If true, use the 'Authorization' header instead of 'Proxy-Authorization'
77
+
78
+ Returns:
79
+ True if token retrieved from cache, False if fresh from API
80
+
81
+
82
+ """
83
+ id_token = None
84
+ from_cache = False
85
+ if use_oauth2:
86
+ token_refresh_struct: TokenRefreshStruct = self.get_token_oauth2(bypass_cached=bypass_cached)
87
+ id_token = token_refresh_struct.id_token
88
+ from_cache = token_refresh_struct.from_cache
89
+ else:
90
+ token_struct: TokenStruct = self.get_token_oidc(bypass_cached=bypass_cached)
91
+ id_token = token_struct.id_token
92
+ from_cache = token_struct.from_cache
93
+
94
+ headers.add_token_to_request_headers(
95
+ request_headers=request_headers,
96
+ id_token=id_token,
97
+ use_auth_header=use_auth_header,
98
+ )
99
+
100
+ return from_cache
101
+
102
+ @staticmethod
103
+ def is_url_safe_for_token(
104
+ url: str | ParseResult,
105
+ valid_domains: t.Optional[t.List[str] | t.Set[str] | t.Tuple[str]] = None,
106
+ ):
107
+ if not isinstance(url, ParseResult):
108
+ url = urlparse(url)
109
+
110
+ return is_url_safe_for_token(url_parts=url, allowed_domains=valid_domains)
111
+
112
+ def check_url_and_add_token_header(
113
+ self,
114
+ url: str | ParseResult,
115
+ request_headers: dict,
116
+ valid_domains: t.List[str] | None = None,
117
+ use_oauth2: bool = False,
118
+ use_auth_header: bool = False,
119
+ bypass_cached: bool = False,
120
+ ) -> ResultAddTokenHeader:
121
+ """
122
+ Checks that the supplied URL is valid (i.e.; in valid_domains) and if so, retrieves a
123
+ token and adds it to request_headers.
124
+
125
+ i.e.; A convenience wrapper with logging for is_url_safe_for_token() and get_token_and_add_to_headers()
126
+
127
+ Params:
128
+ url: URL string or urllib.ParseResult to check for validity
129
+ request_headers: Dict of headers to insert into
130
+ valid_domains: List of domains to validate URL against
131
+ use_oauth2: Passed to get_token_and_add_to_headers() to determine if OAuth2.0 is used or OIDC (default)
132
+ """
133
+
134
+ if self.is_url_safe_for_token(url=url, valid_domains=valid_domains):
135
+ token_is_fresh = self.get_token_and_add_to_headers(
136
+ request_headers=request_headers,
137
+ use_oauth2=use_oauth2,
138
+ use_auth_header=use_auth_header,
139
+ bypass_cached=bypass_cached,
140
+ )
141
+ return ResultAddTokenHeader(token_added=True, token_is_fresh=token_is_fresh)
142
+ else:
143
+ LOG.warning(
144
+ "URL is not approved: %s - Token will not be added to headers. Valid domains are: %s",
145
+ url,
146
+ valid_domains,
147
+ )
148
+ return ResultAddTokenHeader(token_added=False, token_is_fresh=False)
149
+
150
+
151
+ class IAPToolkit_OIDC(IAPToolkit):
152
+ """
153
+ Convenience subclass of IAPToolkit for scenarios where OIDC will always be used, never OAuth2
154
+ """
155
+
156
+ def get_token_oauth2(self, *args, **kwargs):
157
+ raise NotImplementedError("Cannot call OAuth2 methods on OIDC-only instance of IAPToolkit.")
158
+
159
+ def get_token_oauth2_str(self, *args, **kwargs):
160
+ raise NotImplementedError("Cannot call OAuth2 methods on OIDC-only instance of IAPToolkit.")
161
+
162
+ def get_token_and_add_to_headers(
163
+ self,
164
+ request_headers: dict,
165
+ use_auth_header: bool = False,
166
+ use_oauth2: bool = False,
167
+ bypass_cached: bool = False,
168
+ ) -> bool:
169
+ return super().get_token_and_add_to_headers(
170
+ request_headers=request_headers,
171
+ use_oauth2=use_oauth2,
172
+ use_auth_header=use_auth_header,
173
+ bypass_cached=bypass_cached,
174
+ )
175
+
176
+ def check_url_and_add_token_header(
177
+ self,
178
+ url: str | ParseResult,
179
+ request_headers: dict,
180
+ valid_domains: t.List[str] | None = None,
181
+ use_auth_header: bool = False,
182
+ bypass_cached: bool = False,
183
+ ) -> ResultAddTokenHeader:
184
+ return super().check_url_and_add_token_header(
185
+ url,
186
+ request_headers=request_headers,
187
+ valid_domains=valid_domains,
188
+ use_oauth2=False,
189
+ use_auth_header=use_auth_header,
190
+ bypass_cached=bypass_cached,
191
+ )
192
+
193
+
194
+ class IAPToolkit_OAuth2(IAPToolkit):
195
+ """
196
+ Convenience subclass of IAPToolkit for scenarios where OAuth2 will always be used, never OIDC
197
+ """
198
+
199
+ _GOOGLE_CLIENT_ID: str
200
+ _GOOGLE_CLIENT_SECRET: str
201
+
202
+ def __init__(
203
+ self,
204
+ google_iap_client_id: str,
205
+ google_client_id: str,
206
+ google_client_secret: str,
207
+ ) -> None:
208
+ super().__init__(google_iap_client_id=google_iap_client_id)
209
+ self._GOOGLE_CLIENT_ID = google_client_id
210
+ self._GOOGLE_CLIENT_SECRET = google_client_secret
211
+
212
+ def get_token_oidc(self, *args, **kwargs):
213
+ raise NotImplementedError("Cannot call OIDC methods on OAuth2-only instance of IAPToolkit.")
214
+
215
+ def get_token_oidc_str(self, *args, **kwargs):
216
+ raise NotImplementedError("Cannot call OIDC methods on OAuth2-only instance of IAPToolkit.")
217
+
218
+ def get_token_and_add_to_headers(
219
+ self,
220
+ request_headers: dict,
221
+ use_auth_header: bool = False,
222
+ use_oauth2: bool = True,
223
+ bypass_cached: bool = False,
224
+ ) -> bool:
225
+ return super().get_token_and_add_to_headers(
226
+ request_headers=request_headers,
227
+ use_oauth2=use_oauth2,
228
+ use_auth_header=use_auth_header,
229
+ bypass_cached=bypass_cached,
230
+ )
231
+
232
+ def check_url_and_add_token_header(
233
+ self,
234
+ url: str | ParseResult,
235
+ request_headers: dict,
236
+ valid_domains: t.List[str] | None = None,
237
+ use_auth_header: bool = False,
238
+ bypass_cached: bool = False,
239
+ ) -> ResultAddTokenHeader:
240
+ return super().check_url_and_add_token_header(
241
+ url=url,
242
+ request_headers=request_headers,
243
+ valid_domains=valid_domains,
244
+ use_oauth2=True,
245
+ use_auth_header=use_auth_header,
246
+ bypass_cached=bypass_cached,
247
+ )
@@ -14,10 +14,6 @@ class IAPBadResponse(IAPToolkitBaseException):
14
14
  pass
15
15
 
16
16
 
17
- class InvalidDomain(IAPToolkitBaseException):
18
- pass
19
-
20
-
21
17
  class TokenException(IAPToolkitBaseException):
22
18
  pass
23
19
 
@@ -26,13 +22,8 @@ class TokenStorageException(TokenException):
26
22
  pass
27
23
 
28
24
 
29
- class GoogleTokenException(TokenException):
30
- """
31
- Wrapper for exceptions from Google's auth lib
32
- """
33
- def __init__(
34
- self, message: str, google_exception: t.Union[DefaultCredentialsError, RefreshError] | None
35
- ):
25
+ class ServiceAccountTokenException(TokenException):
26
+ def __init__(self, message: str, google_exception: t.Union[DefaultCredentialsError, RefreshError] | None):
36
27
  self.google_exception = google_exception
37
28
  credentials_env_var_value = os.environ.get(GOOGLE_CREDENTIALS_FILE_PATH)
38
29
  metadata_server_attempted = not credentials_env_var_value
@@ -49,24 +40,14 @@ class GoogleTokenException(TokenException):
49
40
  def retryable(self):
50
41
  return self.google_exception and self.google_exception._retryable
51
42
 
52
- # SA / OIDC
53
-
54
- class ServiceAccountTokenException(GoogleTokenException):
55
- pass
56
-
57
-
58
- class ServiceAccountTokenFailedRefresh(ServiceAccountTokenException):
59
- pass
60
-
61
43
 
62
44
  class ServiceAccountNoDefaultCredentials(ServiceAccountTokenException):
63
45
  pass
64
46
 
65
- # OAuth2 / User
66
47
 
67
- class OAuth2RefreshFromAuthCodeFailed(TokenException):
48
+ class ServiceAccountTokenFailedRefresh(ServiceAccountTokenException):
68
49
  pass
69
50
 
70
51
 
71
- class OAuth2IDTokenFromRefreshFailed(TokenException):
52
+ class InvalidDomain(IAPToolkitBaseException):
72
53
  pass
@@ -23,9 +23,7 @@ def _sanitize_request_header(headers_dict: dict, header_key: str):
23
23
  def sanitize_request_headers(headers: dict) -> dict:
24
24
  """
25
25
  Sanitizes a headers dict to remove sensitive strings for logging purposes.
26
-
27
- Returns:
28
- A COPY of the dict with sensitive k/v pairs replaced. Does NOT modify in-place/by-reference.
26
+ Returns A COPY of the dict with sensitive k/v pairs replaced. Does NOT modify in-place/by-reference.
29
27
  """
30
28
  log_safe_headers = headers.copy()
31
29
 
@@ -36,9 +34,7 @@ def sanitize_request_headers(headers: dict) -> dict:
36
34
  return log_safe_headers
37
35
 
38
36
 
39
- def add_token_to_request_headers(
40
- request_headers: dict, id_token: str, use_auth_header: bool = False
41
- ) -> dict:
37
+ def add_token_to_request_headers(request_headers: dict, id_token: str, use_auth_header: bool = False) -> dict:
42
38
  """
43
39
  Adds Bearer token to headers dict. Modifies dict in-place.
44
40
  Returns True if added token is a fresh one, or False if token is from cache
@@ -0,0 +1,24 @@
1
+ from kvcommon import logger
2
+
3
+ from iaptoolkit.exceptions import ServiceAccountTokenException
4
+ from iaptoolkit.exceptions import TokenStorageException
5
+ from iaptoolkit.exceptions import TokenException
6
+
7
+ from .structs import TokenStruct
8
+ from .structs import TokenRefreshStruct
9
+
10
+ # from .structs import TokenStructOAuth2 # TODO: OAuth2
11
+ # from .oauth2 import get_token_for_oauth2 # TODO: OAuth2
12
+ # from .service_account import ServiceAccount
13
+ from .service_account import GoogleServiceAccount
14
+
15
+ LOG = logger.get_logger("iaptk")
16
+
17
+
18
+ __all__ = [
19
+ "TokenStruct",
20
+ "TokenRefreshStruct",
21
+ # "TokenStructOAuth2", # TODO: OAuth2
22
+ "TokenException",
23
+ "TokenStorageException",
24
+ ]
@@ -0,0 +1,179 @@
1
+ import datetime
2
+ import typing as t
3
+
4
+ from google.auth.compute_engine import IDTokenCredentials as GoogleIDTokenCredentials
5
+ from google.auth.exceptions import DefaultCredentialsError as GoogleDefaultCredentialsError
6
+ from google.auth.exceptions import RefreshError as GoogleRefreshError
7
+ from google.auth.transport.requests import Request as GoogleRequest
8
+ from google.oauth2 import id_token as google_id_token_lib
9
+
10
+ from kvcommon import logger
11
+
12
+ # TODO: Don't hardcode the association between OIDC/SA and dict-datastore
13
+ from iaptoolkit.tokens.token_datastore import datastore
14
+ from iaptoolkit.exceptions import ServiceAccountTokenException
15
+ from iaptoolkit.exceptions import ServiceAccountTokenFailedRefresh
16
+ from iaptoolkit.exceptions import ServiceAccountNoDefaultCredentials
17
+ from iaptoolkit.exceptions import TokenException
18
+ from iaptoolkit.exceptions import TokenStorageException
19
+
20
+ from .structs import TokenStruct
21
+ from .structs import TokenRefreshStruct
22
+
23
+
24
+ LOG = logger.get_logger("iaptk")
25
+ MAX_RECURSE = 3
26
+
27
+
28
+ class ServiceAccount(object):
29
+ """Base class for interacting with service accounts and OIDC tokens for IAP"""
30
+
31
+ # TODO: This is a static namespace for SA functions. Turn it into a per-iap-client-id client
32
+ # TODO: Move Google-specific logic to GoogleServiceAccount
33
+
34
+ @staticmethod
35
+ def _store_token(iap_client_id: str, id_token: str, token_expiry: datetime.datetime):
36
+ try:
37
+ datastore.store_service_account_token(iap_client_id, id_token, token_expiry)
38
+ except Exception as ex: # Err on the side of not letting token-caching break requests.
39
+ raise TokenStorageException(f"Exception when trying to store token. exception={ex}")
40
+
41
+ @staticmethod
42
+ def get_stored_token(iap_client_id: str) -> t.Optional[TokenStruct]:
43
+ try:
44
+ token_dict = datastore.get_stored_service_account_token(iap_client_id)
45
+ if not token_dict or not token_dict.get("id_token", None) or not token_dict.get("token_expiry", None):
46
+ LOG.debug("No stored service account token for current iap_client_id")
47
+ return
48
+
49
+ id_token_from_dict: str = token_dict.get("id_token", "")
50
+ token_expiry_from_dict: str = token_dict.get("token_expiry", "")
51
+
52
+ token_expiry = ""
53
+ try:
54
+ token_expiry = datetime.datetime.fromisoformat(token_expiry_from_dict)
55
+ except (ValueError, TypeError) as ex:
56
+ LOG.debug("Invalid token expiry for stored token - Could not parse from ISO format to datetime.")
57
+ return
58
+
59
+ token_struct = TokenStruct(id_token=id_token_from_dict, expiry=token_expiry, from_cache=True)
60
+ if not token_struct.valid:
61
+ LOG.debug("Stored service account token for current iap_client_id is INVALID")
62
+ return
63
+ if token_struct.expired:
64
+ LOG.debug("Stored service account token for current iap_client_id has EXPIRED")
65
+ return
66
+
67
+ return token_struct
68
+
69
+ except Exception as ex:
70
+ # Err on the side of not letting token-caching break requests, hence blanket except
71
+ raise TokenStorageException(f"Exception when trying to retrieve stored token. exception={ex}")
72
+
73
+ @staticmethod
74
+ def _get_fresh_credentials(iap_client_id: str) -> GoogleIDTokenCredentials:
75
+
76
+ try:
77
+ request = GoogleRequest()
78
+ credentials: GoogleIDTokenCredentials = google_id_token_lib.fetch_id_token_credentials(
79
+ iap_client_id, request
80
+ ) # type: ignore
81
+ credentials.refresh(request)
82
+
83
+ except GoogleDefaultCredentialsError as ex:
84
+ # The exceptions that google's libs raise in this case are somewhat vague; wrap them.
85
+ raise ServiceAccountNoDefaultCredentials(
86
+ message="Failed to get ServiceAccount token: Lacking default credentials.",
87
+ google_exception=ex,
88
+ )
89
+ except GoogleRefreshError as ex:
90
+ # Likely attempting to get a token for a service account in an environment that
91
+ # doesn't have one attached.
92
+ raise ServiceAccountTokenFailedRefresh(
93
+ message="Failed to get ServiceAccount token: Refreshing token failed.",
94
+ google_exception=ex,
95
+ )
96
+ return credentials
97
+
98
+ @staticmethod
99
+ def _get_fresh_token(iap_client_id: str) -> TokenStruct:
100
+ google_credentials = ServiceAccount._get_fresh_credentials(iap_client_id)
101
+ id_token: str = str(google_credentials.token)
102
+ if not id_token:
103
+ raise TokenException("Invalid [empty] token retrieved for Service Account.")
104
+
105
+ # Google lib uses deprecated 'utcfromtimestamp' func as of v2.29.x
106
+ # e.g.: datetime.datetime.utcfromtimestamp(payload["exp"])
107
+ # This creates a TZ-naive datetime in UTC from a POSIX timestamp.
108
+ # Python datetimes assume local TZ, and we want to explicitly only work in UTC here.
109
+ token_expiry = google_credentials.expiry.replace(tzinfo=datetime.timezone.utc)
110
+
111
+ return TokenStruct(id_token=id_token, expiry=token_expiry, from_cache=False)
112
+
113
+ @staticmethod
114
+ def get_token(iap_client_id: str, bypass_cached: bool = False, attempts: int = 0) -> TokenStruct:
115
+ """Retrieves an OIDC token for the current environment either from environment variable or from
116
+ metadata service.
117
+
118
+ 1. If the environment variable ``GOOGLE_APPLICATION_CREDENTIALS`` is set
119
+ to the path of a valid service account JSON file, then ID token is
120
+ acquired using this service account credentials.
121
+ 2. If the application is running in Compute Engine, App Engine or Cloud Run,
122
+ then the ID token is obtained from the metadata server.
123
+
124
+ Args:
125
+ iap_client_id: The client ID used by IAP. Can be thought of as JWT audience.
126
+
127
+ Returns:
128
+ An OIDC token for use in connecting through IAP.
129
+
130
+ Raises:
131
+ :class:`ServiceAccountTokenException` if a token could not be retrieved due to either
132
+ missing credentials from env-var/JSON or inability to talk to metadata server.
133
+ """
134
+
135
+ use_cache = not bypass_cached
136
+
137
+ try:
138
+ token_struct: TokenStruct | None = None
139
+
140
+ if use_cache:
141
+ token_struct = ServiceAccount.get_stored_token(iap_client_id)
142
+
143
+ if not token_struct:
144
+ token_struct = ServiceAccount._get_fresh_token(iap_client_id)
145
+ if use_cache:
146
+ ServiceAccount._store_token(iap_client_id, token_struct.id_token, token_struct.expiry)
147
+
148
+ return token_struct
149
+
150
+ except ServiceAccountTokenException as ex:
151
+ attempts += 1
152
+ if attempts > MAX_RECURSE or not ex.retryable:
153
+ raise
154
+ return ServiceAccount.get_token(iap_client_id, bypass_cached=False, attempts=attempts)
155
+
156
+ except TokenStorageException as ex:
157
+ if attempts > 1:
158
+ raise
159
+ attempts += 1
160
+ # Try again without involving the cache
161
+ return ServiceAccount.get_token(iap_client_id, bypass_cached=True, attempts=attempts)
162
+
163
+
164
+ class GoogleServiceAccount(ServiceAccount):
165
+ """For interacting with Google service accounts and OIDC tokens for Google IAP"""
166
+
167
+ def __init__(self, iap_client_id: str) -> None:
168
+ if not iap_client_id or not isinstance(iap_client_id, str):
169
+ raise ServiceAccountTokenException("Invalid iap_client_id for GoogleServiceAccount", google_exception=None)
170
+ self._iap_client_id = iap_client_id
171
+ super().__init__()
172
+
173
+ def get_stored_token(self) -> t.Optional[TokenStruct]:
174
+ return ServiceAccount.get_stored_token(self._iap_client_id)
175
+
176
+ def get_token(self, bypass_cached: bool = False, attempts: int = 0) -> TokenStruct:
177
+ return ServiceAccount.get_token(
178
+ iap_client_id=self._iap_client_id, bypass_cached=bypass_cached, attempts=attempts
179
+ )
@@ -8,11 +8,18 @@ from kvcommon import logger
8
8
  LOG = logger.get_logger("iaptk")
9
9
 
10
10
 
11
+ def validate_token(token: str | None) -> bool:
12
+ if not isinstance(token, str) or token.strip() == "":
13
+ return False
14
+
15
+ return True
16
+
17
+
11
18
  @dataclass(kw_only=True)
12
19
  class TokenStruct:
13
20
  id_token: str
14
21
  expiry: datetime.datetime
15
- token_is_new: bool = True
22
+ from_cache: bool = False
16
23
 
17
24
  @property
18
25
  def expired(self):
@@ -31,22 +38,31 @@ class TokenStruct:
31
38
  LOG.error("Exception when checking token expiry. exception=%s", ex)
32
39
  return True
33
40
 
41
+ @property
42
+ def valid(self):
43
+ return validate_token(self.id_token)
44
+
34
45
 
35
46
  @dataclass(kw_only=True)
36
47
  class TokenRefreshStruct:
37
- """
38
- """
39
48
  id_token: str
40
- token_is_new: bool = True
49
+ from_cache: bool = False
41
50
 
51
+ @property
52
+ def valid(self):
53
+ return validate_token(self.id_token)
42
54
 
43
55
  @dataclass(kw_only=True)
44
- class TokenStructOAuth2:
56
+ class TokenStructOAuth2(TokenStruct):
45
57
  refresh_token: str
46
- token_is_new: bool = False
58
+ from_cache: bool = False
59
+
60
+ @property
61
+ def valid(self):
62
+ return validate_token(self.refresh_token)
47
63
 
48
64
 
49
65
  @dataclass(kw_only=True)
50
66
  class ResultAddTokenHeader:
51
67
  token_added: bool
52
- token_is_new: bool
68
+ token_is_fresh: bool