"""
PKCE (Proof Key for Code Exchange) Implementation
RFC 7636 - OAuth 2.1 Mandatory
"""
import base64
import hashlib
import secrets
from typing import Literal
def generate_code_verifier(length: int = 64) -> str:
"""
Generate PKCE code verifier.
Args:
length: Length of verifier (43-128 characters)
Returns:
URL-safe random string
OAuth 2.1 Spec:
- Length: 43-128 characters
- Character set: [A-Za-z0-9-._~] (unreserved characters)
"""
if not 43 <= length <= 128:
raise ValueError("Code verifier length must be between 43-128")
# Generate random bytes and encode as URL-safe base64
random_bytes = secrets.token_bytes(length)
verifier = base64.urlsafe_b64encode(random_bytes).decode("utf-8")
# Remove padding and truncate to desired length
verifier = verifier.rstrip("=")[:length]
return verifier
def generate_code_challenge(code_verifier: str, method: Literal["S256"] = "S256") -> str:
"""
Generate PKCE code challenge from verifier.
Args:
code_verifier: Code verifier string
method: Challenge method (OAuth 2.1: only S256)
Returns:
Base64 URL-encoded SHA256 hash of verifier
OAuth 2.1 Spec:
- Only S256 method is allowed (plain is removed)
- code_challenge = BASE64URL(SHA256(code_verifier))
"""
if method != "S256":
raise ValueError("OAuth 2.1 only supports S256 method (plain is removed)")
if not code_verifier:
raise ValueError("Code verifier cannot be empty")
# SHA256 hash
digest = hashlib.sha256(code_verifier.encode("utf-8")).digest()
# Base64 URL encode (no padding)
challenge = base64.urlsafe_b64encode(digest).decode("utf-8").rstrip("=")
return challenge
def validate_code_challenge(
code_verifier: str, code_challenge: str, method: Literal["S256"] = "S256"
) -> bool:
"""
Validate PKCE code verifier against challenge.
Args:
code_verifier: Code verifier from client
code_challenge: Code challenge from authorization request
method: Challenge method
Returns:
True if valid, False otherwise
"""
if method != "S256":
raise ValueError("Only S256 method is supported")
try:
# Generate challenge from verifier
expected_challenge = generate_code_challenge(code_verifier, method)
# Constant-time comparison to prevent timing attacks
return secrets.compare_digest(expected_challenge, code_challenge)
except Exception:
return False
