iaptoolkit 0.2.5__tar.gz → 0.3.0__tar.gz

{iaptoolkit-0.2.5 → iaptoolkit-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: iaptoolkit
-Version: 0.2.5
+Version: 0.3.0
 Summary: Library of common utils for interacting with Identity-Aware Proxies
 Author: Rob Voigt
 Author-email: code@ravoigt.com
@@ -8,6 +8,7 @@ Requires-Python: >=3.11,<4.0
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Requires-Dist: google-auth (>=2.29.0,<3.0.0)
 Requires-Dist: kvcommon (>=0.1.3,<0.2.0)
 Requires-Dist: pytest (>=7.4.4,<8.0.0)

{iaptoolkit-0.2.5 → iaptoolkit-0.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "iaptoolkit"
-version = "0.2.5"
+version = "0.3.0"
 description = "Library of common utils for interacting with Identity-Aware Proxies"
 authors = ["Rob Voigt <code@ravoigt.com>"]
 readme = "README.md"
@@ -16,7 +16,7 @@ Repository = "https://github.com/RAVoigt/iaptoolkit"
 # ================================
 # Tools etc.
 [tool.black]
-line-length = 100
+line-length = 120
 target-version = ['py311']
 include = '\.pyi?$'
 

{iaptoolkit-0.2.5 → iaptoolkit-0.3.0}/src/iaptoolkit/__init__.py

@@ -16,6 +16,7 @@ from iaptoolkit.tokens.service_account import ServiceAccount
 from iaptoolkit.tokens.structs import ResultAddTokenHeader
 
 from iaptoolkit.tokens.structs import TokenRefreshStruct
+from iaptoolkit.tokens.structs import TokenStruct
 from iaptoolkit.utils.urls import is_url_safe_for_token
 
 LOG = logger.get_logger("iaptk")
@@ -38,11 +39,9 @@ class IAPToolkit:
     def sanitize_request_headers(request_headers: dict) -> dict:
         return headers.sanitize_request_headers(request_headers)
 
-    def get_token_oidc(self, bypass_cached: bool = False) -> TokenRefreshStruct:
+    def get_token_oidc(self, bypass_cached: bool = False) -> TokenStruct:
         try:
-            return ServiceAccount.get_token(
-                iap_client_id=self._GOOGLE_IAP_CLIENT_ID, bypass_cached=bypass_cached
-            )
+            return ServiceAccount.get_token(iap_client_id=self._GOOGLE_IAP_CLIENT_ID, bypass_cached=bypass_cached)
         except ServiceAccountTokenException as ex:
             LOG.debug(ex)
             raise
@@ -60,7 +59,11 @@ class IAPToolkit:
         return struct.id_token
 
     def get_token_and_add_to_headers(
-        self, request_headers: dict, use_oauth2: bool = False, use_auth_header: bool = False
+        self,
+        request_headers: dict,
+        use_oauth2: bool = False,
+        use_auth_header: bool = False,
+        bypass_cached: bool = False,
     ) -> bool:
         """
         Retrieves an auth token and inserts it into the supplied request_headers dict.
@@ -72,20 +75,29 @@ class IAPToolkit:
                 As a general guideline, OIDC is the assumed default approach for ServiceAccounts.
             use_auth_header: If true, use the 'Authorization' header instead of 'Proxy-Authorization'
 
+        Returns:
+            True if token retrieved from cache, False if fresh from API
+
 
         """
-        if not use_oauth2:
-            token_refresh_struct: TokenRefreshStruct = self.get_token_oidc()
+        id_token = None
+        from_cache = False
+        if use_oauth2:
+            token_refresh_struct: TokenRefreshStruct = self.get_token_oauth2(bypass_cached=bypass_cached)
+            id_token = token_refresh_struct.id_token
+            from_cache = token_refresh_struct.from_cache
         else:
-            token_refresh_struct: TokenRefreshStruct = self.get_token_oauth2()
+            token_struct: TokenStruct = self.get_token_oidc(bypass_cached=bypass_cached)
+            id_token = token_struct.id_token
+            from_cache = token_struct.from_cache
 
         headers.add_token_to_request_headers(
             request_headers=request_headers,
-            id_token=token_refresh_struct.id_token,
+            id_token=id_token,
             use_auth_header=use_auth_header,
         )
 
-        return token_refresh_struct.token_is_new
+        return from_cache
 
     @staticmethod
     def is_url_safe_for_token(
@@ -104,6 +116,7 @@ class IAPToolkit:
         valid_domains: t.List[str] | None = None,
         use_oauth2: bool = False,
         use_auth_header: bool = False,
+        bypass_cached: bool = False,
     ) -> ResultAddTokenHeader:
         """
         Checks that the supplied URL is valid (i.e.; in valid_domains) and if so, retrieves a
@@ -123,10 +136,11 @@ class IAPToolkit:
                 request_headers=request_headers,
                 use_oauth2=use_oauth2,
                 use_auth_header=use_auth_header,
+                bypass_cached=bypass_cached,
             )
             return ResultAddTokenHeader(token_added=True, token_is_fresh=token_is_fresh)
         else:
-            LOG.warn(
+            LOG.warning(
                 "URL is not approved: %s - Token will not be added to headers. Valid domains are: %s",
                 url,
                 valid_domains,
@@ -146,10 +160,17 @@ class IAPToolkit_OIDC(IAPToolkit):
         raise NotImplementedError("Cannot call OAuth2 methods on OIDC-only instance of IAPToolkit.")
 
     def get_token_and_add_to_headers(
-        self, request_headers: dict, use_auth_header: bool = False, use_oauth2: bool = False,
+        self,
+        request_headers: dict,
+        use_auth_header: bool = False,
+        use_oauth2: bool = False,
+        bypass_cached: bool = False,
     ) -> bool:
         return super().get_token_and_add_to_headers(
-            request_headers=request_headers, use_oauth2=use_oauth2, use_auth_header=use_auth_header
+            request_headers=request_headers,
+            use_oauth2=use_oauth2,
+            use_auth_header=use_auth_header,
+            bypass_cached=bypass_cached,
         )
 
     def check_url_and_add_token_header(
@@ -158,6 +179,7 @@ class IAPToolkit_OIDC(IAPToolkit):
         request_headers: dict,
         valid_domains: t.List[str] | None = None,
         use_auth_header: bool = False,
+        bypass_cached: bool = False,
     ) -> ResultAddTokenHeader:
         return super().check_url_and_add_token_header(
             url,
@@ -165,6 +187,7 @@ class IAPToolkit_OIDC(IAPToolkit):
             valid_domains=valid_domains,
             use_oauth2=False,
             use_auth_header=use_auth_header,
+            bypass_cached=bypass_cached,
         )
 
 
@@ -193,10 +216,17 @@ class IAPToolkit_OAuth2(IAPToolkit):
         raise NotImplementedError("Cannot call OIDC methods on OAuth2-only instance of IAPToolkit.")
 
     def get_token_and_add_to_headers(
-        self, request_headers: dict, use_auth_header: bool = False, use_oauth2: bool = True,
+        self,
+        request_headers: dict,
+        use_auth_header: bool = False,
+        use_oauth2: bool = True,
+        bypass_cached: bool = False,
     ) -> bool:
         return super().get_token_and_add_to_headers(
-            request_headers=request_headers, use_oauth2=use_oauth2, use_auth_header=use_auth_header
+            request_headers=request_headers,
+            use_oauth2=use_oauth2,
+            use_auth_header=use_auth_header,
+            bypass_cached=bypass_cached,
         )
 
     def check_url_and_add_token_header(
@@ -205,6 +235,7 @@ class IAPToolkit_OAuth2(IAPToolkit):
         request_headers: dict,
         valid_domains: t.List[str] | None = None,
         use_auth_header: bool = False,
+        bypass_cached: bool = False,
     ) -> ResultAddTokenHeader:
         return super().check_url_and_add_token_header(
             url=url,
@@ -212,4 +243,5 @@ class IAPToolkit_OAuth2(IAPToolkit):
             valid_domains=valid_domains,
             use_oauth2=True,
             use_auth_header=use_auth_header,
+            bypass_cached=bypass_cached,
         )

{iaptoolkit-0.2.5 → iaptoolkit-0.3.0}/src/iaptoolkit/exceptions.py

@@ -23,9 +23,7 @@ class TokenStorageException(TokenException):
 
 
 class ServiceAccountTokenException(TokenException):
-    def __init__(
-        self, message: str, google_exception: t.Union[DefaultCredentialsError, RefreshError] | None
-    ):
+    def __init__(self, message: str, google_exception: t.Union[DefaultCredentialsError, RefreshError] | None):
         self.google_exception = google_exception
         credentials_env_var_value = os.environ.get(GOOGLE_CREDENTIALS_FILE_PATH)
         metadata_server_attempted = not credentials_env_var_value

{iaptoolkit-0.2.5 → iaptoolkit-0.3.0}/src/iaptoolkit/headers.py

@@ -34,9 +34,7 @@ def sanitize_request_headers(headers: dict) -> dict:
     return log_safe_headers
 
 
-def add_token_to_request_headers(
-    request_headers: dict, id_token: str, use_auth_header: bool = False
-) -> dict:
+def add_token_to_request_headers(request_headers: dict, id_token: str, use_auth_header: bool = False) -> dict:
     """
     Adds Bearer token to headers dict. Modifies dict in-place.
     Returns True if added token is a fresh one, or False if token is from cache

{iaptoolkit-0.2.5 → iaptoolkit-0.3.0}/src/iaptoolkit/tokens/service_account.py

@@ -14,6 +14,7 @@ from iaptoolkit.tokens.token_datastore import datastore
 from iaptoolkit.exceptions import ServiceAccountTokenException
 from iaptoolkit.exceptions import ServiceAccountTokenFailedRefresh
 from iaptoolkit.exceptions import ServiceAccountNoDefaultCredentials
+from iaptoolkit.exceptions import TokenException
 from iaptoolkit.exceptions import TokenStorageException
 
 from .structs import TokenStruct
@@ -41,39 +42,33 @@ class ServiceAccount(object):
     def get_stored_token(iap_client_id: str) -> t.Optional[TokenStruct]:
         try:
             token_dict = datastore.get_stored_service_account_token(iap_client_id)
-            if (
-                not token_dict
-                or not token_dict.get("id_token", None)
-                or not token_dict.get("token_expiry", None)
-            ):
+            if not token_dict or not token_dict.get("id_token", None) or not token_dict.get("token_expiry", None):
                 LOG.debug("No stored service account token for current iap_client_id")
                 return
 
-            id_token_from_dict = token_dict.get("id_token")
-            token_expiry_from_dict = token_dict.get("token_expiry", "")
-
-            if not id_token_from_dict:
-                LOG.warning("Invalid stored ID token")
-                return
+            id_token_from_dict: str = token_dict.get("id_token", "")
+            token_expiry_from_dict: str = token_dict.get("token_expiry", "")
 
             token_expiry = ""
             try:
                 token_expiry = datetime.datetime.fromisoformat(token_expiry_from_dict)
             except (ValueError, TypeError) as ex:
-                LOG.debug("Invalid token expiry for current iap_client_id")
+                LOG.debug("Invalid token expiry for stored token - Could not parse from ISO format to datetime.")
                 return
 
-            token_struct = TokenStruct(id_token=id_token_from_dict, expiry=token_expiry)
+            token_struct = TokenStruct(id_token=id_token_from_dict, expiry=token_expiry, from_cache=True)
+            if not token_struct.valid:
+                LOG.debug("Stored service account token for current iap_client_id is INVALID")
+                return
             if token_struct.expired:
                 LOG.debug("Stored service account token for current iap_client_id has EXPIRED")
                 return
+
             return token_struct
 
         except Exception as ex:
             # Err on the side of not letting token-caching break requests, hence blanket except
-            raise TokenStorageException(
-                f"Exception when trying to retrieve stored token. exception={ex}"
-            )
+            raise TokenStorageException(f"Exception when trying to retrieve stored token. exception={ex}")
 
     @staticmethod
     def _get_fresh_credentials(iap_client_id: str) -> GoogleIDTokenCredentials:
@@ -104,6 +99,8 @@ class ServiceAccount(object):
     def _get_fresh_token(iap_client_id: str) -> TokenStruct:
         google_credentials = ServiceAccount._get_fresh_credentials(iap_client_id)
         id_token: str = str(google_credentials.token)
+        if not id_token:
+            raise TokenException("Invalid [empty] token retrieved for Service Account.")
 
         # Google lib uses deprecated 'utcfromtimestamp' func as of v2.29.x
         # e.g.: datetime.datetime.utcfromtimestamp(payload["exp"])
@@ -111,12 +108,10 @@ class ServiceAccount(object):
         # Python datetimes assume local TZ, and we want to explicitly only work in UTC here.
         token_expiry = google_credentials.expiry.replace(tzinfo=datetime.timezone.utc)
 
-        return TokenStruct(id_token=id_token, expiry=token_expiry)
+        return TokenStruct(id_token=id_token, expiry=token_expiry, from_cache=False)
 
     @staticmethod
-    def get_token(
-        iap_client_id: str, bypass_cached: bool = False, attempts: int = 0
-    ) -> TokenRefreshStruct:
+    def get_token(iap_client_id: str, bypass_cached: bool = False, attempts: int = 0) -> TokenStruct:
         """Retrieves an OIDC token for the current environment either from environment variable or from
         metadata service.
 
@@ -140,19 +135,17 @@ class ServiceAccount(object):
         use_cache = not bypass_cached
 
         try:
-            token_from_cache = False
-            token_struct = (use_cache and ServiceAccount.get_stored_token(iap_client_id)) or None
-            if use_cache and token_struct:
-                token_from_cache = True
-            else:
-                token_struct = ServiceAccount._get_fresh_token(iap_client_id)
+            token_struct: TokenStruct | None = None
 
-            ServiceAccount._store_token(iap_client_id, token_struct.id_token, token_struct.expiry)
+            if use_cache:
+                token_struct = ServiceAccount.get_stored_token(iap_client_id)
 
-            token_refresh_struct = TokenRefreshStruct(
-                id_token=token_struct.id_token, token_is_new=not token_from_cache
-            )
-            return token_refresh_struct
+            if not token_struct:
+                token_struct = ServiceAccount._get_fresh_token(iap_client_id)
+                if use_cache:
+                    ServiceAccount._store_token(iap_client_id, token_struct.id_token, token_struct.expiry)
+
+            return token_struct
 
         except ServiceAccountTokenException as ex:
             attempts += 1
@@ -173,16 +166,14 @@ class GoogleServiceAccount(ServiceAccount):
 
     def __init__(self, iap_client_id: str) -> None:
         if not iap_client_id or not isinstance(iap_client_id, str):
-            raise ServiceAccountTokenException(
-                "Invalid iap_client_id for GoogleServiceAccount", google_exception=None
-            )
+            raise ServiceAccountTokenException("Invalid iap_client_id for GoogleServiceAccount", google_exception=None)
         self._iap_client_id = iap_client_id
         super().__init__()
 
     def get_stored_token(self) -> t.Optional[TokenStruct]:
         return ServiceAccount.get_stored_token(self._iap_client_id)
 
-    def get_token(self, bypass_cached: bool = False, attempts: int = 0) -> TokenRefreshStruct:
+    def get_token(self, bypass_cached: bool = False, attempts: int = 0) -> TokenStruct:
         return ServiceAccount.get_token(
             iap_client_id=self._iap_client_id, bypass_cached=bypass_cached, attempts=attempts
         )

{iaptoolkit-0.2.5 → iaptoolkit-0.3.0}/src/iaptoolkit/tokens/structs.py

@@ -8,10 +8,18 @@ from kvcommon import logger
 LOG = logger.get_logger("iaptk")
 
 
+def validate_token(token: str | None) -> bool:
+    if not isinstance(token, str) or token.strip() == "":
+        return False
+
+    return True
+
+
 @dataclass(kw_only=True)
 class TokenStruct:
     id_token: str
     expiry: datetime.datetime
+    from_cache: bool = False
 
     @property
     def expired(self):
@@ -30,17 +38,28 @@ class TokenStruct:
             LOG.error("Exception when checking token expiry. exception=%s", ex)
             return True
 
+    @property
+    def valid(self):
+        return validate_token(self.id_token)
+
 
 @dataclass(kw_only=True)
 class TokenRefreshStruct:
     id_token: str
-    token_is_new: bool = True
+    from_cache: bool = False
 
+    @property
+    def valid(self):
+        return validate_token(self.id_token)
 
 @dataclass(kw_only=True)
 class TokenStructOAuth2(TokenStruct):
     refresh_token: str
-    new_refresh_token: bool = False
+    from_cache: bool = False
+
+    @property
+    def valid(self):
+        return validate_token(self.refresh_token)
 
 
 @dataclass(kw_only=True)

{iaptoolkit-0.2.5 → iaptoolkit-0.3.0}/src/iaptoolkit/tokens/token_datastore.py

@@ -4,9 +4,11 @@ import typing as t
 from kvcommon import logger
 from kvcommon.datastore.backend import DatastoreBackend
 from kvcommon.datastore.backend import DictBackend
+
 # from kvcommon.datastore.backend import TOMLBackend
 from kvcommon.datastore import VersionedDatastore
 
+from iaptoolkit.exceptions import TokenException
 from iaptoolkit.constants import IAPTOOLKIT_CONFIG_VERSION
 
 
@@ -38,9 +40,10 @@ class TokenDatastore(VersionedDatastore):
             return
         return token_struct_dict
 
-    def store_service_account_token(
-        self, iap_client_id: str, id_token: str, token_expiry: datetime.datetime
-    ):
+    def store_service_account_token(self, iap_client_id: str, id_token: str, token_expiry: datetime.datetime):
+        if not id_token:
+            raise TokenException("TokenDatastore: Attempting to store invalid [empty] token")
+
         tokens_dict = self.get_or_create_nested_dict(self._service_account_tokens_key)
         tokens_dict[iap_client_id] = dict(id_token=id_token, token_expiry=token_expiry.isoformat())
 
@@ -62,6 +65,7 @@ class TokenDatastore(VersionedDatastore):
     #     # TODO: OAuth2
     #     raise NotImplementedError()
 
+
 datastore = TokenDatastore(DictBackend)
 
 # if PERSISTENT_DATASTORE_ENABLED:

{iaptoolkit-0.2.5 → iaptoolkit-0.3.0}/src/iaptoolkit/utils/urls.py

@@ -21,7 +21,8 @@ def is_url_safe_for_token(
             f"Invalid url_parts - Expected a ParseResult - Got: "
             f"'{str(url_parts)}' (type#: {type(url_parts).__name__})"
         )
-    if allowed_domains is not None and not isinstance(allowed_domains, (list, set, tuple)) :
+
+    if allowed_domains is not None and not isinstance(allowed_domains, (list, set, tuple)):
         raise TypeError("allowed_domains must be a list, set or tuple if not None")
 
     netloc = get_netloc_without_port_from_url_parts(url_parts)
@@ -34,8 +35,7 @@ def is_url_safe_for_token(
     for domain in allowed_domains:
         if domain == "" or not isinstance(domain, str):
             raise InvalidDomain(
-                f"Empty or non-string domain in allowed_domains: "
-                f"'{str(domain)}' (type#: {type(domain).__name__})"
+                f"Empty or non-string domain in allowed_domains: " f"'{str(domain)}' (type#: {type(domain).__name__})"
             )
 
         if netloc.endswith(domain):