diff --git a/.claude/settings.local.json b/.claude/settings.local.json new file mode 100644 index 0000000..b12585e --- /dev/null +++ b/.claude/settings.local.json @@ -0,0 +1,9 @@ +{ + "permissions": { + "allow": [ + "Bash(python3:*)" + ], + "deny": [], + "ask": [] + } +} diff --git a/examples/oauth_pkce_example.py b/examples/oauth_pkce_example.py new file mode 100644 index 0000000..e07e3e3 --- /dev/null +++ b/examples/oauth_pkce_example.py @@ -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() diff --git a/src/openrouter/utils/__init__.py b/src/openrouter/utils/__init__.py index 87192dd..2443805 100644 --- a/src/openrouter/utils/__init__.py +++ b/src/openrouter/utils/__init__.py @@ -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", diff --git a/src/openrouter/utils/oauth_create_authorization_url.py b/src/openrouter/utils/oauth_create_authorization_url.py new file mode 100644 index 0000000..c403863 --- /dev/null +++ b/src/openrouter/utils/oauth_create_authorization_url.py @@ -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)}" diff --git a/src/openrouter/utils/oauth_create_sha256_code_challenge.py b/src/openrouter/utils/oauth_create_sha256_code_challenge.py new file mode 100644 index 0000000..564862e --- /dev/null +++ b/src/openrouter/utils/oauth_create_sha256_code_challenge.py @@ -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, + )