iaptoolkit 0.2.2__tar.gz → 0.3.0a0__tar.gz

{iaptoolkit-0.2.2 → iaptoolkit-0.3.0a0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: iaptoolkit
-Version: 0.2.2
+Version: 0.3.0a0
 Summary: Library of common utils for interacting with Identity-Aware Proxies
 Author: Rob Voigt
 Author-email: code@ravoigt.com
@@ -9,7 +9,7 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Requires-Dist: google-auth (>=2.29.0,<3.0.0)
-Requires-Dist: kvcommon (>=0.1.3,<0.2.0)
+Requires-Dist: kvcommon (>=0.1.4,<0.2.0)
 Requires-Dist: pytest (>=7.4.4,<8.0.0)
 Requires-Dist: requests (>=2.31.0,<3.0.0)
 Requires-Dist: toml (>=0.10.2,<0.11.0)
@@ -34,9 +34,9 @@ https://pypi.org/project/iaptoolkit/
 ```python
 import requests
 
-from iaptoolkit import IAPToolkit
+from iaptoolkit import IAPToolkit_OIDC
 
-iaptk = IAPToolkit(google_iap_client_id="EXAMPLE_ID_123456789ABCDEF")
+iaptk_oidc = IAPToolkit_OIDC(google_iap_client_id="EXAMPLE_ID_123456789ABCDEF")
 allowed_domains = ["example.com", ]
 
 

{iaptoolkit-0.2.2 → iaptoolkit-0.3.0a0}/README.md

@@ -17,9 +17,9 @@ https://pypi.org/project/iaptoolkit/
 ```python
 import requests
 
-from iaptoolkit import IAPToolkit
+from iaptoolkit import IAPToolkit_OIDC
 
-iaptk = IAPToolkit(google_iap_client_id="EXAMPLE_ID_123456789ABCDEF")
+iaptk_oidc = IAPToolkit_OIDC(google_iap_client_id="EXAMPLE_ID_123456789ABCDEF")
 allowed_domains = ["example.com", ]
 
 

{iaptoolkit-0.2.2 → iaptoolkit-0.3.0a0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "iaptoolkit"
-version = "0.2.2"
+version = "0.3.0a"
 description = "Library of common utils for interacting with Identity-Aware Proxies"
 authors = ["Rob Voigt <code@ravoigt.com>"]
 readme = "README.md"
@@ -28,7 +28,7 @@ google-auth = "^2.29.0"
 requests = "^2.31.0"
 pytest = "^7.4.4"
 toml = "^0.10.2"
-kvcommon = "^0.1.3"
+kvcommon = "^0.1.4"
 
 [tool.poetry.dev-dependencies]
 black = "*"

{iaptoolkit-0.2.2 → iaptoolkit-0.3.0a0}/src/iaptoolkit/__init__.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from abc import ABC, abstractmethod
 import logging
 
 logging.getLogger(__name__).addHandler(logging.NullHandler())
@@ -8,59 +9,50 @@ import typing as t
 from urllib.parse import ParseResult
 from urllib.parse import urlparse
 
-from kvcommon import logger
+from kvcommon.logger import get_logger
 
 from iaptoolkit import headers
 from iaptoolkit.exceptions import ServiceAccountTokenException
-from iaptoolkit.tokens.service_account import ServiceAccount
+from iaptoolkit.tokens.base import BaseTokenInterface
+from iaptoolkit.tokens.oauth2 import OAuth2
+from iaptoolkit.tokens.oidc import OIDC
 from iaptoolkit.tokens.structs import ResultAddTokenHeader
-
 from iaptoolkit.tokens.structs import TokenRefreshStruct
 from iaptoolkit.utils.urls import is_url_safe_for_token
 
-LOG = logger.get_logger("iaptk")
+LOG = get_logger("iaptk")
 
 
-class IAPToolkit:
+class IAPToolkit(ABC):
     """
-    Class to encapsulate client-specific vars and forward them to static functions
+    Abstract base class wrapping up core iaptoolkit functionality in a single interface
+
+    In practice, you should use IAPToolkit_OIDC or IAPToolkit_OAuth2 for
+        OIDC(ServiceAccounts) or OAuth2(Users) respectively.
     """
 
     _GOOGLE_IAP_CLIENT_ID: str
+    _interface: BaseTokenInterface
 
-    def __init__(
-        self,
-        google_iap_client_id: str,
-    ) -> None:
+    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) -> TokenRefreshStruct:
-        try:
-            return ServiceAccount.get_token(
-                iap_client_id=self._GOOGLE_IAP_CLIENT_ID, bypass_cached=bypass_cached
-            )
-        except ServiceAccountTokenException as ex:
-            LOG.debug(ex)
-            raise
-
-    def get_token_oidc_str(self, bypass_cached: bool = False) -> str:
-        struct = self.get_token_oidc(bypass_cached=bypass_cached)
-        return struct.id_token
-
-    def get_token_oauth2(self, bypass_cached: bool = False) -> TokenRefreshStruct:
-        # TODO
+    @abstractmethod
+    def get_token(
+        self, refresh_token: str | None, bypass_cached: bool = False
+    ) -> TokenRefreshStruct:
         raise NotImplementedError()
 
-    def get_token_oauth2_str(self, bypass_cached: bool = False) -> str:
-        struct = self.get_token_oauth2(bypass_cached=bypass_cached)
+    def get_token_str(self, refresh_token: str | None, bypass_cached: bool = False) -> str:
+        struct = self.get_token(refresh_token=refresh_token, bypass_cached=bypass_cached)
         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_auth_header: bool = False, refresh_token: str | None = None
     ) -> bool:
         """
         Retrieves an auth token and inserts it into the supplied request_headers dict.
@@ -71,13 +63,9 @@ class IAPToolkit:
             use_oauth2: Use OAuth2.0 credentials and respective token, else use OIDC (default)
                 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'
-
-
+            refresh_token: Refresh token for OAuth2.0 (Unused by OIDC)
         """
