iaptoolkit 0.3.9__tar.gz → 0.4.0__tar.gz

{iaptoolkit-0.3.9 → iaptoolkit-0.4.0}/PKG-INFO

@@ -1,7 +1,8 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: iaptoolkit
-Version: 0.3.9
+Version: 0.4.0
 Summary: Library of common utils for interacting with Identity-Aware Proxies
+License-File: LICENSE
 Author: Rob Voigt
 Author-email: code@ravoigt.com
 Requires-Python: >=3.11,<4.0
@@ -9,8 +10,10 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
 Requires-Dist: google-auth (>=2.29.0,<3.0.0)
-Requires-Dist: kvcommon[k8s] (>=0.4.1,<0.5.0)
+Requires-Dist: google-cloud-iam (>=2.20.0,<3.0.0)
+Requires-Dist: kvcommon[k8s] (>=0.4.5,<0.5.0)
 Requires-Dist: requests (>=2.32.4)
 Requires-Dist: toml (>=0.10.2,<0.11.0)
 Project-URL: Homepage, https://github.com/RAVoigt/iaptoolkit
@@ -48,6 +51,7 @@ def example1(url: str):
     result = iaptk.check_url_and_add_token_header(
         url=url,
         request_headers=headers,
+        iap_audience="some_iap_client_id_string" # OAuth Client ID for the IAP-protected resource as 'audience'
         valid_domains=allowed_domains
     )
     # result.token_added (bool) indicates if the token was added, depending on whether or not URL was valid
@@ -58,20 +62,57 @@ def example1(url: str):
 
 
 # Example #2 - Separate Calls - Functionally the same as Example 1 but more flexibility in URL validation
-def example1(url: str):
+def example2(url: str):
     is_url_safe: bool = iaptk.is_url_safe_for_token(url=url, valid_domains=valid_domains)
 
     if not is_url_safe:
         raise ExampleBadURLException("This URL isn't safe to send token headers to!")
 
     headers = dict()
-    token_is_fresh: bool = iaptk.get_token_and_add_to_headers(request_headers=headers)
+    token_is_fresh: bool = iaptk.get_token_and_add_to_headers(
+        request_headers=headers,
+        iap_audience="some_iap_client_id_string" # OAuth Client ID for the IAP-protected resource as 'audience'
+    )
     # token_is_fresh indicates if token was newly retrieved (True), or if a cached token was reused (False)
     # headers dict now contains the appropriate Bearer Token header for Google IAP
 
     # Make HTTP GET request with requests lib, with our headers containing bearer token to auth with IAP
     response = requests.request("GET", url, headers=headers)
 
