add oauth helpers

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,
+    )