222 lines
6.5 KiB
Python
222 lines
6.5 KiB
Python
# Copyright 2025 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from enum import Enum
|
|
from typing import Any
|
|
from typing import Dict
|
|
from typing import List
|
|
from typing import Optional
|
|
|
|
from pydantic import BaseModel
|
|
from pydantic import Field
|
|
|
|
|
|
class BaseModelWithConfig(BaseModel):
|
|
model_config = {"extra": "allow"}
|
|
|
|
|
|
class HttpCredentials(BaseModelWithConfig):
|
|
"""Represents the secret token value for HTTP authentication, like user name, password, oauth token, etc."""
|
|
|
|
username: Optional[str] = None
|
|
password: Optional[str] = None
|
|
token: Optional[str] = None
|
|
|
|
@classmethod
|
|
def model_validate(cls, data: Dict[str, Any]) -> "HttpCredentials":
|
|
return cls(
|
|
username=data.get("username"),
|
|
password=data.get("password"),
|
|
token=data.get("token"),
|
|
)
|
|
|
|
|
|
class HttpAuth(BaseModelWithConfig):
|
|
"""The credentials and metadata for HTTP authentication."""
|
|
|
|
# The name of the HTTP Authorization scheme to be used in the Authorization
|
|
# header as defined in RFC7235. The values used SHOULD be registered in the
|
|
# IANA Authentication Scheme registry.
|
|
# Examples: 'basic', 'bearer'
|
|
scheme: str
|
|
credentials: HttpCredentials
|
|
|
|
|
|
class OAuth2Auth(BaseModelWithConfig):
|
|
"""Represents credential value and its metadata for a OAuth2 credential."""
|
|
|
|
client_id: Optional[str] = None
|
|
client_secret: Optional[str] = None
|
|
# tool or adk can generate the auth_uri with the state info thus client
|
|
# can verify the state
|
|
auth_uri: Optional[str] = None
|
|
state: Optional[str] = None
|
|
# tool or adk can decide the redirect_uri if they don't want client to decide
|
|
redirect_uri: Optional[str] = None
|
|
auth_response_uri: Optional[str] = None
|
|
auth_code: Optional[str] = None
|
|
access_token: Optional[str] = None
|
|
refresh_token: Optional[str] = None
|
|
|
|
|
|
class ServiceAccountCredential(BaseModelWithConfig):
|
|
"""Represents Google Service Account configuration.
|
|
|
|
Attributes:
|
|
type: The type should be "service_account".
|
|
project_id: The project ID.
|
|
private_key_id: The ID of the private key.
|
|
private_key: The private key.
|
|
client_email: The client email.
|
|
client_id: The client ID.
|
|
auth_uri: The authorization URI.
|
|
token_uri: The token URI.
|
|
auth_provider_x509_cert_url: URL for auth provider's X.509 cert.
|
|
client_x509_cert_url: URL for the client's X.509 cert.
|
|
universe_domain: The universe domain.
|
|
|
|
Example:
|
|
|
|
config = ServiceAccountCredential(
|
|
type_="service_account",
|
|
project_id="your_project_id",
|
|
private_key_id="your_private_key_id",
|
|
private_key="-----BEGIN PRIVATE KEY-----...",
|
|
client_email="...@....iam.gserviceaccount.com",
|
|
client_id="your_client_id",
|
|
auth_uri="https://accounts.google.com/o/oauth2/auth",
|
|
token_uri="https://oauth2.googleapis.com/token",
|
|
auth_provider_x509_cert_url="https://www.googleapis.com/oauth2/v1/certs",
|
|
client_x509_cert_url="https://www.googleapis.com/robot/v1/metadata/x509/...",
|
|
universe_domain="googleapis.com"
|
|
)
|
|
|
|
|
|
config = ServiceAccountConfig.model_construct(**{
|
|
...service account config dict
|
|
})
|
|
"""
|
|
|
|
type_: str = Field("", alias="type")
|
|
project_id: str
|
|
private_key_id: str
|
|
private_key: str
|
|
client_email: str
|
|
client_id: str
|
|
auth_uri: str
|
|
token_uri: str
|
|
auth_provider_x509_cert_url: str
|
|
client_x509_cert_url: str
|
|
universe_domain: str
|
|
|
|
|
|
class ServiceAccount(BaseModelWithConfig):
|
|
"""Represents Google Service Account configuration."""
|
|
|
|
service_account_credential: Optional[ServiceAccountCredential] = None
|
|
scopes: List[str]
|
|
use_default_credential: Optional[bool] = False
|
|
|
|
|
|
class AuthCredentialTypes(str, Enum):
|
|
"""Represents the type of authentication credential."""
|
|
|
|
# API Key credential:
|
|
# https://swagger.io/docs/specification/v3_0/authentication/api-keys/
|
|
API_KEY = "apiKey"
|
|
|
|
# Credentials for HTTP Auth schemes:
|
|
# https://www.iana.org/assignments/http-authschemes/http-authschemes.xhtml
|
|
HTTP = "http"
|
|
|
|
# OAuth2 credentials:
|
|
# https://swagger.io/docs/specification/v3_0/authentication/oauth2/
|
|
OAUTH2 = "oauth2"
|
|
|
|
# OpenID Connect credentials:
|
|
# https://swagger.io/docs/specification/v3_0/authentication/openid-connect-discovery/
|
|
OPEN_ID_CONNECT = "openIdConnect"
|
|
|
|
# Service Account credentials:
|
|
# https://cloud.google.com/iam/docs/service-account-creds
|
|
SERVICE_ACCOUNT = "serviceAccount"
|
|
|
|
|
|
class AuthCredential(BaseModelWithConfig):
|
|
"""Data class representing an authentication credential.
|
|
|
|
To exchange for the actual credential, please use
|
|
CredentialExchanger.exchange_credential().
|
|
|
|
Examples: API Key Auth
|
|
AuthCredential(
|
|
auth_type=AuthCredentialTypes.API_KEY,
|
|
api_key="1234",
|
|
)
|
|
|
|
Example: HTTP Auth
|
|
AuthCredential(
|
|
auth_type=AuthCredentialTypes.HTTP,
|
|
http=HttpAuth(
|
|
scheme="basic",
|
|
credentials=HttpCredentials(username="user", password="password"),
|
|
),
|
|
)
|
|
|
|
Example: OAuth2 Bearer Token in HTTP Header
|
|
AuthCredential(
|
|
auth_type=AuthCredentialTypes.HTTP,
|
|
http=HttpAuth(
|
|
scheme="bearer",
|
|
credentials=HttpCredentials(token="eyAkaknabna...."),
|
|
),
|
|
)
|
|
|
|
Example: OAuth2 Auth with Authorization Code Flow
|
|
AuthCredential(
|
|
auth_type=AuthCredentialTypes.OAUTH2,
|
|
oauth2=OAuth2Auth(
|
|
client_id="1234",
|
|
client_secret="secret",
|
|
),
|
|
)
|
|
|
|
Example: OpenID Connect Auth
|
|
AuthCredential(
|
|
auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
|
|
oauth2=OAuth2Auth(
|
|
client_id="1234",
|
|
client_secret="secret",
|
|
redirect_uri="https://example.com",
|
|
scopes=["scope1", "scope2"],
|
|
),
|
|
)
|
|
|
|
Example: Auth with resource reference
|
|
AuthCredential(
|
|
auth_type=AuthCredentialTypes.API_KEY,
|
|
resource_ref="projects/1234/locations/us-central1/resources/resource1",
|
|
)
|
|
"""
|
|
|
|
auth_type: AuthCredentialTypes
|
|
# Resource reference for the credential.
|
|
# This will be supported in the future.
|
|
resource_ref: Optional[str] = None
|
|
|
|
api_key: Optional[str] = None
|
|
http: Optional[HttpAuth] = None
|
|
service_account: Optional[ServiceAccount] = None
|
|
oauth2: Optional[OAuth2Auth] = None
|