add oauth helpers

This commit is contained in:
Matt Apperson
2025-11-13 13:12:32 -05:00
parent dede1d21ec
commit 05f81a553c
5 changed files with 316 additions and 0 deletions
+9
View File
@@ -0,0 +1,9 @@
{
"permissions": {
"allow": [
"Bash(python3:*)"
],
"deny": [],
"ask": []
}
}
+67
View File
@@ -0,0 +1,67 @@
"""
Example of using OAuth PKCE helper functions
This example demonstrates how to:
1. Generate a SHA-256 code challenge and verifier
2. Create an authorization URL for OAuth flow
"""
from openrouter import OpenRouter
from openrouter.utils import (
oauth_create_sha256_code_challenge,
oauth_create_authorization_url,
CreateSHA256CodeChallengeRequest,
CreateAuthorizationUrlRequestWithPKCE,
)
def main():
# Step 1: Generate a code challenge and verifier
# You can optionally provide your own code_verifier, or let it generate one
result = oauth_create_sha256_code_challenge()
print("Code Challenge:", result.code_challenge)
print("Code Verifier:", result.code_verifier)
print()
# Or provide your own code verifier (must be 43-128 chars, [A-Za-z0-9-._~])
custom_result = oauth_create_sha256_code_challenge(
CreateSHA256CodeChallengeRequest(
code_verifier="my-custom-verifier-that-is-at-least-43-characters-long-abcdefghij"
)
)
print("Custom Code Challenge:", custom_result.code_challenge)
print("Custom Code Verifier:", custom_result.code_verifier)
print()
# Step 2: Create an authorization URL
client = OpenRouter(api_key="your-api-key")
# Create authorization URL with PKCE
auth_url = oauth_create_authorization_url(
client,
CreateAuthorizationUrlRequestWithPKCE(
callback_url="https://your-app.com/callback",
code_challenge=result.code_challenge,
code_challenge_method="S256",
limit=10.0, # Optional credit limit
)
)
print("Authorization URL:", auth_url)
print()
# Step 3: User would visit the authorization URL and authorize the app
# Step 4: After authorization, the callback URL receives an authorization code
# Step 5: Exchange the code for an API key using the SDK's exchange method
# code = "authorization-code-from-callback"
# api_key_response = client.o_auth.exchange_auth_code_for_api_key(
# code=code,
# code_verifier=result.code_verifier,
# code_challenge_method="S256",
# )
# print("API Key:", api_key_response.key)
if __name__ == "__main__":
main()
+25
View File
@@ -53,9 +53,25 @@ if TYPE_CHECKING:
cast_partial,
)
from .logger import Logger, get_body_content, get_default_logger
from .oauth_create_sha256_code_challenge import (
oauth_create_sha256_code_challenge,
CreateSHA256CodeChallengeRequest,
CreateSHA256CodeChallengeResponse,
)
from .oauth_create_authorization_url import (
oauth_create_authorization_url,
CreateAuthorizationUrlRequest,
CreateAuthorizationUrlRequestBase,
CreateAuthorizationUrlRequestWithPKCE,
)
__all__ = [
"BackoffStrategy",
"CreateAuthorizationUrlRequest",
"CreateAuthorizationUrlRequestBase",
"CreateAuthorizationUrlRequestWithPKCE",
"CreateSHA256CodeChallengeRequest",
"CreateSHA256CodeChallengeResponse",
"FieldMetadata",
"find_metadata",
"FormMetadata",
@@ -78,6 +94,8 @@ __all__ = [
"match_status_codes",
"match_response",
"MultipartFormMetadata",
"oauth_create_authorization_url",
"oauth_create_sha256_code_challenge",
"OpenEnumMeta",
"PathParamMetadata",
"QueryParamMetadata",
@@ -110,6 +128,11 @@ __all__ = [
_dynamic_imports: dict[str, str] = {
"BackoffStrategy": ".retries",
"CreateAuthorizationUrlRequest": ".oauth_create_authorization_url",
"CreateAuthorizationUrlRequestBase": ".oauth_create_authorization_url",
"CreateAuthorizationUrlRequestWithPKCE": ".oauth_create_authorization_url",
"CreateSHA256CodeChallengeRequest": ".oauth_create_sha256_code_challenge",
"CreateSHA256CodeChallengeResponse": ".oauth_create_sha256_code_challenge",
"FieldMetadata": ".metadata",
"find_metadata": ".metadata",
"FormMetadata": ".metadata",
@@ -132,6 +155,8 @@ _dynamic_imports: dict[str, str] = {
"match_status_codes": ".values",
"match_response": ".values",
"MultipartFormMetadata": ".metadata",
"oauth_create_authorization_url": ".oauth_create_authorization_url",
"oauth_create_sha256_code_challenge": ".oauth_create_sha256_code_challenge",
"OpenEnumMeta": ".enums",
"PathParamMetadata": ".metadata",
"QueryParamMetadata": ".metadata",
@@ -0,0 +1,95 @@
"""Generate OAuth2 authorization URL"""
from dataclasses import dataclass
from typing import TYPE_CHECKING, Literal, Optional, Union
from urllib.parse import urlencode, urlparse
if TYPE_CHECKING:
from openrouter.sdk import OpenRouter
@dataclass
class CreateAuthorizationUrlRequestBase:
"""Base request parameters for creating an authorization URL"""
callback_url: Union[str, "urlparse"]
limit: Optional[float] = None
@dataclass
class CreateAuthorizationUrlRequestWithPKCE(CreateAuthorizationUrlRequestBase):
"""Request parameters with PKCE for creating an authorization URL"""
code_challenge_method: Literal["S256", "plain"]
code_challenge: str
# Union type for request - either with PKCE or without
CreateAuthorizationUrlRequest = Union[
CreateAuthorizationUrlRequestWithPKCE,
CreateAuthorizationUrlRequestBase,
]
def _get_server_url(client: "OpenRouter") -> str:
"""
Get the server URL from the client configuration
Args:
client: OpenRouter client instance
Returns:
The server URL
Raises:
ValueError: If no server URL is configured
"""
server_url, _ = client.sdk_configuration.get_server_details()
if not server_url:
raise ValueError("No server URL configured")
return server_url
def oauth_create_authorization_url(
client: "OpenRouter",
params: CreateAuthorizationUrlRequest,
) -> str:
"""
Generate an OAuth2 authorization URL
Generates a URL to redirect users to for authorizing your application. The
URL includes the provided callback URL and, if applicable, the code
challenge parameters for PKCE.
Args:
client: OpenRouter client instance
params: Request parameters including callback URL and optional PKCE parameters
Returns:
The authorization URL as a string
Raises:
ValueError: If no server URL is configured or parameters are invalid
See Also:
- https://openrouter.ai/docs/use-cases/oauth-pkce
"""
base_url = _get_server_url(client)
# Build the auth URL
auth_url = f"{base_url}/auth"
# Build query parameters
query_params = {
"callback_url": str(params.callback_url),
}
# Add PKCE parameters if present
if isinstance(params, CreateAuthorizationUrlRequestWithPKCE):
query_params["code_challenge"] = params.code_challenge
query_params["code_challenge_method"] = params.code_challenge_method
# Add limit if present
if params.limit is not None:
query_params["limit"] = str(params.limit)
# Construct final URL with query parameters
return f"{auth_url}?{urlencode(query_params)}"
@@ -0,0 +1,120 @@
"""Generate SHA-256 code challenge for PKCE OAuth flow"""
import base64
import hashlib
import re
import secrets
from dataclasses import dataclass
from typing import Optional
@dataclass
class CreateSHA256CodeChallengeRequest:
"""
Request parameters for creating a SHA-256 code challenge.
If not provided, a random code verifier will be generated.
If provided, must be 43-128 characters and contain only unreserved
characters [A-Za-z0-9-._~] per RFC 7636.
"""
code_verifier: Optional[str] = None
@dataclass
class CreateSHA256CodeChallengeResponse:
"""Response containing the code challenge and verifier"""
code_challenge: str
code_verifier: str
def _array_buffer_to_base64_url(data: bytes) -> str:
"""
Convert bytes to base64url encoding (RFC 4648)
Args:
data: Bytes to encode
Returns:
Base64url encoded string
"""
return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")
def _generate_code_verifier() -> str:
"""
Generate a cryptographically random code verifier per RFC 7636
RFC 7636 recommends 32 octets of random data, base64url encoded = 43 chars
Returns:
A random code verifier string
"""
random_bytes = secrets.token_bytes(32)
return _array_buffer_to_base64_url(random_bytes)
def _validate_code_verifier(code_verifier: str) -> None:
"""
Validate code verifier according to RFC 7636
Args:
code_verifier: The code verifier to validate
Raises:
ValueError: If the code verifier is invalid
"""
if len(code_verifier) < 43:
raise ValueError("Code verifier must be at least 43 characters")
if len(code_verifier) > 128:
raise ValueError("Code verifier must be at most 128 characters")
if not re.match(r"^[A-Za-z0-9\-._~]+$", code_verifier):
raise ValueError(
"Code verifier must only contain unreserved characters: [A-Za-z0-9-._~]"
)
def oauth_create_sha256_code_challenge(
params: Optional[CreateSHA256CodeChallengeRequest] = None,
) -> CreateSHA256CodeChallengeResponse:
"""
Generate a SHA-256 code challenge for PKCE
Generates a SHA-256 code challenge and corresponding code verifier for use
in the PKCE extension to OAuth2. If no code verifier is provided, a random
one will be generated according to RFC 7636 (32 random bytes, base64url
encoded). If a code verifier is provided, it must be 43-128 characters and
contain only unreserved characters [A-Za-z0-9-._~].
Args:
params: Optional request parameters. If None, a random code verifier will be generated.
Returns:
CreateSHA256CodeChallengeResponse containing the code challenge and verifier
Raises:
ValueError: If the provided code verifier is invalid
See Also:
- https://openrouter.ai/docs/use-cases/oauth-pkce
- https://datatracker.ietf.org/doc/html/rfc7636
"""
if params is None:
params = CreateSHA256CodeChallengeRequest()
code_verifier = params.code_verifier
if code_verifier is None:
code_verifier = _generate_code_verifier()
else:
_validate_code_verifier(code_verifier)
# Generate SHA-256 hash
data = code_verifier.encode("utf-8")
hash_digest = hashlib.sha256(data).digest()
# Convert hash to base64url
code_challenge = _array_buffer_to_base64_url(hash_digest)
return CreateSHA256CodeChallengeResponse(
code_challenge=code_challenge,
code_verifier=code_verifier,
)