-        if not use_oauth2:
-            token_refresh_struct: TokenRefreshStruct = self.get_token_oidc()
-        else:
-            token_refresh_struct: TokenRefreshStruct = self.get_token_oauth2()
+        token_refresh_struct = self.get_token(refresh_token=refresh_token)
 
         headers.add_token_to_request_headers(
             request_headers=request_headers,
@@ -102,8 +90,8 @@ class IAPToolkit:
         url: str | ParseResult,
         request_headers: dict,
         valid_domains: t.List[str] | None = None,
-        use_oauth2: bool = False,
         use_auth_header: bool = False,
+        refresh_token: str | None = None,
     ) -> ResultAddTokenHeader:
         """
         Checks that the supplied URL is valid (i.e.; in valid_domains) and if so, retrieves a
@@ -115,66 +103,54 @@ class IAPToolkit:
             url: URL string or urllib.ParseResult to check for validity
             request_headers: Dict of headers to insert into
             valid_domains: List of domains to validate URL against
-            use_oauth2: Passed to get_token_and_add_to_headers() to determine if OAuth2.0 is used or OIDC (default)
+            use_auth_header: If true, use the 'Authorization' header instead of 'Proxy-Authorization' for IAP
+            refresh_token: Refresh token for OAuth2.0 (Unused by OIDC)
         """
 
         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,
-                use_oauth2=use_oauth2,
                 use_auth_header=use_auth_header,
+                refresh_token=refresh_token,
             )
