mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
chore: Add sample agent for testing oatuh2 client credentials grant type
PiperOrigin-RevId: 815863131
This commit is contained in:
committed by
Copybara-Service
parent
4bb089d386
commit
90e1e3e10c
@@ -0,0 +1,143 @@
|
|||||||
|
# OAuth2 Client Credentials Weather Agent
|
||||||
|
|
||||||
|
This sample demonstrates OAuth2 client credentials flow with ADK's `AuthenticatedFunctionTool` using a practical weather assistant agent.
|
||||||
|
|
||||||
|
## Overview
|
||||||
|
|
||||||
|
The OAuth2 client credentials grant type is used for server-to-server authentication where no user interaction is required. This demo shows:
|
||||||
|
|
||||||
|
- How to configure OAuth2 client credentials in ADK
|
||||||
|
- Using `AuthenticatedFunctionTool` for automatic token management
|
||||||
|
- Transparent authentication in a practical weather assistant
|
||||||
|
- Testing the OAuth2 client credentials implementation
|
||||||
|
|
||||||
|
## Architecture
|
||||||
|
|
||||||
|
```
|
||||||
|
[WeatherAssistant] -> [AuthenticatedFunctionTool] -> [OAuth2CredentialExchanger] -> [OAuth2 Server] -> [Weather API]
|
||||||
|
```
|
||||||
|
|
||||||
|
1. **WeatherAssistant** calls weather tool when user asks for weather data
|
||||||
|
2. **AuthenticatedFunctionTool** automatically handles OAuth2 flow
|
||||||
|
3. **OAuth2CredentialExchanger** exchanges client credentials for access token
|
||||||
|
4. **Authenticated requests** are made to weather API
|
||||||
|
|
||||||
|
## Files
|
||||||
|
|
||||||
|
### `agent.py` - WeatherAssistant Agent
|
||||||
|
|
||||||
|
Weather assistant agent that demonstrates OAuth2 client credentials flow transparently:
|
||||||
|
|
||||||
|
- **OAuth2 Configuration**: Client credentials setup with token URL and scopes
|
||||||
|
- **Weather Tool**: Single `get_weather_data` tool for fetching weather information
|
||||||
|
- **Agent Definition**: ADK LLM agent focused on providing weather information
|
||||||
|
|
||||||
|
**Key Features:**
|
||||||
|
- Automatic token exchange using client ID and secret
|
||||||
|
- Bearer token authentication
|
||||||
|
- Transparent OAuth2 handling (invisible to the model)
|
||||||
|
- Practical use case demonstrating machine-to-machine authentication
|
||||||
|
|
||||||
|
### `main.py` - CLI Interface
|
||||||
|
|
||||||
|
Command-line interface for running the WeatherAssistant agent:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Ask for weather
|
||||||
|
python contributing/samples/oauth2_client_credentials/main.py "What's the weather in Tokyo?"
|
||||||
|
```
|
||||||
|
|
||||||
|
**Requirements:**
|
||||||
|
- LLM API key (Google AI or Vertex AI)
|
||||||
|
- OAuth2 test server running
|
||||||
|
|
||||||
|
### `oauth2_test_server.py` - Local OAuth2 Server
|
||||||
|
|
||||||
|
Mock OAuth2 server for testing the client credentials flow:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python contributing/samples/oauth2_client_credentials/oauth2_test_server.py
|
||||||
|
```
|
||||||
|
|
||||||
|
**Features:**
|
||||||
|
- OIDC discovery endpoint (`/.well-known/openid_configuration`)
|
||||||
|
- Client credentials token exchange (`/token`)
|
||||||
|
- Protected weather API (`/api/weather`)
|
||||||
|
- Supports both `authorization_code` and `client_credentials` grant types
|
||||||
|
- Test credentials: `client_id="test_client"`, `client_secret="test_secret"`
|
||||||
|
|
||||||
|
**Endpoints:**
|
||||||
|
- `GET /.well-known/openid_configuration` - OIDC discovery
|
||||||
|
- `POST /token` - Token exchange
|
||||||
|
- `GET /api/weather` - Weather API (requires Bearer token)
|
||||||
|
- `GET /` - Server info
|
||||||
|
|
||||||
|
## Quick Start
|
||||||
|
|
||||||
|
1. **Start the OAuth2 server:**
|
||||||
|
```bash
|
||||||
|
python contributing/samples/oauth2_client_credentials/oauth2_test_server.py &
|
||||||
|
```
|
||||||
|
2. Create a `.env` file in the project root with your API credentials:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Choose Model Backend: 0 -> ML Dev, 1 -> Vertex
|
||||||
|
GOOGLE_GENAI_USE_VERTEXAI=1
|
||||||
|
|
||||||
|
# ML Dev backend config
|
||||||
|
GOOGLE_API_KEY=your_google_api_key_here
|
||||||
|
|
||||||
|
# Vertex backend config
|
||||||
|
GOOGLE_CLOUD_PROJECT=your_project_id
|
||||||
|
GOOGLE_CLOUD_LOCATION=us-central1
|
||||||
|
```
|
||||||
|
|
||||||
|
3. **Run the agent:**
|
||||||
|
```bash
|
||||||
|
# Ask for weather
|
||||||
|
python contributing/samples/oauth2_client_credentials/main.py "What's the weather in Tokyo?"
|
||||||
|
```
|
||||||
|
|
||||||
|
3. **Interactive demo (use ADK commands):**
|
||||||
|
```bash
|
||||||
|
# Interactive CLI
|
||||||
|
adk run contributing/samples/oauth2_client_credentials
|
||||||
|
|
||||||
|
# Interactive web UI
|
||||||
|
adk web contributing/samples
|
||||||
|
```
|
||||||
|
|
||||||
|
## OAuth2 Configuration
|
||||||
|
|
||||||
|
The agent uses these OAuth2 settings (configured in `agent.py`):
|
||||||
|
|
||||||
|
```python
|
||||||
|
flows = OAuthFlows(
|
||||||
|
clientCredentials=OAuthFlowClientCredentials(
|
||||||
|
tokenUrl="http://localhost:8000/token",
|
||||||
|
scopes={
|
||||||
|
"read": "Read access to weather data",
|
||||||
|
"write": "Write access for data updates",
|
||||||
|
"admin": "Administrative access",
|
||||||
|
},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
raw_credential = AuthCredential(
|
||||||
|
auth_type=AuthCredentialTypes.OAUTH2,
|
||||||
|
oauth2=OAuth2Auth(
|
||||||
|
client_id="test_client",
|
||||||
|
client_secret="test_secret",
|
||||||
|
),
|
||||||
|
)
|
||||||
|
```
|
||||||
|
|
||||||
|
## Authentication Flow
|
||||||
|
|
||||||
|
1. **Weather Request**: User asks WeatherAssistant for weather information
|
||||||
|
2. **Tool Invocation**: Agent calls `get_weather_data` authenticated function tool
|
||||||
|
3. **Credential Loading**: CredentialManager loads OAuth2 configuration
|
||||||
|
4. **Token Exchange**: OAuth2CredentialExchanger uses client credentials to get access token
|
||||||
|
5. **Request Enhancement**: AuthenticatedFunctionTool adds `Authorization: Bearer <token>` header
|
||||||
|
6. **API Call**: Weather API accessed with valid token
|
||||||
|
7. **Response**: Weather data returned to user
|
||||||
@@ -0,0 +1,15 @@
|
|||||||
|
# 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 . import agent
|
||||||
@@ -0,0 +1,134 @@
|
|||||||
|
# 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.
|
||||||
|
|
||||||
|
"""Weather Assistant Agent.
|
||||||
|
|
||||||
|
This agent provides weather information for cities worldwide.
|
||||||
|
It demonstrates OAuth2 client credentials flow transparently
|
||||||
|
through AuthenticatedFunctionTool usage.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from fastapi.openapi.models import OAuth2
|
||||||
|
from fastapi.openapi.models import OAuthFlowClientCredentials
|
||||||
|
from fastapi.openapi.models import OAuthFlows
|
||||||
|
from google.adk.agents.llm_agent import Agent
|
||||||
|
from google.adk.auth.auth_credential import AuthCredential
|
||||||
|
from google.adk.auth.auth_credential import AuthCredentialTypes
|
||||||
|
from google.adk.auth.auth_credential import OAuth2Auth
|
||||||
|
from google.adk.auth.auth_tool import AuthConfig
|
||||||
|
from google.adk.tools.authenticated_function_tool import AuthenticatedFunctionTool
|
||||||
|
import requests
|
||||||
|
|
||||||
|
|
||||||
|
# OAuth2 configuration for weather API access
|
||||||
|
def create_auth_config() -> AuthConfig:
|
||||||
|
"""Create OAuth2 auth configuration for weather API."""
|
||||||
|
|
||||||
|
# Define OAuth2 scheme with client credentials flow
|
||||||
|
flows = OAuthFlows(
|
||||||
|
clientCredentials=OAuthFlowClientCredentials(
|
||||||
|
tokenUrl="http://localhost:8080/token",
|
||||||
|
scopes={
|
||||||
|
"read": "Read access to weather data",
|
||||||
|
"write": "Write access for data updates",
|
||||||
|
"admin": "Administrative access",
|
||||||
|
},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
auth_scheme = OAuth2(flows=flows)
|
||||||
|
|
||||||
|
# Create credential with client ID and secret
|
||||||
|
raw_credential = AuthCredential(
|
||||||
|
auth_type=AuthCredentialTypes.OAUTH2,
|
||||||
|
oauth2=OAuth2Auth(
|
||||||
|
client_id="test_client",
|
||||||
|
client_secret="test_secret",
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
return AuthConfig(
|
||||||
|
auth_scheme=auth_scheme,
|
||||||
|
raw_auth_credential=raw_credential,
|
||||||
|
credential_key="weather_api_client",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def get_weather_data(city: str = "San Francisco", credential=None) -> str:
|
||||||
|
"""Get current weather data for a specified city.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
city: City name to get weather for
|
||||||
|
credential: API credential (automatically injected by AuthenticatedFunctionTool)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Current weather information for the city.
|
||||||
|
"""
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Use the credential to make authenticated requests to weather API
|
||||||
|
headers = {}
|
||||||
|
if credential and credential.oauth2 and credential.oauth2.access_token:
|
||||||
|
headers["Authorization"] = f"Bearer {credential.oauth2.access_token}"
|
||||||
|
|
||||||
|
# Call weather API endpoint
|
||||||
|
params = {"city": city, "units": "metric"}
|
||||||
|
response = requests.get(
|
||||||
|
"http://localhost:8080/api/weather",
|
||||||
|
headers=headers,
|
||||||
|
params=params,
|
||||||
|
timeout=10,
|
||||||
|
)
|
||||||
|
|
||||||
|
if response.status_code == 200:
|
||||||
|
data = response.json()
|
||||||
|
result = f"🌤️ Weather for {city}:\n"
|
||||||
|
result += f"Temperature: {data.get('temperature', 'N/A')}°C\n"
|
||||||
|
result += f"Condition: {data.get('condition', 'N/A')}\n"
|
||||||
|
result += f"Humidity: {data.get('humidity', 'N/A')}%\n"
|
||||||
|
result += f"Wind Speed: {data.get('wind_speed', 'N/A')} km/h\n"
|
||||||
|
result += f"Last Updated: {data.get('timestamp', 'N/A')}\n"
|
||||||
|
return result
|
||||||
|
else:
|
||||||
|
return (
|
||||||
|
f"❌ Failed to get weather data: {response.status_code} -"
|
||||||
|
f" {response.text}"
|
||||||
|
)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
return f"❌ Error getting weather data: {str(e)}"
|
||||||
|
|
||||||
|
|
||||||
|
# Create the weather assistant agent
|
||||||
|
root_agent = Agent(
|
||||||
|
name="WeatherAssistant",
|
||||||
|
description=(
|
||||||
|
"Weather assistant that provides current weather information for cities"
|
||||||
|
" worldwide."
|
||||||
|
),
|
||||||
|
model="gemini-2.5-pro",
|
||||||
|
instruction=(
|
||||||
|
"You are a helpful Weather Assistant that provides current weather"
|
||||||
|
" information for any city worldwide.\n\nWhen users ask for weather:\n•"
|
||||||
|
" Ask for the city name if not provided\n• Provide temperature in"
|
||||||
|
" Celsius\n• Include helpful details like humidity, wind speed, and"
|
||||||
|
" conditions\n• Be friendly and conversational about the weather\n\nIf"
|
||||||
|
" there are any issues getting weather data, apologize and suggest"
|
||||||
|
" trying again or checking for a different city name."
|
||||||
|
),
|
||||||
|
tools=[
|
||||||
|
AuthenticatedFunctionTool(
|
||||||
|
func=get_weather_data, auth_config=create_auth_config()
|
||||||
|
),
|
||||||
|
],
|
||||||
|
)
|
||||||
@@ -0,0 +1,152 @@
|
|||||||
|
"""WeatherAssistant Agent main script.
|
||||||
|
|
||||||
|
This script demonstrates OAuth2 client credentials flow using a practical
|
||||||
|
weather assistant agent with AuthenticatedFunctionTool.
|
||||||
|
"""
|
||||||
|
|
||||||
|
# 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.
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import asyncio
|
||||||
|
import logging
|
||||||
|
import sys
|
||||||
|
import time
|
||||||
|
|
||||||
|
import agent
|
||||||
|
from dotenv import load_dotenv
|
||||||
|
from google.adk.cli.utils import logs
|
||||||
|
from google.adk.runners import InMemoryRunner
|
||||||
|
|
||||||
|
APP_NAME = "weather_assistant_app"
|
||||||
|
USER_ID = "weather_user"
|
||||||
|
|
||||||
|
logs.setup_adk_logger(level=logging.INFO)
|
||||||
|
|
||||||
|
|
||||||
|
def process_arguments():
|
||||||
|
"""Parses command-line arguments."""
|
||||||
|
parser = argparse.ArgumentParser(
|
||||||
|
description=(
|
||||||
|
"WeatherAssistant Agent - demonstrates OAuth2 client credentials"
|
||||||
|
" authentication transparently through weather queries."
|
||||||
|
),
|
||||||
|
epilog=(
|
||||||
|
"Example usage:\n\tpython main.py"
|
||||||
|
' "What\'s the weather in Tokyo?"\n\n'
|
||||||
|
"For interactive usage, use ADK commands:\n"
|
||||||
|
"\tadk run .\n"
|
||||||
|
"\tadk web .\n"
|
||||||
|
),
|
||||||
|
formatter_class=argparse.RawTextHelpFormatter,
|
||||||
|
)
|
||||||
|
|
||||||
|
parser.add_argument(
|
||||||
|
"message",
|
||||||
|
type=str,
|
||||||
|
help=(
|
||||||
|
"Ask the weather assistant a question or request weather information."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
async def process_message(runner, session_id, message):
|
||||||
|
"""Process a single message with the weather assistant."""
|
||||||
|
print(f"🌤️ Weather Assistant: ")
|
||||||
|
|
||||||
|
response = await call_agent_async(runner, USER_ID, session_id, message)
|
||||||
|
print(f"{response}\n")
|
||||||
|
|
||||||
|
|
||||||
|
async def call_agent_async(runner, user_id, session_id, prompt):
|
||||||
|
"""Helper function to call agent asynchronously."""
|
||||||
|
from google.adk.agents.run_config import RunConfig
|
||||||
|
from google.genai import types
|
||||||
|
|
||||||
|
content = types.Content(
|
||||||
|
role="user", parts=[types.Part.from_text(text=prompt)]
|
||||||
|
)
|
||||||
|
final_response_text = ""
|
||||||
|
|
||||||
|
async for event in runner.run_async(
|
||||||
|
user_id=user_id,
|
||||||
|
session_id=session_id,
|
||||||
|
new_message=content,
|
||||||
|
run_config=RunConfig(save_input_blobs_as_artifacts=False),
|
||||||
|
):
|
||||||
|
if event.content and event.content.parts:
|
||||||
|
if text := "".join(part.text or "" for part in event.content.parts):
|
||||||
|
if event.author != "user":
|
||||||
|
final_response_text += text
|
||||||
|
|
||||||
|
return final_response_text
|
||||||
|
|
||||||
|
|
||||||
|
async def main():
|
||||||
|
"""Main function."""
|
||||||
|
# Load environment variables from .env file
|
||||||
|
load_dotenv()
|
||||||
|
|
||||||
|
args = process_arguments()
|
||||||
|
|
||||||
|
print("🌤️ WeatherAssistant Agent")
|
||||||
|
print("=" * 40)
|
||||||
|
print("Ask me about weather in any city around the world!")
|
||||||
|
print("(OAuth2 client credentials authentication happens transparently)\n")
|
||||||
|
|
||||||
|
# Create runner and session
|
||||||
|
runner = InMemoryRunner(
|
||||||
|
agent=agent.root_agent,
|
||||||
|
app_name=APP_NAME,
|
||||||
|
)
|
||||||
|
|
||||||
|
session = await runner.session_service.create_session(
|
||||||
|
app_name=APP_NAME, user_id=USER_ID
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
await process_message(runner, session.id, args.message)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
print(f"❌ Error: {e}", file=sys.stderr)
|
||||||
|
return 1
|
||||||
|
|
||||||
|
return 0
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
start_time = time.time()
|
||||||
|
print(
|
||||||
|
"⏰ Started at"
|
||||||
|
f" {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(start_time))}"
|
||||||
|
)
|
||||||
|
print("-" * 50)
|
||||||
|
|
||||||
|
try:
|
||||||
|
exit_code = asyncio.run(main())
|
||||||
|
except KeyboardInterrupt:
|
||||||
|
print("\n⏹️ Interrupted by user")
|
||||||
|
exit_code = 1
|
||||||
|
|
||||||
|
end_time = time.time()
|
||||||
|
print("-" * 50)
|
||||||
|
print(
|
||||||
|
"⏰ Finished at"
|
||||||
|
f" {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(end_time))}"
|
||||||
|
)
|
||||||
|
print(f"⌛ Total execution time: {end_time - start_time:.2f} seconds")
|
||||||
|
|
||||||
|
sys.exit(exit_code)
|
||||||
@@ -0,0 +1,350 @@
|
|||||||
|
# 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.
|
||||||
|
"""
|
||||||
|
Weather API OAuth2 Test Server
|
||||||
|
|
||||||
|
A simple FastAPI server that implements OAuth2 flows for weather API testing:
|
||||||
|
- Client Credentials Flow
|
||||||
|
- Authorization Code Flow
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
python oauth2_test_server.py
|
||||||
|
|
||||||
|
Endpoints:
|
||||||
|
GET /auth - Authorization endpoint (auth code flow)
|
||||||
|
POST /token - Token endpoint (both flows)
|
||||||
|
GET /.well-known/openid_configuration - OpenID Connect discovery
|
||||||
|
GET /api/weather - Weather API (requires Bearer token)
|
||||||
|
"""
|
||||||
|
|
||||||
|
import secrets
|
||||||
|
import time
|
||||||
|
from typing import Dict
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
|
from fastapi import FastAPI
|
||||||
|
from fastapi import Form
|
||||||
|
from fastapi import HTTPException
|
||||||
|
from fastapi import Query
|
||||||
|
from fastapi import Request
|
||||||
|
from fastapi import status
|
||||||
|
from fastapi.responses import HTMLResponse
|
||||||
|
from fastapi.responses import RedirectResponse
|
||||||
|
from pydantic import BaseModel
|
||||||
|
|
||||||
|
app = FastAPI(title="Weather API OAuth2 Server", version="1.0.0")
|
||||||
|
|
||||||
|
# In-memory storage (for testing only)
|
||||||
|
clients = {
|
||||||
|
"test_client": {
|
||||||
|
"client_secret": "test_secret",
|
||||||
|
"redirect_uris": [
|
||||||
|
"http://localhost:8080/callback",
|
||||||
|
"urn:ietf:wg:oauth:2.0:oob",
|
||||||
|
],
|
||||||
|
"scopes": ["read", "write", "admin"],
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
authorization_codes = {} # code -> {client_id, redirect_uri, scope, expires_at}
|
||||||
|
access_tokens = {} # token -> {client_id, scope, expires_at, token_type}
|
||||||
|
|
||||||
|
|
||||||
|
class TokenResponse(BaseModel):
|
||||||
|
access_token: str
|
||||||
|
token_type: str = "Bearer"
|
||||||
|
expires_in: int = 3600
|
||||||
|
refresh_token: Optional[str] = None
|
||||||
|
scope: Optional[str] = None
|
||||||
|
|
||||||
|
|
||||||
|
@app.get("/.well-known/openid_configuration")
|
||||||
|
async def openid_configuration():
|
||||||
|
"""OpenID Connect Discovery endpoint."""
|
||||||
|
return {
|
||||||
|
"issuer": "http://localhost:8080",
|
||||||
|
"authorization_endpoint": "http://localhost:8080/auth",
|
||||||
|
"token_endpoint": "http://localhost:8080/token",
|
||||||
|
"userinfo_endpoint": "http://localhost:8080/userinfo",
|
||||||
|
"revocation_endpoint": "http://localhost:8080/revoke",
|
||||||
|
"scopes_supported": ["openid", "read", "write", "admin"],
|
||||||
|
"response_types_supported": ["code"],
|
||||||
|
"grant_types_supported": ["authorization_code", "client_credentials"],
|
||||||
|
"token_endpoint_auth_methods_supported": [
|
||||||
|
"client_secret_basic",
|
||||||
|
"client_secret_post",
|
||||||
|
],
|
||||||
|
"subject_types_supported": ["public"],
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@app.get("/auth")
|
||||||
|
async def authorize(
|
||||||
|
response_type: str = Query(...),
|
||||||
|
client_id: str = Query(...),
|
||||||
|
redirect_uri: str = Query(...),
|
||||||
|
scope: str = Query(default="read"),
|
||||||
|
state: str = Query(default=""),
|
||||||
|
):
|
||||||
|
"""Authorization endpoint for OAuth2 authorization code flow."""
|
||||||
|
|
||||||
|
# Validate client
|
||||||
|
if client_id not in clients:
|
||||||
|
raise HTTPException(status_code=400, detail="Invalid client_id")
|
||||||
|
|
||||||
|
client = clients[client_id]
|
||||||
|
if redirect_uri not in client["redirect_uris"]:
|
||||||
|
raise HTTPException(status_code=400, detail="Invalid redirect_uri")
|
||||||
|
|
||||||
|
if response_type != "code":
|
||||||
|
raise HTTPException(status_code=400, detail="Unsupported response_type")
|
||||||
|
|
||||||
|
# Generate authorization code
|
||||||
|
auth_code = secrets.token_urlsafe(32)
|
||||||
|
authorization_codes[auth_code] = {
|
||||||
|
"client_id": client_id,
|
||||||
|
"redirect_uri": redirect_uri,
|
||||||
|
"scope": scope,
|
||||||
|
"expires_at": time.time() + 600, # 10 minutes
|
||||||
|
}
|
||||||
|
|
||||||
|
# Simulate user consent - in real implementation, this would show a consent form
|
||||||
|
params = f"code={auth_code}"
|
||||||
|
if state:
|
||||||
|
params += f"&state={state}"
|
||||||
|
|
||||||
|
return RedirectResponse(url=f"{redirect_uri}?{params}")
|
||||||
|
|
||||||
|
|
||||||
|
@app.post("/token")
|
||||||
|
async def token_endpoint(
|
||||||
|
request: Request,
|
||||||
|
grant_type: str = Form(...),
|
||||||
|
client_id: str = Form(default=None),
|
||||||
|
client_secret: str = Form(default=None),
|
||||||
|
code: str = Form(default=None),
|
||||||
|
redirect_uri: str = Form(default=None),
|
||||||
|
scope: str = Form(default="read"),
|
||||||
|
):
|
||||||
|
"""Token endpoint for both client credentials and authorization code flows."""
|
||||||
|
|
||||||
|
# Support both HTTP Basic auth and form-based client authentication
|
||||||
|
auth_header = request.headers.get("Authorization")
|
||||||
|
|
||||||
|
if auth_header and auth_header.startswith("Basic "):
|
||||||
|
# HTTP Basic authentication
|
||||||
|
import base64
|
||||||
|
|
||||||
|
try:
|
||||||
|
encoded_credentials = auth_header[6:] # Remove "Basic " prefix
|
||||||
|
decoded = base64.b64decode(encoded_credentials).decode("utf-8")
|
||||||
|
basic_client_id, basic_client_secret = decoded.split(":", 1)
|
||||||
|
client_id = client_id or basic_client_id
|
||||||
|
client_secret = client_secret or basic_client_secret
|
||||||
|
except Exception:
|
||||||
|
raise HTTPException(
|
||||||
|
status_code=401, detail="Invalid authorization header"
|
||||||
|
)
|
||||||
|
|
||||||
|
if not client_id or not client_secret:
|
||||||
|
raise HTTPException(status_code=400, detail="Client credentials required")
|
||||||
|
|
||||||
|
# Validate client credentials
|
||||||
|
if client_id not in clients:
|
||||||
|
raise HTTPException(status_code=401, detail="Invalid client")
|
||||||
|
|
||||||
|
client = clients[client_id]
|
||||||
|
if client["client_secret"] != client_secret:
|
||||||
|
raise HTTPException(status_code=401, detail="Invalid client credentials")
|
||||||
|
|
||||||
|
if grant_type == "client_credentials":
|
||||||
|
return await handle_client_credentials(client_id, scope)
|
||||||
|
elif grant_type == "authorization_code":
|
||||||
|
return await handle_authorization_code(client_id, code, redirect_uri, scope)
|
||||||
|
else:
|
||||||
|
raise HTTPException(status_code=400, detail="Unsupported grant_type")
|
||||||
|
|
||||||
|
|
||||||
|
async def handle_client_credentials(
|
||||||
|
client_id: str, scope: str
|
||||||
|
) -> TokenResponse:
|
||||||
|
"""Handle client credentials flow."""
|
||||||
|
|
||||||
|
# Generate access token
|
||||||
|
access_token = secrets.token_urlsafe(32)
|
||||||
|
expires_at = time.time() + 3600 # 1 hour
|
||||||
|
|
||||||
|
# Store token
|
||||||
|
access_tokens[access_token] = {
|
||||||
|
"client_id": client_id,
|
||||||
|
"scope": scope,
|
||||||
|
"expires_at": expires_at,
|
||||||
|
"token_type": "Bearer",
|
||||||
|
}
|
||||||
|
|
||||||
|
return TokenResponse(
|
||||||
|
access_token=access_token,
|
||||||
|
token_type="Bearer",
|
||||||
|
expires_in=3600,
|
||||||
|
scope=scope,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def handle_authorization_code(
|
||||||
|
client_id: str, code: str, redirect_uri: str, scope: str
|
||||||
|
) -> TokenResponse:
|
||||||
|
"""Handle authorization code flow."""
|
||||||
|
|
||||||
|
if not code:
|
||||||
|
raise HTTPException(status_code=400, detail="Missing authorization code")
|
||||||
|
|
||||||
|
if code not in authorization_codes:
|
||||||
|
raise HTTPException(status_code=400, detail="Invalid authorization code")
|
||||||
|
|
||||||
|
auth_data = authorization_codes[code]
|
||||||
|
|
||||||
|
# Validate authorization code
|
||||||
|
if time.time() > auth_data["expires_at"]:
|
||||||
|
del authorization_codes[code]
|
||||||
|
raise HTTPException(status_code=400, detail="Authorization code expired")
|
||||||
|
|
||||||
|
if auth_data["client_id"] != client_id:
|
||||||
|
raise HTTPException(status_code=400, detail="Client mismatch")
|
||||||
|
|
||||||
|
if redirect_uri and auth_data["redirect_uri"] != redirect_uri:
|
||||||
|
raise HTTPException(status_code=400, detail="Redirect URI mismatch")
|
||||||
|
|
||||||
|
# Generate tokens
|
||||||
|
access_token = secrets.token_urlsafe(32)
|
||||||
|
refresh_token = secrets.token_urlsafe(32)
|
||||||
|
expires_at = time.time() + 3600 # 1 hour
|
||||||
|
|
||||||
|
# Store token
|
||||||
|
access_tokens[access_token] = {
|
||||||
|
"client_id": client_id,
|
||||||
|
"scope": auth_data["scope"],
|
||||||
|
"expires_at": expires_at,
|
||||||
|
"token_type": "Bearer",
|
||||||
|
}
|
||||||
|
|
||||||
|
# Clean up authorization code (one-time use)
|
||||||
|
del authorization_codes[code]
|
||||||
|
|
||||||
|
return TokenResponse(
|
||||||
|
access_token=access_token,
|
||||||
|
token_type="Bearer",
|
||||||
|
expires_in=3600,
|
||||||
|
refresh_token=refresh_token,
|
||||||
|
scope=auth_data["scope"],
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@app.get("/api/weather")
|
||||||
|
async def get_weather(
|
||||||
|
request: Request, city: str = "San Francisco", units: str = "metric"
|
||||||
|
):
|
||||||
|
"""Weather API endpoint that returns weather data for a city."""
|
||||||
|
|
||||||
|
# Check authentication
|
||||||
|
auth_header = request.headers.get("Authorization")
|
||||||
|
if not auth_header or not auth_header.startswith("Bearer "):
|
||||||
|
raise HTTPException(
|
||||||
|
status_code=401, detail="Missing or invalid authorization header"
|
||||||
|
)
|
||||||
|
|
||||||
|
token = auth_header[7:] # Remove "Bearer " prefix
|
||||||
|
|
||||||
|
if token not in access_tokens:
|
||||||
|
raise HTTPException(status_code=401, detail="Invalid access token")
|
||||||
|
|
||||||
|
token_data = access_tokens[token]
|
||||||
|
|
||||||
|
if time.time() > token_data["expires_at"]:
|
||||||
|
del access_tokens[token]
|
||||||
|
raise HTTPException(status_code=401, detail="Access token expired")
|
||||||
|
|
||||||
|
# Return weather data (simulated)
|
||||||
|
from datetime import datetime
|
||||||
|
import random
|
||||||
|
|
||||||
|
conditions = ["Sunny", "Partly Cloudy", "Cloudy", "Light Rain", "Clear"]
|
||||||
|
|
||||||
|
weather_data = {
|
||||||
|
"city": city,
|
||||||
|
"temperature": random.randint(15, 30),
|
||||||
|
"condition": random.choice(conditions),
|
||||||
|
"humidity": random.randint(40, 80),
|
||||||
|
"wind_speed": random.randint(5, 25),
|
||||||
|
"timestamp": datetime.now().isoformat(),
|
||||||
|
"units": units,
|
||||||
|
"api_client": token_data["client_id"],
|
||||||
|
}
|
||||||
|
|
||||||
|
return weather_data
|
||||||
|
|
||||||
|
|
||||||
|
@app.get("/")
|
||||||
|
async def root():
|
||||||
|
"""Root endpoint with server information."""
|
||||||
|
return HTMLResponse("""
|
||||||
|
<html>
|
||||||
|
<head><title>Weather API OAuth2 Server</title></head>
|
||||||
|
<body>
|
||||||
|
<h1>Weather API OAuth2 Server</h1>
|
||||||
|
<h2>Available Endpoints:</h2>
|
||||||
|
<ul>
|
||||||
|
<li><strong>GET /auth</strong> - Authorization endpoint</li>
|
||||||
|
<li><strong>POST /token</strong> - Token endpoint</li>
|
||||||
|
<li><strong>GET /.well-known/openid_configuration</strong> - Discovery</li>
|
||||||
|
<li><strong>GET /api/weather</strong> - Weather API (requires Bearer token)</li>
|
||||||
|
</ul>
|
||||||
|
|
||||||
|
<h2>Test Client Credentials:</h2>
|
||||||
|
<ul>
|
||||||
|
<li><strong>Client ID:</strong> test_client</li>
|
||||||
|
<li><strong>Client Secret:</strong> test_secret</li>
|
||||||
|
<li><strong>Scopes:</strong> read, write, admin</li>
|
||||||
|
</ul>
|
||||||
|
|
||||||
|
<h2>Example cURL Commands:</h2>
|
||||||
|
<h3>Client Credentials Flow:</h3>
|
||||||
|
<pre>
|
||||||
|
curl -X POST http://localhost:8080/token \\
|
||||||
|
-d "grant_type=client_credentials" \\
|
||||||
|
-d "client_id=test_client" \\
|
||||||
|
-d "client_secret=test_secret" \\
|
||||||
|
-d "scope=read write"
|
||||||
|
</pre>
|
||||||
|
|
||||||
|
<h3>Test Weather API:</h3>
|
||||||
|
<pre>
|
||||||
|
curl -H "Authorization: Bearer YOUR_ACCESS_TOKEN" \\
|
||||||
|
"http://localhost:8080/api/weather?city=Tokyo"
|
||||||
|
</pre>
|
||||||
|
</body>
|
||||||
|
</html>
|
||||||
|
""")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
import uvicorn
|
||||||
|
|
||||||
|
print("🌤️ Starting Weather API OAuth2 Server...")
|
||||||
|
print("📖 Documentation: http://localhost:8080/docs")
|
||||||
|
print("🏠 Server Info: http://localhost:8080")
|
||||||
|
print(
|
||||||
|
'🔧 Test with: curl -H "Authorization: Bearer TOKEN"'
|
||||||
|
' "http://localhost:8080/api/weather?city=Tokyo"'
|
||||||
|
)
|
||||||
|
uvicorn.run(app, host="0.0.0.0", port=8080, log_level="info")
|
||||||
Reference in New Issue
Block a user