lint

Author: twishabansal

packages/toolbox-core/src/toolbox_core/auth_methods.py

@@ -33,33 +33,41 @@
     tools = await toolbox.load_toolset()
 """
 
+import asyncio
 from datetime import datetime, timedelta, timezone
 from typing import Any, Dict
+
 import google.auth
 from google.auth.exceptions import GoogleAuthError
-from google.auth.transport.requests import Request, AuthorizedSession
+from google.auth.transport.requests import AuthorizedSession, Request
 from google.oauth2 import id_token
-import asyncio
 
 # --- Constants ---
 BEARER_TOKEN_PREFIX = "Bearer "
 CACHE_REFRESH_MARGIN = timedelta(seconds=60)
 
-_token_cache: Dict[str, Any] = {"token": None, "expires_at": datetime.min.replace(tzinfo=timezone.utc)}
+_token_cache: Dict[str, Any] = {
+    "token": None,
+    "expires_at": datetime.min.replace(tzinfo=timezone.utc),
+}
+
 
 def _is_token_valid() -> bool:
     """Checks if the cached token exists and is not nearing expiry."""
     if not _token_cache["token"]:
         return False
-    return datetime.now(timezone.utc) < (_token_cache["expires_at"] - CACHE_REFRESH_MARGIN)
+    return datetime.now(timezone.utc) < (
+        _token_cache["expires_at"] - CACHE_REFRESH_MARGIN
+    )
+
 
 def _update_cache(new_token: str) -> None:
     """
     Validates a new token, extracts its expiry, and updates the cache.
-    
+
     Args:
         new_token: The new JWT ID token string.
-    
+
     Raises:
         ValueError: If the token is invalid or its expiry cannot be determined.
     """
@@ -68,13 +76,15 @@ def _update_cache(new_token: str) -> None:
         # signature and claims against Google's public keys.
         # It's a synchronous, CPU-bound operation, safe for async contexts.
         claims = id_token.verify_oauth2_token(new_token, Request())
-        
+
         expiry_timestamp = claims.get("exp")
         if not expiry_timestamp:
             raise ValueError("Token does not contain an 'exp' claim.")
-            
+
         _token_cache["token"] = new_token
-        _token_cache["expires_at"] = datetime.fromtimestamp(expiry_timestamp, tz=timezone.utc)
+        _token_cache["expires_at"] = datetime.fromtimestamp(
+            expiry_timestamp, tz=timezone.utc
+        )
 
     except (ValueError, GoogleAuthError) as e:
         # Clear cache on failure to prevent using a stale or invalid token
@@ -83,13 +93,12 @@ def _update_cache(new_token: str) -> None:
         raise ValueError(f"Failed to validate and cache the new token: {e}") from e
 
 
-# --- Public API Functions ---
-
 def get_google_id_token(audience: str) -> str:
     """
     Synchronously fetches a Google ID token for a specific audience.
-
-    This function uses Application Default Credentials and caches the token in memory.
+    This function uses Application Default Credentials for local systems
+    and standard google auth libraries for Google Cloud environments.
+    It caches the token in memory.
 
     Args:
         audience: The audience for the ID token (e.g., a service URL or client ID).
@@ -103,7 +112,7 @@ def get_google_id_token(audience: str) -> str:
     """
     if _is_token_valid():
         return BEARER_TOKEN_PREFIX + _token_cache["token"]
-    
+
     # Get local user credentials
     credentials, _ = google.auth.default()
     session = AuthorizedSession(credentials)
@@ -122,10 +131,29 @@ def get_google_id_token(audience: str) -> str:
         new_token = id_token.fetch_id_token(request, audience)
         _update_cache(new_token)
         return BEARER_TOKEN_PREFIX + _token_cache["token"]
-        
+
     except GoogleAuthError as e:
-        raise GoogleAuthError(f"Failed to fetch Google ID token for audience '{audience}': {e}") from e
-    
+        raise GoogleAuthError(
+            f"Failed to fetch Google ID token for audience '{audience}': {e}"
+        ) from e
+
+
 async def aget_google_id_token(audience: str) -> str:
-    token = await asyncio.to_thread(get_google_id_token, audience) 
-    return token
+    """
+    Asynchronously fetches a Google ID token for a specific audience.
+    This function uses Application Default Credentials for local systems
+    and standard google auth libraries for Google Cloud environments.
+    It caches the token in memory.
+
+    Args:
+        audience: The audience for the ID token (e.g., a service URL or client ID).
+
+    Returns:
+        A string in the format "Bearer <google_id_token>".
+
+    Raises:
+        GoogleAuthError: If fetching credentials or the token fails.
+        ValueError: If the fetched token is invalid.
+    """
+    token = await asyncio.to_thread(get_google_id_token, audience)
+    return token