-            return ResultAddTokenHeader(token_added=True, token_is_fresh=token_is_fresh)
+            return ResultAddTokenHeader(token_added=True, token_is_new=token_is_fresh)
         else:
             LOG.warn(
                 "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_new=False)
 
 
 class IAPToolkit_OIDC(IAPToolkit):
     """
-    Convenience subclass of IAPToolkit for scenarios where OIDC will always be used, never OAuth2
+    OIDC-only implementation of IAPToolkit
     """
+    _interface: OIDC
 
-    def get_token_oauth2(self, *args, **kwargs):
-        raise NotImplementedError("Cannot call OAuth2 methods on OIDC-only instance of IAPToolkit.")
+    def __init__(self, google_iap_client_id: str) -> None:
+        super().__init__(google_iap_client_id)
+        self._interface = OIDC(iap_client_id=google_iap_client_id)
 
-    def get_token_oauth2_str(self, *args, **kwargs):
-        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
-    ) -> bool:
-        return super().get_token_and_add_to_headers(
-            request_headers=request_headers, use_oauth2=False, use_auth_header=use_auth_header
-        )
-
-    def check_url_and_add_token_header(
-        self,
-        url: str | ParseResult,
-        request_headers: dict,
-        valid_domains: t.List[str] | None = None,
-        use_auth_header: bool = False,
-    ) -> ResultAddTokenHeader:
-        return super().check_url_and_add_token_header(
-            url,
-            request_headers=request_headers,
-            valid_domains=valid_domains,
-            use_oauth2=False,
-            use_auth_header=use_auth_header,
-        )
+    def get_token(self, bypass_cached: bool = False) -> TokenRefreshStruct:
+        try:
+            return self._interface.get_token(
+                iap_client_id=self._GOOGLE_IAP_CLIENT_ID, bypass_cached=bypass_cached
+            )
+        except ServiceAccountTokenException as ex:
+            LOG.debug(ex)
+            raise
 
 
 class IAPToolkit_OAuth2(IAPToolkit):
     """
-    Convenience subclass of IAPToolkit for scenarios where OAuth2 will always be used, never OIDC
+    OAuth2.0-only implementation of IAPToolkit
     """
 
     _GOOGLE_CLIENT_ID: str
     _GOOGLE_CLIENT_SECRET: str
+    _interface: OAuth2
 
     def __init__(
         self,
@@ -185,24 +161,34 @@ class IAPToolkit_OAuth2(IAPToolkit):
         super().__init__(google_iap_client_id=google_iap_client_id)
         self._GOOGLE_CLIENT_ID = google_client_id
         self._GOOGLE_CLIENT_SECRET = google_client_secret
+        self._interface = OAuth2(iap_client_id=google_iap_client_id, client_id=google_client_id)
 
-    def get_token_oidc(self, *args, **kwargs):
-        raise NotImplementedError("Cannot call OIDC methods on OAuth2-only instance of IAPToolkit.")
+    def get_refresh_token(self, bypass_cached: bool = False) -> t.Any:
+        pass
 
-    def get_token_oidc_str(self, *args, **kwargs):
-        raise NotImplementedError("Cannot call OIDC methods on OAuth2-only instance of IAPToolkit.")
+    def get_token(self, refresh_token: str, bypass_cached: bool = False) -> TokenRefreshStruct:
+        if not self._GOOGLE_CLIENT_ID or not self._GOOGLE_CLIENT_SECRET:
+            raise ValueError()  # TODO
 
-    def get_token_and_add_to_headers(
-        self, request_headers: dict, use_auth_header: bool = False
-    ) -> bool:
-        return super().get_token_and_add_to_headers(
-            request_headers=request_headers, use_oauth2=True, use_auth_header=use_auth_header
-        )
+        # TODO: Get from cache and check expiry
+        expired = True
+
+        if expired or bypass_cached:
+            token: str = self._interface.get_id_token_from_refresh_token(
+                client_id=self._GOOGLE_CLIENT_ID,
+                client_secret=self._GOOGLE_CLIENT_SECRET,
+                refresh_token=refresh_token,
+                iap_client_id=self._GOOGLE_IAP_CLIENT_ID,
+            )
+
+        # TODO: Move this when implementing cache
+        return TokenRefreshStruct(id_token=token, token_is_new=expired or bypass_cached)
 
     def check_url_and_add_token_header(
         self,
         url: str | ParseResult,
         request_headers: dict,
+        refresh_token: str,
         valid_domains: t.List[str] | None = None,
         use_auth_header: bool = False,
     ) -> ResultAddTokenHeader:
@@ -210,6 +196,6 @@ class IAPToolkit_OAuth2(IAPToolkit):
             url=url,
             request_headers=request_headers,
             valid_domains=valid_domains,
-            use_oauth2=True,
             use_auth_header=use_auth_header,
+            refresh_token=refresh_token,
         )

{iaptoolkit-0.2.2 → iaptoolkit-0.3.0a0}/src/iaptoolkit/exceptions.py

@@ -14,6 +14,10 @@ class IAPBadResponse(IAPToolkitBaseException):
     pass
 
 
