Skip to content

Commit a9d9b0c

Browse files
committed
feat: Implement core data classes
1 parent 661f62f commit a9d9b0c

File tree

2 files changed

+170
-0
lines changed

2 files changed

+170
-0
lines changed
Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from dataclasses import dataclass
16+
from typing import Any, Optional, List, Dict
17+
from enum import Enum
18+
19+
try:
20+
from google.auth import credentials as google_creds
21+
except ImportError:
22+
google_creds = None # pragma: no cover
23+
24+
25+
class CredentialType(Enum):
26+
TOOLBOX_IDENTITY = "TOOLBOX_IDENTITY"
27+
APPLICATION_DEFAULT_CREDENTIALS = "APPLICATION_DEFAULT_CREDENTIALS"
28+
USER_IDENTITY = "USER_IDENTITY"
29+
MANUAL_TOKEN = "MANUAL_TOKEN"
30+
MANUAL_CREDS = "MANUAL_CREDS"
31+
32+
33+
@dataclass
34+
class CredentialConfig:
35+
type: CredentialType
36+
# For APPLICATION_DEFAULT_CREDENTIALS
37+
target_audience: Optional[str] = None
38+
# For USER_IDENTITY
39+
client_id: Optional[str] = None
40+
client_secret: Optional[str] = None
41+
scopes: Optional[List[str]] = None
42+
# For MANUAL_TOKEN
43+
token: Optional[str] = None
44+
scheme: Optional[str] = None
45+
# For MANUAL_CREDS
46+
credentials: Optional[Any] = None # google.auth.credentials.Credentials
47+
48+
49+
class CredentialStrategy:
50+
"""Factory for creating credential configurations for ToolboxToolset."""
51+
52+
@staticmethod
53+
def TOOLBOX_IDENTITY() -> CredentialConfig:
54+
"""
55+
No credentials are sent. Relies on the remote Toolbox server's own identity.
56+
"""
57+
return CredentialConfig(type=CredentialType.TOOLBOX_IDENTITY)
58+
59+
@staticmethod
60+
def APPLICATION_DEFAULT_CREDENTIALS(target_audience: str) -> CredentialConfig:
61+
"""
62+
Uses the agent ADC to generate a Google-signed ID token for the specified audience.
63+
This is suitable for Cloud Run, GKE, or local development with `gcloud auth login`.
64+
"""
65+
return CredentialConfig(
66+
type=CredentialType.APPLICATION_DEFAULT_CREDENTIALS,
67+
target_audience=target_audience,
68+
)
69+
70+
@staticmethod
71+
def WORKLOAD_IDENTITY(target_audience: str) -> CredentialConfig:
72+
"""
73+
Alias for APPLICATION_DEFAULT_CREDENTIALS.
74+
"""
75+
return CredentialStrategy.APPLICATION_DEFAULT_CREDENTIALS(target_audience)
76+
77+
@staticmethod
78+
def USER_IDENTITY(
79+
client_id: str, client_secret: str, scopes: Optional[List[str]] = None
80+
) -> CredentialConfig:
81+
"""
82+
Configures the ADK-native interactive 3-legged OAuth flow to get consent
83+
and credentials from the end-user at runtime.
84+
"""
85+
return CredentialConfig(
86+
type=CredentialType.USER_IDENTITY,
87+
client_id=client_id,
88+
client_secret=client_secret,
89+
scopes=scopes,
90+
)
91+
92+
@staticmethod
93+
def MANUAL_TOKEN(token: str, scheme: str = "Bearer") -> CredentialConfig:
94+
"""
95+
Send a hardcoded token string in the Authorization header.
96+
"""
97+
return CredentialConfig(
98+
type=CredentialType.MANUAL_TOKEN,
99+
token=token,
100+
scheme=scheme,
101+
)
102+
103+
@staticmethod
104+
def MANUAL_CREDS(credentials: Any) -> CredentialConfig:
105+
"""
106+
Uses a provided Google Auth Credentials object.
107+
"""
108+
return CredentialConfig(
109+
type=CredentialType.MANUAL_CREDS,
110+
credentials=credentials,
111+
)
Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from toolbox_adk.credentials import CredentialStrategy, CredentialType
16+
17+
class TestCredentialStrategy:
18+
def test_toolbox_identity(self):
19+
config = CredentialStrategy.TOOLBOX_IDENTITY()
20+
assert config.type == CredentialType.TOOLBOX_IDENTITY
21+
22+
def test_application_default_credentials(self):
23+
audience = "https://example.com"
24+
config = CredentialStrategy.APPLICATION_DEFAULT_CREDENTIALS(audience)
25+
assert config.type == CredentialType.APPLICATION_DEFAULT_CREDENTIALS
26+
assert config.target_audience == audience
27+
28+
def test_workload_identity_alias(self):
29+
audience = "https://example.com"
30+
config = CredentialStrategy.WORKLOAD_IDENTITY(audience)
31+
assert config.type == CredentialType.APPLICATION_DEFAULT_CREDENTIALS
32+
assert config.target_audience == audience
33+
34+
def test_user_identity(self):
35+
config = CredentialStrategy.USER_IDENTITY(
36+
client_id="id",
37+
client_secret="secret",
38+
scopes=["scope1"]
39+
)
40+
assert config.type == CredentialType.USER_IDENTITY
41+
assert config.client_id == "id"
42+
assert config.client_secret == "secret"
43+
assert config.scopes == ["scope1"]
44+
45+
def test_manual_token(self):
46+
config = CredentialStrategy.MANUAL_TOKEN(token="abc", scheme="Custom")
47+
assert config.type == CredentialType.MANUAL_TOKEN
48+
assert config.token == "abc"
49+
assert config.scheme == "Custom"
50+
51+
def test_manual_token_defaults(self):
52+
config = CredentialStrategy.MANUAL_TOKEN(token="abc")
53+
assert config.scheme == "Bearer"
54+
55+
def test_manual_creds(self):
56+
fake_creds = object()
57+
config = CredentialStrategy.MANUAL_CREDS(fake_creds)
58+
assert config.type == CredentialType.MANUAL_CREDS
59+
assert config.credentials == fake_creds

0 commit comments

Comments
 (0)