google-adk-extras 0.1.1__py3-none-any.whl → 0.2.5__py3-none-any.whl

google_adk_extras/credentials/jwt_credential_service.py

@@ -0,0 +1,345 @@
+"""JWT credential service implementation."""
+
+from typing import Optional, Dict, Any
+import logging
+import jwt
+from datetime import datetime, timedelta, timezone
+
+from google.adk.auth.credential_service.session_state_credential_service import SessionStateCredentialService
+from google.adk.auth.credential_service.base_credential_service import CallbackContext
+from google.adk.auth import AuthConfig, AuthCredential, AuthCredentialTypes
+from google.adk.auth.auth_credential import HttpAuth, HttpCredentials
+from fastapi.openapi.models import HTTPBearer
+
+from .base_custom_credential_service import BaseCustomCredentialService
+
+logger = logging.getLogger(__name__)
+
+
+class JWTCredentialService(BaseCustomCredentialService):
+    """JWT credential service for handling JSON Web Token authentication.
+    
+    This service generates and manages JWT tokens for API authentication.
+    It supports both short-lived and long-lived tokens with automatic refresh.
+    
+    Args:
+        secret: The secret key used to sign JWT tokens.
+        algorithm: The algorithm used for JWT signing. Default is 'HS256'.
+        issuer: The issuer of the JWT token. Optional.
+        audience: The intended audience of the JWT token. Optional.
+        expiration_minutes: Token expiration time in minutes. Default is 60 minutes.
+        custom_claims: Additional custom claims to include in the JWT payload.
+        use_session_state: If True, stores credentials in session state. If False,
+            uses in-memory storage. Default is True for persistence.
+    
+    Example:
+        ```python
+        credential_service = JWTCredentialService(
+            secret="your-jwt-secret",
+            algorithm="HS256",
+            issuer="my-app",
+            audience="api.example.com",
+            expiration_minutes=120,
+            custom_claims={"role": "admin", "permissions": ["read", "write"]}
+        )
+        await credential_service.initialize()
+        
+        # Use with Runner
+        runner = Runner(
+            agent=agent,
+            session_service=session_service,
+            credential_service=credential_service,
+            app_name="my_app"
+        )
+        ```
+    """
+
+    SUPPORTED_ALGORITHMS = {
+        'HS256', 'HS384', 'HS512',  # HMAC with SHA
+        'RS256', 'RS384', 'RS512',  # RSA with SHA
+        'ES256', 'ES384', 'ES512'   # ECDSA with SHA
+    }
+
+    def __init__(
+        self,
+        secret: str,
+        algorithm: str = 'HS256',
+        issuer: Optional[str] = None,
+        audience: Optional[str] = None,
+        expiration_minutes: int = 60,
+        custom_claims: Optional[Dict[str, Any]] = None,
+        use_session_state: bool = True
+    ):
+        """Initialize the JWT credential service.
+        
+        Args:
+            secret: JWT signing secret.
+            algorithm: JWT signing algorithm.
+            issuer: JWT issuer.
+            audience: JWT audience.
+            expiration_minutes: Token expiration in minutes.
+            custom_claims: Additional claims to include in JWT.
+            use_session_state: Whether to use session state for credential storage.
+        """
+        super().__init__()
+        self.secret = secret
+        self.algorithm = algorithm
+        self.issuer = issuer
+        self.audience = audience
+        self.expiration_minutes = expiration_minutes
+        self.custom_claims = custom_claims or {}
+        self.use_session_state = use_session_state
+        
+        # Underlying credential service for storage
+        if use_session_state:
+            self._storage_service = SessionStateCredentialService()
+        else:
+            from google.adk.auth.credential_service.in_memory_credential_service import InMemoryCredentialService
+            self._storage_service = InMemoryCredentialService()
+
+    async def _initialize_impl(self) -> None:
+        """Initialize the JWT credential service.
+        
+        Validates the configuration parameters.
+        
+        Raises:
+            ValueError: If configuration is invalid.
+        """
+        if not self.secret:
+            raise ValueError("JWT secret is required")
+        
+        if self.algorithm not in self.SUPPORTED_ALGORITHMS:
+            raise ValueError(f"Unsupported JWT algorithm: {self.algorithm}. Supported: {self.SUPPORTED_ALGORITHMS}")
+        
+        if self.expiration_minutes <= 0:
+            raise ValueError("JWT expiration_minutes must be positive")
+            
+        # Test JWT creation to validate secret and algorithm
+        try:
+            test_payload = {"test": "validation"}
+            jwt.encode(test_payload, self.secret, algorithm=self.algorithm)
+            logger.info(f"Initialized JWT credential service with algorithm {self.algorithm}")
+        except Exception as e:
+            raise ValueError(f"Invalid JWT configuration: {e}")
+
+    def generate_jwt_token(self, user_id: str, additional_claims: Optional[Dict[str, Any]] = None) -> str:
+        """Generate a JWT token for the specified user.
+        
+        Args:
+            user_id: The user ID to include in the JWT token.
+            additional_claims: Additional claims to include in this specific token.
+            
+        Returns:
+            str: The generated JWT token.
+            
+        Raises:
+            RuntimeError: If the service is not initialized.
+        """
+        self._check_initialized()
+        
+        now = datetime.now(timezone.utc)
+        exp = now + timedelta(minutes=self.expiration_minutes)
+        
+        payload = {
+            "sub": user_id,  # Subject (user ID)
+            "iat": now,      # Issued at
+            "exp": exp,      # Expiration
+        }
+        
+        # Add optional standard claims
+        if self.issuer:
+            payload["iss"] = self.issuer
+        if self.audience:
+            payload["aud"] = self.audience
+            
+        # Add custom claims
+        payload.update(self.custom_claims)
+        if additional_claims:
+            payload.update(additional_claims)
+            
+        token = jwt.encode(payload, self.secret, algorithm=self.algorithm)
+        logger.debug(f"Generated JWT token for user {user_id} expiring at {exp}")
+        
+        return token
+
+    def verify_jwt_token(self, token: str) -> Dict[str, Any]:
+        """Verify and decode a JWT token.
+        
+        Args:
+            token: The JWT token to verify.
+            
+        Returns:
+            Dict[str, Any]: The decoded token payload.
+            
+        Raises:
+            jwt.InvalidTokenError: If the token is invalid or expired.
+            RuntimeError: If the service is not initialized.
+        """
+        self._check_initialized()
+        
+        options = {
+            "verify_signature": True,
+            "verify_exp": True,
+            "verify_iat": True,
+        }
+        
+        # Set audience verification if configured
+        audience = self.audience if self.audience else None
+        issuer = self.issuer if self.issuer else None
+        
+        payload = jwt.decode(
+            token,
+            self.secret,
+            algorithms=[self.algorithm],
+            audience=audience,
+            issuer=issuer,
+            options=options
+        )
+        
+        return payload
+
+    def is_token_expired(self, token: str) -> bool:
+        """Check if a JWT token is expired without raising an exception.
+        
+        Args:
+            token: The JWT token to check.
+            
+        Returns:
+            bool: True if the token is expired, False otherwise.
+        """
+        try:
+            self.verify_jwt_token(token)
+            return False
+        except jwt.ExpiredSignatureError:
+            return True
+        except jwt.InvalidTokenError:
+            # Other validation errors also count as "expired" for refresh purposes
+            return True
+
+    def create_auth_config(self, user_id: str, additional_claims: Optional[Dict[str, Any]] = None) -> AuthConfig:
+        """Create an AuthConfig with a generated JWT token.
+        
+        Args:
+            user_id: The user ID for the JWT token.
+            additional_claims: Additional claims for this specific token.
+            
+        Returns:
+            AuthConfig: Configured auth config with JWT Bearer token.
+        """
+        self._check_initialized()
+        
+        # Generate JWT token
+        token = self.generate_jwt_token(user_id, additional_claims)
+        
+        # Create HTTP Bearer auth scheme
+        auth_scheme = HTTPBearer()
+        
+        # Create HTTP Bearer credential
+        auth_credential = AuthCredential(
+            auth_type=AuthCredentialTypes.HTTP,
+            http=HttpAuth(
+                scheme="bearer",
+                credentials=HttpCredentials(token=token)
+            )
+        )
+        
+        return AuthConfig(
+            auth_scheme=auth_scheme,
+            raw_auth_credential=auth_credential
+        )
+
+    async def load_credential(
+        self,
+        auth_config: AuthConfig,
+        callback_context: CallbackContext,
+    ) -> Optional[AuthCredential]:
+        """Load JWT credential from storage and refresh if expired.
+        
+        Args:
+            auth_config: The auth config containing credential key information.
+            callback_context: The current callback context.
+            
+        Returns:
+            Optional[AuthCredential]: The stored credential or refreshed credential.
+        """
+        self._check_initialized()
+        
+        # Load existing credential
+        credential = await self._storage_service.load_credential(auth_config, callback_context)
+        
+        if not credential or not credential.http or not credential.http.credentials:
+            return None
+            
+        # Check if token needs refresh
+        token = credential.http.credentials.token
+        if not token or self.is_token_expired(token):
+            logger.info(f"JWT token expired for user {callback_context._invocation_context.user_id}, generating new token")
+            
+            # Generate new token
+            user_id = callback_context._invocation_context.user_id
+            new_token = self.generate_jwt_token(user_id)
+            
+            # Update credential with new token
+            credential.http.credentials.token = new_token
+            
+            # Save refreshed credential
+            updated_auth_config = AuthConfig(
+                auth_scheme=auth_config.auth_scheme,
+                raw_auth_credential=credential,
+                exchanged_auth_credential=credential
+            )
+            await self._storage_service.save_credential(updated_auth_config, callback_context)
+            
+        return credential
+
+    async def save_credential(
+        self,
+        auth_config: AuthConfig,
+        callback_context: CallbackContext,
+    ) -> None:
+        """Save JWT credential to storage.
+        
+        Args:
+            auth_config: The auth config containing the credential to save.
+            callback_context: The current callback context.
+        """
+        self._check_initialized()
+        await self._storage_service.save_credential(auth_config, callback_context)
+        
+        logger.info(f"Saved JWT credential for user {callback_context._invocation_context.user_id}")
+
+    def get_token_info(self, token: str) -> Dict[str, Any]:
+        """Get information about a JWT token without full verification.
+        
+        Args:
+            token: The JWT token to inspect.
+            
+        Returns:
+            Dict[str, Any]: Token information including claims and expiration.
+        """
+        try:
+            # Decode without verification to get token info
+            payload = jwt.decode(token, options={"verify_signature": False})
+            
+            # Check expiration without requiring initialization
+            expired = True
+            if self._initialized:
+                expired = self.is_token_expired(token)
+            elif "exp" in payload:
+                # Simple expiration check without initialization
+                exp_timestamp = payload["exp"]
+                expired = datetime.now(timezone.utc).timestamp() > exp_timestamp
+            
+            info = {
+                "payload": payload,
+                "expired": expired,
+            }
+            
+            if "exp" in payload:
+                exp_timestamp = payload["exp"]
+                exp_datetime = datetime.fromtimestamp(exp_timestamp, timezone.utc)
+                info["expires_at"] = exp_datetime.isoformat()
+                
+            return info
+        except Exception as e:
+            return {"error": str(e), "expired": True}