+class InvalidDomain(IAPToolkitBaseException):
+    pass
+
+
 class TokenException(IAPToolkitBaseException):
     pass
 
@@ -22,7 +26,10 @@ class TokenStorageException(TokenException):
     pass
 
 
-class ServiceAccountTokenException(TokenException):
+class GoogleTokenException(TokenException):
+    """
+    Wrapper for exceptions from Google's auth lib
+    """
     def __init__(
         self, message: str, google_exception: t.Union[DefaultCredentialsError, RefreshError] | None
     ):
@@ -42,8 +49,9 @@ class ServiceAccountTokenException(TokenException):
     def retryable(self):
         return self.google_exception and self.google_exception._retryable
 
+# SA / OIDC
 
-class ServiceAccountNoDefaultCredentials(ServiceAccountTokenException):
+class ServiceAccountTokenException(GoogleTokenException):
     pass
 
 
@@ -51,5 +59,14 @@ class ServiceAccountTokenFailedRefresh(ServiceAccountTokenException):
     pass
 
 
-class InvalidDomain(IAPToolkitBaseException):
+class ServiceAccountNoDefaultCredentials(ServiceAccountTokenException):
+    pass
+
+# OAuth2 / User
+
+class OAuth2RefreshFromAuthCodeFailed(TokenException):
+    pass
+
+
+class OAuth2IDTokenFromRefreshFailed(TokenException):
     pass

{iaptoolkit-0.2.2 → iaptoolkit-0.3.0a0}/src/iaptoolkit/headers.py

@@ -23,7 +23,9 @@ def _sanitize_request_header(headers_dict: dict, header_key: str):
 def sanitize_request_headers(headers: dict) -> dict:
     """
     Sanitizes a headers dict to remove sensitive strings for logging purposes.
-    Returns A COPY of the dict with sensitive k/v pairs replaced. Does NOT modify in-place/by-reference.
+
+    Returns:
+        A COPY of the dict with sensitive k/v pairs replaced. Does NOT modify in-place/by-reference.
     """
     log_safe_headers = headers.copy()
 

iaptoolkit-0.3.0a0/src/iaptoolkit/tokens/base.py

@@ -0,0 +1,68 @@
+import datetime
+from abc import ABC, abstractmethod
+
+from kvcommon import logger
+
+from iaptoolkit.exceptions import TokenStorageException
+from iaptoolkit.tokens.token_datastore import TokenDatastore
+
+from iaptoolkit.tokens.structs import TokenStruct
+
+
+LOG = logger.get_logger("iaptk")
+
+
+class BaseTokenInterface(ABC):
+    _datastore: TokenDatastore
+    _iap_client_id: str
+
+    def __init__(self, datastore: TokenDatastore, iap_client_id: str) -> None:
+        super().__init__()
+        self._iap_client_id = iap_client_id
+        self._datastore = datastore
+
+    @abstractmethod
+    def _get_stored_token(self) -> dict | None:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def _store_token(self, iap_client_id: str, id_token: str, token_expiry: datetime.datetime):
+        raise NotImplementedError()
+
+    def get_stored_token(self) -> TokenStruct | None:
+        try:
+            token_dict = self._get_stored_token()
+            if (
+                not token_dict
+                or not token_dict.get("id_token", None)
+                or not token_dict.get("token_expiry", None)
+            ):
+                LOG.debug("No stored token for supplied client_id(s)")
+                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
+
+            token_expiry = ""
+            try:
+                token_expiry = datetime.datetime.fromisoformat(token_expiry_from_dict)
+            except (ValueError, TypeError) as ex:
+                LOG.debug("Invalid token expiry for supplied client_id(s)")
+                return
+
+            token_struct = TokenStruct(id_token=id_token_from_dict, expiry=token_expiry)
+            if token_struct.expired:
+                LOG.debug("Stored OAuth2 token for supplied client_id(s) has EXPIRED")
+                return
+            return token_struct
+
+        except Exception as ex:
+            # Err on the side of not letting token-caching break requests, hence blanket except
+            # Caller can `try/except TokenStorageException` for safety
+            raise TokenStorageException(
+                f"Exception when trying to retrieve stored token. exception={ex}"
+            )

iaptoolkit-0.3.0a0/src/iaptoolkit/tokens/oauth2/__init__.py