+# Example #3 - Service Account JWT (instead of OIDC Token)
+def example3(url: str):
+    headers = dict()
+    result = iaptk.check_url_and_add_jwt_header(
+        url=url,
+        request_headers=headers,
+        service_account_email="service-account@PROJECT_ID.iam.gserviceaccount.com",
+        url_audience="https://some-iap-protected.resource/path",
+        valid_domains=allowed_domains
+    )
+    # result.token_added (bool) indicates if the token was added, depending on whether or not URL was valid
+    # headers dict now contains the appropriate Bearer JWT header for Google IAP
+
+    # Make HTTP GET request with requests lib, with our headers containing bearer token to auth with IAP
+    response = requests.request("GET", url, headers=headers)
+
+# Example #4 - Separate Calls - Service Account JWT - Functionally the same as Example 3 but more flexibility in URL validation
+def example4(url: str):
+    is_url_safe: bool = iaptk.is_url_safe_for_token(url=url, valid_domains=valid_domains)
+
+    if not is_url_safe:
+        raise ExampleBadURLException("This URL isn't safe to send token headers to!")
+
+    headers = dict()
+    token_is_fresh: bool = iaptk.get_jwt_and_add_to_headers(
+        request_headers=headers,
+        service_account_email="service-account@PROJECT_ID.iam.gserviceaccount.com",
+        url_audience="https://some-iap-protected.resource/path"
+    )
+    # token_is_fresh indicates if token was newly retrieved (True), or if a cached token was reused (False)
+    # headers dict now contains the appropriate Bearer Token header for Google IAP
+
+    # Make HTTP GET request with requests lib, with our headers containing bearer token to auth with IAP
+    response = requests.request("GET", url, headers=headers)
 ```
 
 ## Disclaimer

iaptoolkit-0.4.0/README.md

@@ -0,0 +1,98 @@
+# IAP Toolkit
+
+A library of utils to ease programmatic authentication with Google IAP (and ideally other IAPs in future).
+
+# PyPi
+https://pypi.org/project/iaptoolkit/
+
+# Installation
+### With Poetry:
+`poetry add iaptoolkit`
+
+### With pip:
+`pip install iaptoolkit`
+
+## Quick Start / Example Usage
+
+```python
+import requests
+
+from iaptoolkit import IAPToolkit
+
+iaptk = IAPToolkit(google_iap_client_id="EXAMPLE_ID_123456789ABCDEF")
+allowed_domains = ["example.com", ]
+
+
+# Example #1 - Combined Calls
+def example1(url: str):
+    headers = dict()
+    result = iaptk.check_url_and_add_token_header(
+        url=url,
+        request_headers=headers,
+        iap_audience="some_iap_client_id_string" # OAuth Client ID for the IAP-protected resource as 'audience'
+        valid_domains=allowed_domains
+    )
+    # result.token_added (bool) indicates if the token was added, depending on whether or not URL was valid
+    # headers dict now contains the appropriate Bearer Token header for Google IAP
+
+    # Make HTTP GET request with requests lib, with our headers containing bearer token to auth with IAP
+    response = requests.request("GET", url, headers=headers)
+
+
+# Example #2 - Separate Calls - Functionally the same as Example 1 but more flexibility in URL validation
+def example2(url: str):
+    is_url_safe: bool = iaptk.is_url_safe_for_token(url=url, valid_domains=valid_domains)
+
+    if not is_url_safe:
+        raise ExampleBadURLException("This URL isn't safe to send token headers to!")
+
+    headers = dict()
+    token_is_fresh: bool = iaptk.get_token_and_add_to_headers(
+        request_headers=headers,
+        iap_audience="some_iap_client_id_string" # OAuth Client ID for the IAP-protected resource as 'audience'
+    )
+    # token_is_fresh indicates if token was newly retrieved (True), or if a cached token was reused (False)
+    # headers dict now contains the appropriate Bearer Token header for Google IAP
+
+    # Make HTTP GET request with requests lib, with our headers containing bearer token to auth with IAP
+    response = requests.request("GET", url, headers=headers)
+
+# Example #3 - Service Account JWT (instead of OIDC Token)
+def example3(url: str):
+    headers = dict()
+    result = iaptk.check_url_and_add_jwt_header(
+        url=url,
+        request_headers=headers,
+        service_account_email="service-account@PROJECT_ID.iam.gserviceaccount.com",
+        url_audience="https://some-iap-protected.resource/path",
+        valid_domains=allowed_domains
+    )
+    # result.token_added (bool) indicates if the token was added, depending on whether or not URL was valid
+    # headers dict now contains the appropriate Bearer JWT header for Google IAP
+
+    # Make HTTP GET request with requests lib, with our headers containing bearer token to auth with IAP
+    response = requests.request("GET", url, headers=headers)
+
+# Example #4 - Separate Calls - Service Account JWT - Functionally the same as Example 3 but more flexibility in URL validation
+def example4(url: str):
+    is_url_safe: bool = iaptk.is_url_safe_for_token(url=url, valid_domains=valid_domains)
+
+    if not is_url_safe:
+        raise ExampleBadURLException("This URL isn't safe to send token headers to!")
+
+    headers = dict()
+    token_is_fresh: bool = iaptk.get_jwt_and_add_to_headers(
+        request_headers=headers,
+        service_account_email="service-account@PROJECT_ID.iam.gserviceaccount.com",
+        url_audience="https://some-iap-protected.resource/path"
+    )
+    # token_is_fresh indicates if token was newly retrieved (True), or if a cached token was reused (False)
+    # headers dict now contains the appropriate Bearer Token header for Google IAP
+
+    # Make HTTP GET request with requests lib, with our headers containing bearer token to auth with IAP
+    response = requests.request("GET", url, headers=headers)
+```
+
+## Disclaimer
+
+This project is not affiliated with Google. No trademark infringement intended.

{iaptoolkit-0.3.9 → iaptoolkit-0.4.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "iaptoolkit"
-version = "0.3.9"
+version = "0.4.0"
 description = "Library of common utils for interacting with Identity-Aware Proxies"
 authors = ["Rob Voigt <code@ravoigt.com>"]
 readme = "README.md"
@@ -23,9 +23,10 @@ include = '\.pyi?$'
 [tool.poetry.dependencies]
 python = "^3.11"
 google-auth = "^2.29.0"
+google-cloud-iam = "^2.20.0"
 requests = ">=2.32.4"
 toml = "^0.10.2"
-kvcommon = {extras = ["k8s"], version = "^0.4.1"}
+kvcommon = {extras = ["k8s"], version = "^0.4.5"}
 
 [tool.poetry.group.dev.dependencies]
 black = "*"

{iaptoolkit-0.3.9 → iaptoolkit-0.4.0}/src/iaptoolkit/__init__.py

@@ -27,39 +27,54 @@ class IAPToolkit:
     Class to encapsulate client-specific vars and forward them to static functions
     """
 
-    _GOOGLE_IAP_CLIENT_ID: str
-
-    def __init__(self, google_iap_client_id: str) -> None:
-        self._GOOGLE_IAP_CLIENT_ID = google_iap_client_id
-
     @staticmethod
     def sanitize_request_headers(request_headers: dict) -> dict:
         return headers.sanitize_request_headers(request_headers)
 
-    def get_token_oidc(self, bypass_cached: bool = False) -> TokenStruct:
+    @staticmethod
+    def get_service_account_jwt(
+        service_account_email: str, url_audience: str, bypass_cached: bool = False
+    ) -> TokenStruct:
+        try:
+            return ServiceAccount.get_jwt(
+                service_account_email=service_account_email,
+                url_audience=url_audience,
+                bypass_cached=bypass_cached,
+            )
+        except ServiceAccountTokenException as ex:
+            LOG.warning(ex)
+            raise
+
+    @staticmethod
+    def get_token_oidc(iap_audience: str, bypass_cached: bool = False) -> TokenStruct:
         try:
             return ServiceAccount.get_token(
-                iap_client_id=self._GOOGLE_IAP_CLIENT_ID, bypass_cached=bypass_cached,
+                iap_audience=iap_audience,
+                bypass_cached=bypass_cached,
             )
         except ServiceAccountTokenException as ex:
-            LOG.debug(ex)
+            LOG.warning(ex)
             raise
 
-    def get_token_oidc_str(self, bypass_cached: bool = False) -> str:
-        struct = self.get_token_oidc(bypass_cached=bypass_cached)
+    @staticmethod
+    def get_token_oidc_str(iap_audience: str, bypass_cached: bool = False) -> str:
+        struct = IAPToolkit.get_token_oidc(iap_audience=iap_audience, bypass_cached=bypass_cached)
         return struct.id_token
 
-    def get_token_oauth2(self, bypass_cached: bool = False) -> TokenRefreshStruct:
+    @staticmethod
+    def get_token_oauth2(bypass_cached: bool = False) -> TokenRefreshStruct:
         # TODO
         raise NotImplementedError()
 
-    def get_token_oauth2_str(self, bypass_cached: bool = False) -> str:
-        struct = self.get_token_oauth2(bypass_cached=bypass_cached)
+    @staticmethod
+    def get_token_oauth2_str(bypass_cached: bool = False) -> str:
+        struct = IAPToolkit.get_token_oauth2(bypass_cached=bypass_cached)
         return struct.id_token
 
+    @staticmethod
     def get_token_and_add_to_headers(
-        self,
         request_headers: dict,
+        iap_audience: str,
         use_oauth2: bool = False,
         use_auth_header: bool = False,
         bypass_cached: bool = False,
@@ -82,23 +97,68 @@ class IAPToolkit:
         id_token = None
         from_cache = False
         if use_oauth2:
-            token_refresh_struct: TokenRefreshStruct = self.get_token_oauth2(bypass_cached=bypass_cached)
+            token_refresh_struct: TokenRefreshStruct = IAPToolkit.get_token_oauth2(bypass_cached=bypass_cached)
             id_token = token_refresh_struct.id_token
             from_cache = token_refresh_struct.from_cache
         else:
-            token_struct: TokenStruct = self.get_token_oidc(bypass_cached=bypass_cached)
+            token_struct: TokenStruct = IAPToolkit.get_token_oidc(
+                iap_audience=iap_audience, 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=id_token, use_auth_header=use_auth_header,
+            request_headers=request_headers,
+            id_token=id_token,
+            use_auth_header=use_auth_header,
+        )
+
+        return from_cache
+
+    @staticmethod
+    def get_jwt_and_add_to_headers(
+        request_headers: dict,
+        service_account_email: str,
+        url_audience: str,
+        use_auth_header: bool = False,
+        bypass_cached: bool = False,
+    ) -> bool:
+        """
+        Retrieves a Service Account JWT and inserts it into the supplied request_headers dict.
+
+        request_headers is modified in-place
+
+        Params:
+            request_headers: dict of headers to insert into
+            service_account_email: Email address of the Service Account requesting access to IAP resource
+            url_audience: Target URL/Audience of IAP resource
+            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
+
+
+        """
+        from_cache = False
+
+        token_struct: TokenStruct = IAPToolkit.get_service_account_jwt(
+            service_account_email=service_account_email, url_audience=url_audience, bypass_cached=bypass_cached
+        )
+        signed_jwt = token_struct.id_token
+        from_cache = token_struct.from_cache
+
+        headers.add_token_to_request_headers(
+            request_headers=request_headers,
+            id_token=signed_jwt,
+            use_auth_header=use_auth_header,
         )
 
         return from_cache
 
     @staticmethod
     def is_url_safe_for_token(
-        url: str | ParseResult, valid_domains: t.Optional[t.List[str] | t.Set[str] | t.Tuple[str]] = None,
+        url: str | ParseResult,
+        valid_domains: t.Optional[t.List[str] | t.Set[str] | t.Tuple[str]] = None,
     ):
         if not isinstance(url, ParseResult):
             url = urlparse(url)
@@ -109,6 +169,7 @@ class IAPToolkit:
         self,
         url: str | ParseResult,
         request_headers: dict,
+        iap_audience: str,
         valid_domains: t.List[str] | None = None,
         use_oauth2: bool = False,
         use_auth_header: bool = False,
@@ -130,18 +191,60 @@ class IAPToolkit:
         if self.is_url_safe_for_token(url=url, valid_domains=valid_domains):
             token_is_fresh = self.get_token_and_add_to_headers(
                 request_headers=request_headers,
+                iap_audience=iap_audience,
                 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, token_is_jwt=False)
+        else:
+            LOG.warning(
+                "URL is not approved: %s - Token will not be added to headers. Valid domains are: %s",
+                url,
+                valid_domains,
+            )
+            return ResultAddTokenHeader(token_added=False, token_is_fresh=False, token_is_jwt=False)
+
+    def check_url_and_add_jwt_header(
+        self,
+        url: str | ParseResult,
+        request_headers: dict,
+        service_account_email: str,
+        url_audience: str,
+        valid_domains: t.List[str] | None = None,
+        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
+        Service Account JWT and adds it to request_headers.
+
+        i.e.; A convenience wrapper with logging for is_url_safe_for_token() and get_jwt_and_add_to_headers()
+
+        Params:
+            url: URL string or urllib.ParseResult to check for validity
+            service_account_email: Email address of the Service Account requesting access to IAP resource
+            url_audience: Target URL/Audience of IAP resource - Should be the same as url or a substring thereof
+            request_headers: Dict of headers to insert into
+            valid_domains: List of domains to validate URL against
+        """
+
+        if self.is_url_safe_for_token(url=url, valid_domains=valid_domains):
+            token_is_fresh = self.get_jwt_and_add_to_headers(
+                request_headers=request_headers,
+                service_account_email=service_account_email,
+                url_audience=url_audience,
+                use_auth_header=use_auth_header,
                 bypass_cached=bypass_cached,
             )
-            return ResultAddTokenHeader(token_added=True, token_is_fresh=token_is_fresh)
+            return ResultAddTokenHeader(token_added=True, token_is_fresh=token_is_fresh, token_is_jwt=True)
         else:
             LOG.warning(
                 "URL is not approved: %s - Token will not be added to headers. Valid domains are: %s",
                 url,
                 valid_domains,
             )
-            return ResultAddTokenHeader(token_added=False, token_is_fresh=False)
+            return ResultAddTokenHeader(token_added=False, token_is_fresh=False, token_is_jwt=True)
 
 
 class IAPToolkit_OIDC(IAPToolkit):
@@ -158,13 +261,14 @@ class IAPToolkit_OIDC(IAPToolkit):
     def get_token_and_add_to_headers(
         self,
         request_headers: dict,
+        iap_audience: str,
         use_auth_header: bool = False,
-        use_oauth2: bool = False,
         bypass_cached: bool = False,
     ) -> bool:
         return super().get_token_and_add_to_headers(
+            iap_audience=iap_audience,
             request_headers=request_headers,
-            use_oauth2=use_oauth2,
+            use_oauth2=False,
             use_auth_header=use_auth_header,
             bypass_cached=bypass_cached,
         )
@@ -173,6 +277,7 @@ class IAPToolkit_OIDC(IAPToolkit):
         self,
         url: str | ParseResult,
         request_headers: dict,
+        iap_audience: str,
         valid_domains: t.List[str] | None = None,
         use_auth_header: bool = False,
         bypass_cached: bool = False,
@@ -180,6 +285,7 @@ class IAPToolkit_OIDC(IAPToolkit):
         return super().check_url_and_add_token_header(
             url,
             request_headers=request_headers,
+            iap_audience=iap_audience,
             valid_domains=valid_domains,
             use_oauth2=False,
             use_auth_header=use_auth_header,
@@ -195,8 +301,12 @@ class IAPToolkit_OAuth2(IAPToolkit):
     _GOOGLE_CLIENT_ID: str
     _GOOGLE_CLIENT_SECRET: str
 
-    def __init__(self, google_iap_client_id: str, google_client_id: str, google_client_secret: str,) -> None:
-        super().__init__(google_iap_client_id=google_iap_client_id)
+    def __init__(
+        self,
+        google_client_id: str,
+        google_client_secret: str,
+    ) -> None:
+        super().__init__()
         self._GOOGLE_CLIENT_ID = google_client_id
         self._GOOGLE_CLIENT_SECRET = google_client_secret
 
@@ -207,15 +317,12 @@ 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,
-        bypass_cached: bool = False,
+        self, request_headers: dict, iap_audience: str, use_auth_header: bool = False, bypass_cached: bool = False
     ) -> bool:
         return super().get_token_and_add_to_headers(
             request_headers=request_headers,
-            use_oauth2=use_oauth2,
+            iap_audience=iap_audience,
+            use_oauth2=True,
             use_auth_header=use_auth_header,
             bypass_cached=bypass_cached,
         )
@@ -224,6 +331,7 @@ class IAPToolkit_OAuth2(IAPToolkit):
         self,
         url: str | ParseResult,
         request_headers: dict,
+        iap_audience: str,
         valid_domains: t.List[str] | None = None,
         use_auth_header: bool = False,
         bypass_cached: bool = False,
@@ -231,6 +339,7 @@ class IAPToolkit_OAuth2(IAPToolkit):
         return super().check_url_and_add_token_header(
             url=url,
             request_headers=request_headers,
+            iap_audience=iap_audience,
             valid_domains=valid_domains,
             use_oauth2=True,
             use_auth_header=use_auth_header,

{iaptoolkit-0.3.9 → iaptoolkit-0.4.0}/src/iaptoolkit/exceptions.py

@@ -1,6 +1,7 @@
 import os
 import typing as t
 
+from google.api_core.exceptions import PermissionDenied
 from google.auth.environment_vars import CREDENTIALS as GOOGLE_CREDENTIALS_FILE_PATH
 from google.auth.exceptions import DefaultCredentialsError
 from google.auth.exceptions import RefreshError
@@ -27,7 +28,9 @@ class IAPClientIDException(IAPToolkitBaseException):
 
 
 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, PermissionDenied] | 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
@@ -42,7 +45,9 @@ class ServiceAccountTokenException(TokenException):
 
     @property
     def retryable(self):
-        return self.google_exception and self.google_exception._retryable
+        if not self.google_exception:
+            return False
+        return getattr(self.google_exception, "_retryable", False)
 
 
 class ServiceAccountNoDefaultCredentials(ServiceAccountTokenException):
@@ -53,6 +58,10 @@ class ServiceAccountTokenFailedRefresh(ServiceAccountTokenException):
     pass
 
 
+class JWTPermissionException(ServiceAccountTokenException):
+    pass
+
+
 class InvalidDomain(IAPToolkitBaseException):
     pass
 

iaptoolkit-0.4.0/src/iaptoolkit/tokens/service_account.py

@@ -0,0 +1,276 @@
+import datetime
+import json
+import typing as t
+
+import google.auth
+import google.api_core.exceptions
+from google.auth.compute_engine import IDTokenCredentials as GoogleIDTokenCredentials
+from google.auth.exceptions import DefaultCredentialsError as GoogleDefaultCredentialsError
+from google.auth.exceptions import RefreshError as GoogleRefreshError
+from google.auth.transport.requests import Request as GoogleRequest
+from google.cloud import iam_credentials_v1
+from google.oauth2 import id_token as google_id_token_lib
+
+from kvcommon import logger
+
+from iaptoolkit import exceptions
+from iaptoolkit.tokens.token_datastore import datastore
+
+from .structs import TokenStruct
+
+
+LOG = logger.get_logger("iaptk")
+MAX_RECURSE = 3
+
+
+def _utcnow() -> datetime.datetime:
+    return datetime.datetime.now(tz=datetime.UTC)
+
+
+class ServiceAccount(object):
+    """Base class for interacting with service accounts and OIDC tokens for IAP"""
+
+    # TODO: This is a static namespace for SA functions. Turn it into a per-iap-client-id client
+    # TODO: Move Google-specific logic to GoogleServiceAccount
+
+    @staticmethod
+    def _store_token(iap_client_id: str, id_token: str, token_expiry: datetime.datetime):
+        try:
+            datastore.store_service_account_token(iap_client_id, id_token, token_expiry)
+        except Exception as ex:  # Err on the side of not letting token-caching break requests.
+            raise exceptions.TokenStorageException(f"Exception when trying to store token. exception={ex}")
+
+    @staticmethod
+    def _store_jwt(service_account_email: str, url_audience: str, signed_jwt: str, expiry: datetime.datetime):
+        try:
+            datastore.store_service_account_jwt(
+                service_account_email=service_account_email,
+                url_audience=url_audience,
+                signed_jwt=signed_jwt,
+                expiry=expiry,
+            )
+        except Exception as ex:  # Err on the side of not letting token-caching break requests.
+            raise exceptions.TokenStorageException(f"Exception when trying to store token. exception={ex}")
+
+    @staticmethod
+    def get_stored_token(iap_client_id: str) -> TokenStruct | None:
+        try:
+            return datastore.get_stored_service_account_token(iap_client_id)
+
+        except Exception as ex:
+            # Err on the side of not letting token-caching break requests, hence blanket except
+            raise exceptions.TokenStorageException(f"Exception when trying to retrieve stored token. exception={ex}")
+
+    @staticmethod
+    def get_stored_jwt(service_account_email: str, url_audience: str) -> TokenStruct | None:
+        try:
+            return datastore.get_stored_service_account_jwt(
+                service_account_email=service_account_email, url_audience=url_audience
+            )
+
+        except Exception as ex:
+            # Err on the side of not letting token-caching break requests, hence blanket except
+            raise exceptions.TokenStorageException(f"Exception when trying to retrieve stored token. exception={ex}")
+
+    @staticmethod
+    def _get_fresh_credentials(iap_client_id: str) -> GoogleIDTokenCredentials:
+
+        try:
+            request = GoogleRequest()
+            credentials: GoogleIDTokenCredentials = google_id_token_lib.fetch_id_token_credentials(
+                iap_client_id, request
+            )  # type: ignore
+            credentials.refresh(request)
+
+        except GoogleDefaultCredentialsError as ex:
+            # The exceptions that google's libs raise in this case are somewhat vague; wrap them.
+            raise exceptions.ServiceAccountNoDefaultCredentials(
+                message="Failed to get ServiceAccount token: Lacking default credentials.",
+                google_exception=ex,
+            )
+        except GoogleRefreshError as ex:
+            # Likely attempting to get a token for a service account in an environment that
+            # doesn't have one attached.
+            raise exceptions.ServiceAccountTokenFailedRefresh(
+                message="Failed to get ServiceAccount token: Refreshing token failed.",
+                google_exception=ex,
+            )
+        return credentials
+
+    @staticmethod
+    def _get_fresh_token(iap_client_id: str, use_jwt: bool = False) -> TokenStruct:
+        google_credentials = ServiceAccount._get_fresh_credentials(iap_client_id)
+        id_token: str = str(google_credentials.token)
+        if not id_token:
+            raise exceptions.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"])
+        # This creates a TZ-naive datetime in UTC from a POSIX timestamp.
+        # 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, from_cache=False)
+
+    @staticmethod
+    def _get_jwt(service_account_email: str, url_audience: str) -> TokenStruct:
+        """
+        Returns a signed JWT for the specified service account
+        """
+        now = _utcnow()
+        expiration_delta = datetime.timedelta(seconds=3595)
+        expiry_dt = now + expiration_delta
+
+        issued_at = int(now.timestamp())
+        expiry = int(expiry_dt.timestamp())
+
+        jwt_payload = {
+            "iss": service_account_email,
+            "sub": service_account_email,
+            "aud": url_audience,
+            "iat": issued_at,
+            "exp": expiry,
+        }
+        jwt_payload_str = json.dumps(jwt_payload)
+
+        source_credentials, project_id = google.auth.default()
+        iam_client = iam_credentials_v1.IAMCredentialsClient(credentials=source_credentials)  # type: ignore
+        name = iam_client.service_account_path("-", service_account_email)
+        response = iam_client.sign_jwt(name=name, payload=jwt_payload_str)
+        # return response.signed_jwt
+        return TokenStruct.for_jwt(signed_jwt=response.signed_jwt, expiry=expiry_dt, from_cache=False)
+
+    @staticmethod
+    def get_jwt(
+        service_account_email: str, url_audience: str, bypass_cached: bool = False, attempts: int = 0
+    ) -> TokenStruct:
+        use_cache = not bypass_cached
+
+        try:
+            token_struct: TokenStruct | None = None
+
+            if use_cache:
+                token_struct = ServiceAccount.get_stored_jwt(
+                    service_account_email=service_account_email, url_audience=url_audience
+                )
+
+            if not token_struct:
+                token_struct = ServiceAccount._get_jwt(
+                    service_account_email=service_account_email, url_audience=url_audience
+                )
+                if use_cache:
+                    ServiceAccount._store_jwt(
+                        service_account_email=service_account_email,
+                        url_audience=url_audience,
+                        signed_jwt=token_struct.id_token,
+                        expiry=token_struct.expiry,
+                    )
+
+            return token_struct
+
+        except google.api_core.exceptions.PermissionDenied as ex:
+            raise exceptions.JWTPermissionException(
+                "Permission denied while retrieving signed JWT for service account. "
+                "Service Account requires IAM Role: 'roles/iam.serviceAccountTokenCreator'",
+                google_exception=ex,
+            )
+
+        except exceptions.ServiceAccountTokenException as ex:
+            attempts += 1
+            if attempts > MAX_RECURSE or not ex.retryable:
+                raise
+            return ServiceAccount.get_jwt(
+                service_account_email=service_account_email,
+                url_audience=url_audience,
+                bypass_cached=False,
+                attempts=attempts,
+            )
+
+        except exceptions.TokenStorageException as ex:
+            if attempts > 1:
+                raise
+            attempts += 1
+            # Try again without involving the cache
+            return ServiceAccount.get_jwt(
+                service_account_email=service_account_email,
+                url_audience=url_audience,
+                bypass_cached=True,
+                attempts=attempts,
+            )
+
+    @staticmethod
+    def get_token(iap_audience: 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.
+
+        1. If the environment variable ``GOOGLE_APPLICATION_CREDENTIALS`` is set
+        to the path of a valid service account JSON file, then ID token is
+        acquired using this service account credentials.
+        2. If the application is running in Compute Engine, App Engine or Cloud Run,
+        then the ID token is obtained from the metadata server.
+
+        Args:
+            iap_client_id: The client ID used by IAP. Can be thought of as JWT audience.
+
+        Returns:
+            An OIDC token for use in connecting through IAP.
+
+        Raises:
+            :class:`ServiceAccountTokenException` if a token could not be retrieved due to either
+            missing credentials from env-var/JSON or inability to talk to metadata server.
+        """
+
+        use_cache = not bypass_cached
+
+        try:
+            token_struct: TokenStruct | None = None
+
+            if use_cache:
+                token_struct = ServiceAccount.get_stored_token(iap_audience)
+
+            if not token_struct:
+                token_struct = ServiceAccount._get_fresh_token(iap_audience)
+                if use_cache:
+                    ServiceAccount._store_token(iap_audience, token_struct.id_token, token_struct.expiry)
+
+            return token_struct
+
+        except exceptions.ServiceAccountTokenException as ex:
+            attempts += 1
+            if attempts > MAX_RECURSE or not ex.retryable:
+                raise
+            return ServiceAccount.get_token(iap_audience, bypass_cached=False, attempts=attempts)
+
+        except exceptions.TokenStorageException as ex:
+            if attempts > 1:
+                raise
+            attempts += 1
+            # Try again without involving the cache
+            return ServiceAccount.get_token(iap_audience, bypass_cached=True, attempts=attempts)
+
+
+class GoogleServiceAccount(ServiceAccount):
+    """
+    For interacting with Google service accounts, Service Account JWTs and OIDC tokens for Google IAP
+
+    Service Account requires IAM Role: 'roles/iam.serviceAccountTokenCreator'
+    """
+
+    def __init__(self, service_account_email: str) -> None:
+        if not service_account_email or not isinstance(service_account_email, str):
+            raise exceptions.ServiceAccountTokenException(
+                "Invalid iap_client_id for GoogleServiceAccount", google_exception=None
+            )
+        self._service_account_email = service_account_email
+        super().__init__()
+
+    def get_stored_jwt(self, url_audience: str) -> t.Optional[TokenStruct]:
+        return ServiceAccount.get_stored_jwt(self._service_account_email, url_audience=url_audience)
+
+    def get_jwt(self, url_audience: str, bypass_cached: bool = False, attempts: int = 0) -> TokenStruct:
+        return ServiceAccount.get_jwt(
+            service_account_email=self._service_account_email,
+            url_audience=url_audience,
+            bypass_cached=bypass_cached,
+            attempts=attempts,
+        )

{iaptoolkit-0.3.9 → iaptoolkit-0.4.0}/src/iaptoolkit/tokens/structs.py

@@ -20,6 +20,11 @@ class TokenStruct:
     id_token: str
     expiry: datetime.datetime
     from_cache: bool = False
+    is_jwt: bool = False
+
+    @classmethod
+    def for_jwt(cls, signed_jwt: str, expiry: datetime.datetime, from_cache: bool = False):
+        return TokenStruct(id_token=signed_jwt, is_jwt=True, from_cache=from_cache, expiry=expiry)
 
     @property
     def expired(self):
@@ -67,3 +72,4 @@ class TokenStructOAuth2(TokenStruct):
 class ResultAddTokenHeader:
     token_added: bool
     token_is_fresh: bool
+    token_is_jwt: bool

iaptoolkit-0.4.0/src/iaptoolkit/tokens/token_datastore.py

@@ -0,0 +1,133 @@
+import datetime
+import typing as t
+
+from kvcommon import logger
+from kvcommon.datastore.backend import DatastoreBackend
+from kvcommon.datastore.backend import DictBackend
+
+from kvcommon.datastore import VersionedDatastore
+
+from iaptoolkit.exceptions import TokenStorageException
+from iaptoolkit.constants import IAPTOOLKIT_CONFIG_VERSION
+
+from .structs import TokenStruct
+
+
+LOG = logger.get_logger("iaptk-ds")
+
+
+class TokenDatastore(VersionedDatastore):
+    _service_account_jwts_email_limit = 1000
+
+    def __init__(self, backend: DatastoreBackend | type[DatastoreBackend]) -> None:
+        super().__init__(backend=backend, config_version=IAPTOOLKIT_CONFIG_VERSION)
+        self._ensure_tokens_dict()
+
+    def _ensure_tokens_dict(self):
+        tokens_dict = self.get_or_create_nested_dict("tokens")
+        if "refresh" not in tokens_dict.keys():
+            tokens_dict["refresh"] = None
+        self.set_value("tokens", tokens_dict)
+
+    @property
+    def service_account_tokens(self) -> dict:
+        return self.get_or_create_nested_dict("service_account_tokens")
+
+    @property
+    def service_account_jwts(self) -> dict:
+        return self.get_or_create_nested_dict("service_account_jwts")
+
+    def discard_existing_tokens(self):
+        LOG.debug("Discarding existing tokens.")
+        self.update_data(tokens={})
+
+    def get_stored_service_account_token(self, iap_client_id: str) -> TokenStruct | None:
+        token_data = self.service_account_tokens.get(iap_client_id, None)
+        if not token_data or not token_data.id_token or not token_data.expiry:
+            LOG.debug("No stored service account token for current iap_client_id")
+            return
+        return self._dict_to_tokenstruct(token_data)
+
+    def store_service_account_token(self, iap_client_id: str, id_token: str, token_expiry: datetime.datetime):
+        if not id_token:
+            raise TokenStorageException("TokenDatastore: Attempting to store invalid [empty] token")
+
+        tokens_dict = self.service_account_tokens
+        self.service_account_tokens[iap_client_id] = dict(id_token=id_token, token_expiry=token_expiry.isoformat())
+
+        try:
+            self.update_data(service_account_tokens=tokens_dict)
+        except Exception as ex:
+            LOG.error("Failed to store service account token for re-use. exception=%s", ex)
+
+    def _get_or_create_dict_for_service_account_and_url(self, service_account_email: str, url_audience: str):
+        jwts_dict = self.service_account_jwts
+        jwts_dict_for_email = jwts_dict.get(service_account_email, dict())
+        jwts_dict[service_account_email] = jwts_dict_for_email
+
+        token_dict = jwts_dict_for_email.get(url_audience, dict())
+        return token_dict
+
+    def get_stored_service_account_jwt(self, service_account_email: str, url_audience: str) -> TokenStruct | None:
+        jwts_dict_for_email = self._get_or_create_dict_for_service_account_and_url(service_account_email, url_audience)
+        token_data = jwts_dict_for_email.get(url_audience, None)
+        if not token_data:
+            LOG.debug("No stored service account JWT for service account '%s'", service_account_email)
+            return
+        return self._dict_to_tokenstruct(token_data, is_jwt=True)
+
+    def store_service_account_jwt(
+        self, service_account_email: str, url_audience: str, signed_jwt: str, expiry: datetime.datetime
+    ):
+        if not signed_jwt:
+            raise TokenStorageException("TokenDatastore: Attempting to store invalid [empty] jwt")
+
+        try:
+            token_dict = self._get_or_create_dict_for_service_account_and_url(service_account_email, url_audience)
+            token_dict["signed_jwt"] = signed_jwt
+            token_dict["token_expiry"] = expiry.isoformat()
+        except Exception as ex:
+            LOG.error("Failed to store service account JWT for re-use. exception=%s", ex)
+
+    @staticmethod
+    def _dict_to_tokenstruct(token_data: dict, is_jwt: bool = False) -> TokenStruct | None:
+        if not token_data:
+            return
+
+        id_token_from_dict: str = token_data.get("id_token", "")
+        token_expiry_from_dict: str = token_data.get("token_expiry", "")
+        if not id_token_from_dict or not token_expiry_from_dict:
+            return
+
+        token_expiry = ""
+        try:
+            token_expiry = datetime.datetime.fromisoformat(token_expiry_from_dict)
+        except (ValueError, TypeError) as ex:
+            LOG.warning("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, from_cache=True, is_jwt=is_jwt)
+        if not token_struct.valid:
+            LOG.warning("Stored service account token is INVALID")
+            return
+        if token_struct.expired:
+            LOG.debug("Stored service account token has EXPIRED")
+            return
+
+        return token_struct
+
+    def _migrate_version(self):
+        # Override
+        self.discard_existing_tokens()
+        return super()._migrate_version()
+
+    # def get_stored_oauth2_token(self, iap_client_id: str):
+    #     # TODO: OAuth2
+    #     raise NotImplementedError()
+
+    # def store_oauth2_token(self, iap_client_id: str):
+    #     # TODO: OAuth2
+    #     raise NotImplementedError()
+
+
+datastore = TokenDatastore(DictBackend)

iaptoolkit-0.3.9/README.md

@@ -1,60 +0,0 @@
-# IAP Toolkit
-
-A library of utils to ease programmatic authentication with Google IAP (and ideally other IAPs in future).
-
-# PyPi
-https://pypi.org/project/iaptoolkit/
-
-# Installation
-### With Poetry:
-`poetry add iaptoolkit`
-
-### With pip:
-`pip install iaptoolkit`
-
-## Quick Start / Example Usage
-
-```python
-import requests
-
-from iaptoolkit import IAPToolkit
-
-iaptk = IAPToolkit(google_iap_client_id="EXAMPLE_ID_123456789ABCDEF")
-allowed_domains = ["example.com", ]
-
-
-# Example #1 - Combined Calls
-def example1(url: str):
-    headers = dict()
-    result = iaptk.check_url_and_add_token_header(
-        url=url,
-        request_headers=headers,
-        valid_domains=allowed_domains
-    )
-    # result.token_added (bool) indicates if the token was added, depending on whether or not URL was valid
-    # headers dict now contains the appropriate Bearer Token header for Google IAP
-
-    # Make HTTP GET request with requests lib, with our headers containing bearer token to auth with IAP
-    response = requests.request("GET", url, headers=headers)
-
-
-# Example #2 - Separate Calls - Functionally the same as Example 1 but more flexibility in URL validation
-def example1(url: str):
-    is_url_safe: bool = iaptk.is_url_safe_for_token(url=url, valid_domains=valid_domains)
-
-    if not is_url_safe:
-        raise ExampleBadURLException("This URL isn't safe to send token headers to!")
-
-    headers = dict()
-    token_is_fresh: bool = iaptk.get_token_and_add_to_headers(request_headers=headers)
-    # token_is_fresh indicates if token was newly retrieved (True), or if a cached token was reused (False)
-    # headers dict now contains the appropriate Bearer Token header for Google IAP
-
-    # Make HTTP GET request with requests lib, with our headers containing bearer token to auth with IAP
-    response = requests.request("GET", url, headers=headers)
-
-```
-
-## Disclaimer
-
-This project is not affiliated with Google. No trademark infringement intended.

iaptoolkit-0.3.9/src/iaptoolkit/tokens/service_account.py

@@ -1,186 +0,0 @@
-import datetime
-import typing as t
-
-from google.auth.compute_engine import IDTokenCredentials as GoogleIDTokenCredentials
-from google.auth.exceptions import DefaultCredentialsError as GoogleDefaultCredentialsError
-from google.auth.exceptions import RefreshError as GoogleRefreshError
-from google.auth.transport.requests import Request as GoogleRequest
-from google.oauth2 import id_token as google_id_token_lib
-
-from kvcommon import logger
-
-# TODO: Don't hardcode the association between OIDC/SA and dict-datastore
-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
-from .structs import TokenRefreshStruct
-
-
-LOG = logger.get_logger("iaptk")
-MAX_RECURSE = 3
-
-
-class ServiceAccount(object):
-    """Base class for interacting with service accounts and OIDC tokens for IAP"""
-
-    # TODO: This is a static namespace for SA functions. Turn it into a per-iap-client-id client
-    # TODO: Move Google-specific logic to GoogleServiceAccount
-
-    @staticmethod
-    def _store_token(iap_client_id: str, id_token: str, token_expiry: datetime.datetime):
-        try:
-            datastore.store_service_account_token(iap_client_id, id_token, token_expiry)
-        except Exception as ex:  # Err on the side of not letting token-caching break requests.
-            raise TokenStorageException(f"Exception when trying to store token. exception={ex}")
-
-    @staticmethod
-    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)
-            ):
-                LOG.debug("No stored service account token for current iap_client_id")
-                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 stored token - Could not parse from ISO format to datetime."
-                )
-                return
-
-            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}")
-
-    @staticmethod
-    def _get_fresh_credentials(iap_client_id: str) -> GoogleIDTokenCredentials:
-
-        try:
-            request = GoogleRequest()
-            credentials: GoogleIDTokenCredentials = google_id_token_lib.fetch_id_token_credentials(
-                iap_client_id, request
-            )  # type: ignore
-            credentials.refresh(request)
-
-        except GoogleDefaultCredentialsError as ex:
-            # The exceptions that google's libs raise in this case are somewhat vague; wrap them.
-            raise ServiceAccountNoDefaultCredentials(
-                message="Failed to get ServiceAccount token: Lacking default credentials.",
-                google_exception=ex,
-            )
-        except GoogleRefreshError as ex:
-            # Likely attempting to get a token for a service account in an environment that
-            # doesn't have one attached.
-            raise ServiceAccountTokenFailedRefresh(
-                message="Failed to get ServiceAccount token: Refreshing token failed.", google_exception=ex,
-            )
-        return credentials
-
-    @staticmethod
-    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"])
-        # This creates a TZ-naive datetime in UTC from a POSIX timestamp.
-        # 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, from_cache=False)
-
-    @staticmethod
-    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.
-
-        1. If the environment variable ``GOOGLE_APPLICATION_CREDENTIALS`` is set
-        to the path of a valid service account JSON file, then ID token is
-        acquired using this service account credentials.
-        2. If the application is running in Compute Engine, App Engine or Cloud Run,
-        then the ID token is obtained from the metadata server.
-
-        Args:
-            iap_client_id: The client ID used by IAP. Can be thought of as JWT audience.
-
-        Returns:
-            An OIDC token for use in connecting through IAP.
-
-        Raises:
-            :class:`ServiceAccountTokenException` if a token could not be retrieved due to either
-            missing credentials from env-var/JSON or inability to talk to metadata server.
-        """
-
-        use_cache = not bypass_cached
-
-        try:
-            token_struct: TokenStruct | None = None
-
-            if use_cache:
-                token_struct = ServiceAccount.get_stored_token(iap_client_id)
-
-            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
-            if attempts > MAX_RECURSE or not ex.retryable:
-                raise
-            return ServiceAccount.get_token(iap_client_id, bypass_cached=False, attempts=attempts)
-
-        except TokenStorageException as ex:
-            if attempts > 1:
-                raise
-            attempts += 1
-            # Try again without involving the cache
-            return ServiceAccount.get_token(iap_client_id, bypass_cached=True, attempts=attempts)
-
-
-class GoogleServiceAccount(ServiceAccount):
-    """For interacting with Google service accounts and OIDC tokens for Google IAP"""
-
-    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
-            )
-        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) -> TokenStruct:
-        return ServiceAccount.get_token(
-            iap_client_id=self._iap_client_id, bypass_cached=bypass_cached, attempts=attempts
-        )

iaptoolkit-0.3.9/src/iaptoolkit/tokens/token_datastore.py

@@ -1,72 +0,0 @@
-import datetime
-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
-
-
-LOG = logger.get_logger("iaptk-ds")
-
-
-class TokenDatastore(VersionedDatastore):
-    _service_account_tokens_key = "service_account_tokens"
-
-    def __init__(self, backend: DatastoreBackend | type[DatastoreBackend]) -> None:
-        super().__init__(backend=backend, config_version=IAPTOOLKIT_CONFIG_VERSION)
-        self._ensure_tokens_dict()
-
-    def _ensure_tokens_dict(self):
-        tokens_dict = self.get_or_create_nested_dict("tokens")
-        if "refresh" not in tokens_dict.keys():
-            tokens_dict["refresh"] = None
-        self.set_value("tokens", tokens_dict)
-
-    def discard_existing_tokens(self):
-        LOG.debug("Discarding existing tokens.")
-        self.update_data(tokens={})
-
-    def get_stored_service_account_token(self, iap_client_id: str) -> t.Optional[dict]:
-        tokens_dict = self.get_or_create_nested_dict(self._service_account_tokens_key)
-        token_struct_dict = tokens_dict.get(iap_client_id, None)
-        if not token_struct_dict:
-            LOG.debug("No stored service account token for current iap_client_id")
-            return
-        return token_struct_dict
-
-    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())
-
-        try:
-            self.update_data(service_account_tokens=tokens_dict)
-        except OSError as ex:
-            LOG.error("Failed to store service account token for re-use. exception=%s", ex)
-
-    def _migrate_version(self):
-        # Override
-        self.discard_existing_tokens()
-        return super()._migrate_version()
-
-    # def get_stored_oauth2_token(self, iap_client_id: str):
-    #     # TODO: OAuth2
-    #     raise NotImplementedError()
-
-    # def store_oauth2_token(self, iap_client_id: str):
-    #     # TODO: OAuth2
-    #     raise NotImplementedError()
-
-
-datastore = TokenDatastore(DictBackend)
-
-# if PERSISTENT_DATASTORE_ENABLED:
-#     datastore_toml = TokenDatastore(TOMLBackend(PERSISTENT_DATASTORE_PATH, PERSISTENT_DATASTORE_USERNAME))