google_adk_extras/credentials/microsoft_oauth2_credential_service.py

@@ -0,0 +1,250 @@
+"""Microsoft OAuth2 credential service implementation."""
+
+from typing import Optional, List
+import logging
+
+from google.adk.auth.credential_service.session_state_credential_service import SessionStateCredentialService
+from google.adk.auth.credential_service.base_credential_service import CallbackContext
+from google.adk.auth import AuthConfig, AuthCredential, AuthCredentialTypes
+from google.adk.auth.auth_credential import OAuth2Auth
+from fastapi.openapi.models import OAuth2
+from fastapi.openapi.models import OAuthFlowAuthorizationCode, OAuthFlows
+
+from .base_custom_credential_service import BaseCustomCredentialService
+
+logger = logging.getLogger(__name__)
+
+
+class MicrosoftOAuth2CredentialService(BaseCustomCredentialService):
+    """Microsoft OAuth2 credential service for handling Microsoft authentication flows.
+    
+    This service provides pre-configured OAuth2 flows for Microsoft Graph API including
+    Outlook, Teams, OneDrive, and other Microsoft 365 services.
+    
+    Args:
+        tenant_id: The Azure AD tenant ID. Use "common" for multi-tenant applications.
+        client_id: The Microsoft OAuth2 client ID from Azure AD App Registration.
+        client_secret: The Microsoft OAuth2 client secret from Azure AD App Registration.
+        scopes: List of OAuth2 scopes to request. Common scopes include:
+            - "User.Read" - Read user profile
+            - "Mail.Read" - Read user's mail
+            - "Mail.ReadWrite" - Read and write user's mail
+            - "Calendars.Read" - Read user's calendars
+            - "Calendars.ReadWrite" - Read and write user's calendars
+            - "Files.Read" - Read user's files
+            - "Files.ReadWrite" - Read and write user's files
+        use_session_state: If True, stores credentials in session state. If False,
+            uses in-memory storage. Default is True for persistence.
+    
+    Example:
+        ```python
+        credential_service = MicrosoftOAuth2CredentialService(
+            tenant_id="common",  # or specific tenant ID
+            client_id="your-azure-client-id",
+            client_secret="your-azure-client-secret",
+            scopes=["User.Read", "Mail.Read", "Calendars.ReadWrite"]
+        )
+        await credential_service.initialize()
+        
+        # Use with Runner
+        runner = Runner(
+            agent=agent,
+            session_service=session_service,
+            credential_service=credential_service,
+            app_name="my_app"
+        )
+        ```
+    """
+
+    # Microsoft OAuth2 endpoints (v2.0)
+    def _get_auth_url(self, tenant_id: str) -> str:
+        """Get the Microsoft OAuth2 authorization URL for the tenant."""
+        return f"https://login.microsoftonline.com/{tenant_id}/oauth2/v2.0/authorize"
+    
+    def _get_token_url(self, tenant_id: str) -> str:
+        """Get the Microsoft OAuth2 token URL for the tenant."""
+        return f"https://login.microsoftonline.com/{tenant_id}/oauth2/v2.0/token"
+    
+    # Common Microsoft Graph API scopes
+    COMMON_SCOPES = {
+        # User and profile scopes
+        "User.Read": "Read user profile",
+        "User.ReadWrite": "Read and write user profile",
+        "User.ReadBasic.All": "Read basic profiles of all users",
+        "User.Read.All": "Read all users' full profiles",
+        "User.ReadWrite.All": "Read and write all users' full profiles",
+        
+        # Mail scopes
+        "Mail.Read": "Read user mail",
+        "Mail.ReadWrite": "Read and write user mail", 
+        "Mail.Send": "Send mail as user",
+        "Mail.Read.Shared": "Read user and shared mail",
+        "Mail.ReadWrite.Shared": "Read and write user and shared mail",
+        
+        # Calendar scopes
+        "Calendars.Read": "Read user calendars",
+        "Calendars.ReadWrite": "Read and write user calendars",
+        "Calendars.Read.Shared": "Read user and shared calendars",
+        "Calendars.ReadWrite.Shared": "Read and write user and shared calendars",
+        
+        # Files and OneDrive scopes
+        "Files.Read": "Read user files",
+        "Files.ReadWrite": "Read and write user files",
+        "Files.Read.All": "Read all files that user can access",
+        "Files.ReadWrite.All": "Read and write all files that user can access",
+        "Sites.Read.All": "Read items in all site collections",
+        "Sites.ReadWrite.All": "Read and write items in all site collections",
+        
+        # Groups and Teams scopes
+        "Group.Read.All": "Read all groups",
+        "Group.ReadWrite.All": "Read and write all groups",
+        "Team.ReadBasic.All": "Read names and descriptions of teams",
+        "TeamSettings.Read.All": "Read all teams' settings",
+        "TeamSettings.ReadWrite.All": "Read and write all teams' settings",
+        
+        # Directory scopes
+        "Directory.Read.All": "Read directory data",
+        "Directory.ReadWrite.All": "Read and write directory data",
+        
+        # Application scopes
+        "Application.Read.All": "Read all applications",
+        "Application.ReadWrite.All": "Read and write all applications",
+        
+        # OpenID Connect scopes
+        "openid": "OpenID Connect sign-in",
+        "email": "View user's email address",
+        "profile": "View user's basic profile",
+        "offline_access": "Maintain access to data you have given access to"
+    }
+
+    def __init__(
+        self,
+        tenant_id: str,
+        client_id: str,
+        client_secret: str,
+        scopes: Optional[List[str]] = None,
+        use_session_state: bool = True
+    ):
+        """Initialize the Microsoft OAuth2 credential service.
+        
+        Args:
+            tenant_id: Azure AD tenant ID or "common" for multi-tenant.
+            client_id: Microsoft OAuth2 client ID.
+            client_secret: Microsoft OAuth2 client secret.
+            scopes: List of OAuth2 scopes to request.
+            use_session_state: Whether to use session state for credential storage.
+        """
+        super().__init__()
+        self.tenant_id = tenant_id
+        self.client_id = client_id
+        self.client_secret = client_secret
+        self.scopes = scopes or ["User.Read", "Mail.Read"]
+        self.use_session_state = use_session_state
+        
+        # Underlying credential service for storage
+        if use_session_state:
+            self._storage_service = SessionStateCredentialService()
+        else:
+            from google.adk.auth.credential_service.in_memory_credential_service import InMemoryCredentialService
+            self._storage_service = InMemoryCredentialService()
+
+    async def _initialize_impl(self) -> None:
+        """Initialize the Microsoft OAuth2 credential service.
+        
+        Validates the client credentials and sets up the OAuth2 auth scheme.
+        
+        Raises:
+            ValueError: If required parameters are missing.
+        """
+        if not self.tenant_id:
+            raise ValueError("Microsoft OAuth2 tenant_id is required")
+        if not self.client_id:
+            raise ValueError("Microsoft OAuth2 client_id is required")
+        if not self.client_secret:
+            raise ValueError("Microsoft OAuth2 client_secret is required")
+        if not self.scopes:
+            raise ValueError("At least one OAuth2 scope is required")
+            
+        # Validate scopes against known Microsoft scopes
+        unknown_scopes = set(self.scopes) - set(self.COMMON_SCOPES.keys())
+        if unknown_scopes:
+            logger.warning(f"Unknown Microsoft OAuth2 scopes: {unknown_scopes}")
+            
+        logger.info(f"Initialized Microsoft OAuth2 credential service for tenant {self.tenant_id} with scopes: {self.scopes}")
+
+    def create_auth_config(self) -> AuthConfig:
+        """Create an AuthConfig for Microsoft OAuth2 authentication.
+        
+        Returns:
+            AuthConfig: Configured auth config for Microsoft OAuth2 flow.
+        """
+        self._check_initialized()
+        
+        # Create OAuth2 auth scheme
+        auth_scheme = OAuth2(
+            flows=OAuthFlows(
+                authorizationCode=OAuthFlowAuthorizationCode(
+                    authorizationUrl=self._get_auth_url(self.tenant_id),
+                    tokenUrl=self._get_token_url(self.tenant_id),
+                    scopes={
+                        scope: self.COMMON_SCOPES.get(scope, f"Microsoft Graph scope: {scope}")
+                        for scope in self.scopes
+                    }
+                )
+            )
+        )
+        
+        # Create OAuth2 credential
+        auth_credential = AuthCredential(
+            auth_type=AuthCredentialTypes.OAUTH2,
+            oauth2=OAuth2Auth(
+                client_id=self.client_id,
+                client_secret=self.client_secret
+            )
+        )
+        
+        return AuthConfig(
+            auth_scheme=auth_scheme,
+            raw_auth_credential=auth_credential
+        )
+
+    async def load_credential(
+        self,
+        auth_config: AuthConfig,
+        callback_context: CallbackContext,
+    ) -> Optional[AuthCredential]:
+        """Load Microsoft OAuth2 credential from storage.
+        
+        Args:
+            auth_config: The auth config containing credential key information.
+            callback_context: The current callback context.
+            
+        Returns:
+            Optional[AuthCredential]: The stored credential or None if not found.
+        """
+        self._check_initialized()
+        return await self._storage_service.load_credential(auth_config, callback_context)
+
+    async def save_credential(
+        self,
+        auth_config: AuthConfig,
+        callback_context: CallbackContext,
+    ) -> None:
+        """Save Microsoft OAuth2 credential to storage.
+        
+        Args:
+            auth_config: The auth config containing the credential to save.
+            callback_context: The current callback context.
+        """
+        self._check_initialized()
+        await self._storage_service.save_credential(auth_config, callback_context)
+        
+        logger.info(f"Saved Microsoft OAuth2 credential for user {callback_context._invocation_context.user_id} in tenant {self.tenant_id}")
+
+    def get_supported_scopes(self) -> dict:
+        """Get dictionary of supported Microsoft OAuth2 scopes and their descriptions.
+        
+        Returns:
+            dict: Mapping of scope names to descriptions.
+        """
+        return self.COMMON_SCOPES.copy()