@@ -0,0 +1,110 @@
+import datetime
+import json
+import typing as t
+import requests
+
+from kvcommon import logger
+from kvcommon.datastore.backend import DictBackend
+
+from iaptoolkit.tokens.base import BaseTokenInterface
+from iaptoolkit.exceptions import OAuth2IDTokenFromRefreshFailed
+from iaptoolkit.exceptions import OAuth2RefreshFromAuthCodeFailed
+
+from .datastore_oauth2 import TokenDatastore_OAuth2
+
+
+LOG = logger.get_logger("iaptk-oauth2")
+
+GOOGLE_OAUTH_TOKEN_URL = "https://www.googleapis.com/oauth2/v4/token"
+
+
+def get_localhost_redirect_uri(listen_port: int):
+    return f"http://localhost:{listen_port}"
+
+
+def get_oauth2_auth_url(client_id: str, redirect_uri: str):
+    # TODO: Unhardcode
+    return (
+        f"https://accounts.google.com/o/oauth2/v2/auth?client_id={client_id}"
+        f"&response_type=code&scope=openid%20email&access_type=offline&redirect_uri={redirect_uri}"
+    )
+
+
+class OAuth2(BaseTokenInterface):
+    """
+    Base class for interacting with OAuth2.0 tokens for IAP
+
+    OAuth2.0 access Tokens have a shorter expiry (<60mins)
+    Refresh tokens have a longer expiry and are used to retrieve new access tokens
+
+    # TODO: Move Google-specific logic to GoogleServiceAccount
+    """
+
+    _datastore: TokenDatastore_OAuth2
+    _client_id: str
+
+    def __init__(self, iap_client_id: str, client_id: str) -> None:
+        super().__init__(
+            datastore=TokenDatastore_OAuth2(DictBackend),
+            iap_client_id=iap_client_id,
+        )
+        self._client_id = client_id
+
+    def _store_token(self, id_token: str, token_expiry: datetime.datetime):
+        self._datastore.store_token(self._iap_client_id, self._client_id, id_token, token_expiry)
+
+    def _get_stored_token(self, iap_client_id: str, client_id: str) -> dict | None:
+        return self._datastore.get_stored_token(iap_client_id=iap_client_id, client_id=client_id)
+
+    # TODO: Unstatic
+    @staticmethod
+    def get_id_token_from_refresh_token(
+        client_id: str,
+        client_secret: str,
+        refresh_token: str,
+        iap_client_id: str,
+    ) -> str:
+
+        oauth2_token_url = GOOGLE_OAUTH_TOKEN_URL  # TODO: Unhardcode
+        request_payload = {
+            "client_id": client_id,
+            "client_secret": client_secret,
+            "refresh_token": refresh_token,
+            "grant_type": "refresh_token",
+            "audience": iap_client_id,
+        }
+        response = requests.post(oauth2_token_url, data=request_payload)
+        response_dict = json.loads(response.text)
+        id_token: str = response_dict.get("id_token", None)
+        if response.status_code != 200 or not id_token:
+            raise OAuth2IDTokenFromRefreshFailed(
+                f"Failure in acquiring OAuth2.0 access token from refresh token - HTTP Response:"
+                f"{response.status_code} : {response.reason or 'Unknown'} : {response.text or ''}"
+            )
+        return id_token
+
+    # TODO: Unstatic
+    @staticmethod
+    def get_refresh_token_from_auth_code(
+        client_id: str,
+        client_secret: str,
+        auth_code: str,
+        redirect_uri: str,
+    ) -> str:
+        oauth2_token_url = GOOGLE_OAUTH_TOKEN_URL  # TODO: Unhardcode
+        request_payload = {
+            "code": auth_code,
+            "client_id": client_id,
+            "client_secret": client_secret,
+            "grant_type": "authorization_code",
+            "redirect_uri": redirect_uri,
+        }
+        response = requests.post(oauth2_token_url, data=request_payload)
+        response_dict = json.loads(response.text)
+        refresh_token: str = response_dict.get("refresh_token", None)
+        if response.status_code != 200 or not refresh_token:
+            raise OAuth2RefreshFromAuthCodeFailed(
+                f"Failure in acquiring refresh token from auth code - HTTP Response:"
+                f"{response.status_code} : {response.reason or 'Unknown'} : {response.text or ''}"
+            )
+        return refresh_token

iaptoolkit-0.3.0a0/src/iaptoolkit/tokens/oauth2/datastore_oauth2.py

@@ -0,0 +1,39 @@
+import datetime
+import typing as t
+
+from kvcommon import logger
+
+from iaptoolkit.tokens.token_datastore import TokenDatastore
+
+
+LOG = logger.get_logger("iaptk-ds-oauth2")
+
+
+class TokenDatastore_OAuth2(TokenDatastore):
+    _tokens_key: str = "outh2_tokens"
+
+    def get_stored_token(self, iap_client_id: str, client_id: str) -> dict | None:
+        tokens_dict = self.get_or_create_nested_dict(self._tokens_key)
+        source_key = f"{iap_client_id}{client_id}"
+        token_struct_dict = self._retrieve_hashed_dict_entry(
+            target=tokens_dict, source_key=source_key
+        )
+        if not token_struct_dict:
+            LOG.debug("No stored service account token for current iap_client_id")
+            return
+        return token_struct_dict
+
+    def store_token(
+        self, iap_client_id: str, client_id: str, id_token: str, token_expiry: datetime.datetime
+    ):
+        tokens_dict = self.get_or_create_nested_dict(self._tokens_key)
+
+        # TODO: Encode/encrypt token?
+        value = dict(id_token=id_token, token_expiry=token_expiry.isoformat())
+        source_key = f"{iap_client_id}{client_id}"
+        self._insert_hashed_dict_entry(target=tokens_dict, source_key=source_key, value=value)
+
+        try:
+            self.update_data(outh2_tokens=tokens_dict)
+        except OSError as ex:
+            LOG.error("Failed to store OAuth2 token for re-use. exception=%s", ex)

iaptoolkit-0.3.0a0/src/iaptoolkit/tokens/oauth2/gua.py

@@ -0,0 +1,11 @@
+from . import OAuth2
+
+
+class GoogleUserAccount(OAuth2):
+    """
+    For interacting with Google user accounts and OAuth2.0 tokens for Google IAP
+
+    # TODO: Move Google-specific logic into this class
+    """
+
+    pass

iaptoolkit-0.2.2/src/iaptoolkit/tokens/service_account.py → iaptoolkit-0.3.0a0/src/iaptoolkit/tokens/oidc/__init__.py

@@ -1,188 +1,149 @@
-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 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 = 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
-
-            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")
-                return
-
-            token_struct = TokenStruct(id_token=id_token_from_dict, expiry=token_expiry)
-            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)
-
-        # 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)
-
-    @staticmethod
-    def get_token(
-        iap_client_id: str, bypass_cached: bool = False, attempts: int = 0
-    ) -> TokenRefreshStruct:
-        """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_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)
-
-            ServiceAccount._store_token(iap_client_id, token_struct.id_token, token_struct.expiry)
-
-            token_refresh_struct = TokenRefreshStruct(
-                id_token=token_struct.id_token, token_is_new=not token_from_cache
-            )
-            return token_refresh_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) -> TokenRefreshStruct:
-        return ServiceAccount.get_token(
-            iap_client_id=self._iap_client_id, bypass_cached=bypass_cached, attempts=attempts
-        )
+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.logger import get_logger
+from kvcommon.datastore.backend import DictBackend
+
+from iaptoolkit.exceptions import ServiceAccountTokenException
+from iaptoolkit.exceptions import ServiceAccountTokenFailedRefresh
+from iaptoolkit.exceptions import ServiceAccountNoDefaultCredentials
+from iaptoolkit.exceptions import TokenStorageException
+
+from iaptoolkit.tokens.base import BaseTokenInterface
+from iaptoolkit.tokens.structs import TokenStruct
+from iaptoolkit.tokens.structs import TokenRefreshStruct
+
+from .datastore_oidc import TokenDatastore_OIDC
+
+
+
+LOG = get_logger("iaptk-oidc")
+MAX_RECURSE = 3
+
+
+class OIDC(BaseTokenInterface):
+    """
+    Base class for interacting with service accounts and OIDC tokens for IAP
+
+    # TODO: Move Google-specific logic to GoogleServiceAccount
+    """
+    _datastore: TokenDatastore_OIDC
+
+    def __init__(self, iap_client_id: str) -> None:
+        super().__init__(
+            datastore=TokenDatastore_OIDC(DictBackend),
+            iap_client_id=iap_client_id,
+        )
+
+    def _store_token(self, id_token: str, token_expiry: datetime.datetime):
+        self._datastore.store_token(self._iap_client_id, id_token, token_expiry)
+
+    def _get_stored_token(self) -> dict | None:
+        return self._datastore.get_stored_token(iap_client_id=self._iap_client_id)
+
+    def get_stored_token(self) -> TokenStruct | None:
+        return super().get_stored_token()
+
+    # TODO: Unstatic
+    @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
+
+    # TODO: Unstatic
+    @staticmethod
+    def _get_fresh_token(iap_client_id: str) -> TokenStruct:
+        google_credentials = OIDC._get_fresh_credentials(iap_client_id)
+        id_token: str = str(google_credentials.token)
+
+        # 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)
+
+    # TODO: Unstatic
+    def get_token(
+        self, iap_client_id: str, bypass_cached: bool = False, attempts: int = 0
+    ) -> TokenRefreshStruct:
+        """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.
+            bypass_cached: Force retrieval of fresh tokens, bypassing in-memory cache
+            attempts: For recursive retries
+
+        Returns:
+            A struct containing:
+                id_token: An OIDC auth token for use in connecting through IAP
+                token_is_new: A bool indicating if the refresh token is new (i.e.; the previous had expired)
+
+        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_from_cache = False
+            token_struct = (use_cache and self.get_stored_token()) or None
+            if use_cache and token_struct:
+                token_from_cache = True
+            else:
+                token_struct = OIDC._get_fresh_token(iap_client_id)
+
+            self._store_token(token_struct.id_token, token_struct.expiry)
+
+            token_refresh_struct = TokenRefreshStruct(
+                id_token=token_struct.id_token, token_is_new=not token_from_cache
+            )
+            return token_refresh_struct
+
+        except ServiceAccountTokenException as ex:
+            attempts += 1
+            if attempts > MAX_RECURSE or not ex.retryable:
+                raise
+            return self.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 self.get_token(iap_client_id, bypass_cached=True, attempts=attempts)

iaptoolkit-0.3.0a0/src/iaptoolkit/tokens/oidc/datastore_oidc.py

@@ -0,0 +1,35 @@
+import datetime
+import typing as t
+
+from kvcommon import logger
+
+from iaptoolkit.tokens.token_datastore import TokenDatastore
+
+
+LOG = logger.get_logger("iaptk-ds-oidc")
+
+
+class TokenDatastore_OIDC(TokenDatastore):
+    _tokens_key: str = "oidc_tokens"
+
+    def get_stored_token(self, iap_client_id: str) -> dict | None:
+        tokens_dict = self.get_or_create_nested_dict(self._tokens_key)
+        token_struct_dict = self._retrieve_hashed_dict_entry(
+            target=tokens_dict, source_key=iap_client_id
+        )
+        if not token_struct_dict:
+            LOG.debug("No stored service account token for current iap_client_id")
+            return
+        return token_struct_dict
+
+    def store_token(self, iap_client_id: str, id_token: str, token_expiry: datetime.datetime):
+        tokens_dict = self.get_or_create_nested_dict(self._tokens_key)
+
+        # TODO: Encode/encrypt token?
+        value = dict(id_token=id_token, token_expiry=token_expiry.isoformat())
+        self._insert_hashed_dict_entry(target=tokens_dict, source_key=iap_client_id, value=value)
+
+        try:
+            self.update_data(oidc_tokens=tokens_dict)
+        except OSError as ex:
+            LOG.error("Failed to store service account token for re-use. exception=%s", ex)

iaptoolkit-0.3.0a0/src/iaptoolkit/tokens/oidc/gsa.py

@@ -0,0 +1,11 @@
+from . import OIDC
+
+
+class GoogleServiceAccount(OIDC):
+    """
+    For interacting with Google service accounts and OIDC tokens for Google IAP
+
+    # TODO: Move Google-specific logic into this class
+    """
+
+    pass

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

@@ -12,6 +12,7 @@ LOG = logger.get_logger("iaptk")
 class TokenStruct:
     id_token: str
     expiry: datetime.datetime
+    token_is_new: bool = True
 
     @property
     def expired(self):
@@ -33,6 +34,8 @@ class TokenStruct:
 
 @dataclass(kw_only=True)
 class TokenRefreshStruct:
+    """
+    """
     id_token: str
     token_is_new: bool = True
 
@@ -40,10 +43,10 @@ class TokenRefreshStruct:
 @dataclass(kw_only=True)
 class TokenStructOAuth2(TokenStruct):
     refresh_token: str
-    new_refresh_token: bool = False
+    token_is_new: bool = False
 
 
 @dataclass(kw_only=True)
 class ResultAddTokenHeader:
     token_added: bool
-    token_is_fresh: bool
+    token_is_new: bool

iaptoolkit-0.3.0a0/src/iaptoolkit/tokens/token_datastore.py

@@ -0,0 +1,63 @@
+import datetime
+import hashlib
+import typing as t
+from abc import abstractmethod
+
+from kvcommon import logger
+from kvcommon.datastore.backend import DatastoreBackend
+from kvcommon.datastore import VersionedDatastore
+
+from iaptoolkit.constants import IAPTOOLKIT_CONFIG_VERSION
+
+
+LOG = logger.get_logger("iaptk-ds")
+
+
+class TokenDatastore(VersionedDatastore):
+    _tokens_key: str
+
+    def __init__(self, backend: DatastoreBackend | type[DatastoreBackend]) -> None:
+        super().__init__(backend=backend, config_version=IAPTOOLKIT_CONFIG_VERSION)
+        self._ensure_tokens_dict()
+
+    def _migrate_version(self):
+        # Override
+        self.discard_existing_tokens()
+        return super()._migrate_version()
+
+    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)
+
+    @staticmethod
+    def _insert_hashed_dict_entry(target: dict, source_key: str, value: t.Any):
+        hash_obj = hashlib.sha256(source_key.encode("utf-8"))
+        hex_digest = hash_obj.hexdigest()
+        target[hex_digest] = value
+
+    @staticmethod
+    def _retrieve_hashed_dict_entry(target: dict, source_key: str) -> t.Any:
+        hash_obj = hashlib.sha256(source_key.encode("utf-8"))
+        hex_digest = hash_obj.hexdigest()
+        # TODO: Check collisions/pre-existing keys?
+        target.get(hex_digest)
+
+    def discard_existing_tokens(self):
+        LOG.debug("Discarding existing tokens.")
+        self.update_data(tokens={})
+
+    @abstractmethod
+    def get_stored_token(self, iap_client_id: str) -> dict | None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def store_token(self, iap_client_id: str, id_token: str, token_expiry: datetime.datetime):
+        raise NotImplementedError
+
+
+# datastore = TokenDatastore(DictBackend)
+
+# if PERSISTENT_DATASTORE_ENABLED:
+#     datastore_toml = TokenDatastore(TOMLBackend(PERSISTENT_DATASTORE_PATH, PERSISTENT_DATASTORE_USERNAME))

iaptoolkit-0.2.2/src/iaptoolkit/tokens/__init__.py

@@ -1,24 +0,0 @@
-from kvcommon import logger
-
-from iaptoolkit.exceptions import ServiceAccountTokenException
-from iaptoolkit.exceptions import TokenStorageException
-from iaptoolkit.exceptions import TokenException
-
-from .structs import TokenStruct
-from .structs import TokenRefreshStruct
-
-# from .structs import TokenStructOAuth2  # TODO: OAuth2
-# from .oauth2 import get_token_for_oauth2  # TODO: OAuth2
-# from .service_account import ServiceAccount
-from .service_account import GoogleServiceAccount
-
-LOG = logger.get_logger("iaptk")
-
-
-__all__ = [
-    "TokenStruct",
-    "TokenRefreshStruct",
-    # "TokenStructOAuth2",  # TODO: OAuth2
-    "TokenException",
-    "TokenStorageException",
-]

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

@@ -1,68 +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.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
-    ):
-        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))