import time import uuid from typing import Any, Optional, Union from authlib.jose import JsonWebKey, jwt from .utils import calculate_jwk_thumbprint, normalize_url_for_htu, sha256_base64url # A private RSA JWK for test usage. PRIVATE_JWK = { "kty": "RSA", "alg": "RS256", "use": "sig", "kid": "TEST_KEY", "n": "whYOFK2Ocbbpb_zVypi9SeKiNUqKQH0zTKN1-6fpCTu6ZalGI82s7XK3tan4dJt90ptUPKD2zvxqTzFNfx4HHHsrYCf2-FMLn1VTJfQazA2BvJqAwcpW1bqRUEty8tS_Yv4hRvWfQPcc2Gc3-_fQOOW57zVy-rNoJc744kb30NjQxdGp03J2S3GLQu7oKtSDDPooQHD38PEMNnITf0pj-KgDPjymkMGoJlO3aKppsjfbt_AH6GGdRghYRLOUwQU-h-ofWHR3lbYiKtXPn5dN24kiHy61e3VAQ9_YAZlwXC_99GGtw_NpghFAuM4P1JDn0DppJldy3PGFC0GfBCZASw", "e": "AQAB", "d": "VuVE_KEP6323WjpbBdAIv7HGahGrgGANvbxZsIhm34lsVOPK0XDegZkhAybMZHjRhp-gwVxX5ChC-J3cUpOBH5FNxElgW6HizD2Jcq6t6LoLYgPSrfEHm71iHg8JsgrqfUnGYFzMJmv88C6WdCtpgG_qJV1K00_Ly1G1QKoBffEs-v4fAMJrCbUdCz1qWto-PU-HLMEo-krfEpGgcmtZeRlDADh8cETMQlgQfQX2VWq_aAP4a1SXmo-j0cvRU4W5Fj0RVwNesIpetX2ZFz4p_JmB5sWFEj_fC7h5z2lq-6Bme2T3BHtXkIxoBW0_pYVnASC8P2puO5FnVxDmWuHDYQ", "p": "07rgXd_tLUhVRF_g1OaqRZh5uZ8hiLWUSU0vu9coOaQcatSqjQlIwLW8UdKv_38GrmpIfgcEVQjzq6rFBowUm9zWBO9Eq6enpasYJBOeD8EMeDK-nsST57HjPVOCvoVC5ZX-cozPXna3iRNZ1TVYBY3smn0IaxysIK-zxESf4pM", "q": "6qrE9TPhCS5iNR7QrKThunLu6t4H_8CkYRPLbvOIt2MgZyPLiZCsvdkTVSOX76QQEXt7Y0nTNua69q3K3Jhf-YOkPSJsWTxgrfOnjoDvRKzbW3OExIMm7D99fVBODuNWinjYgUwGSqGAsb_3TKhtI-Gr5ls3fn6B6oEjVL0dpmk", "dp": "mHqjrFdgelT2OyiFRS3dAAPf3cLxJoAGC4gP0UoQyPocEP-Y17sQ7t-ygIanguubBy65iDFLeGXa_g0cmSt2iAzRAHrDzI8P1-pQl2KdWSEg9ssspjBRh_F_AiJLLSPRWn_b3-jySkhawtfxwO8Kte1QsK1My765Y0zFvJnjPws", "dq": "KmjaV4YcsVAUp4z-IXVa5htHWmLuByaFjpXJOjABEUN0467wZdgjn9vPRp-8Ia8AyGgMkJES_uUL_PDDrMJM9gb4c6P4-NeUkVtreLGMjFjA-_IQmIMrUZ7XywHsWXx0c2oLlrJqoKo3W-hZhR0bPFTYgDUT_mRWjk7wV6wl46E", "qi": "iYltkV_4PmQDfZfGFpzn2UtYEKyhy-9t3Vy8Mw2VHLAADKGwJvVK5ficQAr2atIF1-agXY2bd6KV-w52zR8rmZfTr0gobzYIyqHczOm13t7uXJv2WygY7QEC2OGjdxa2Fr9RnvS99ozMa5nomZBqTqT7z5QV33czjPRCjvg6FcE", } async def generate_token( domain: str, user_id: str, audience: Optional[str] = None, issuer: Union[str, bool, None] = None, iat: bool = True, exp: bool = True, claims: Optional[dict[str, Any]] = None, expiration_time: int = 3600, ) -> str: """ Generates a real RS256-signed JWT using the private key above. Args: domain: The Auth0 domain (used if issuer is not False). user_id: The 'sub' claim in the token. audience: The 'aud' claim in the token. If omitted, 'aud' won't be included. issuer: - If a string, it's placed in 'iss' claim. - If None, default is f"https://{domain}/". - If False, skip 'iss' claim entirely. iat: Whether to set the 'iat' (issued at) claim. If False, skip it. exp: Whether to set the 'exp' claim. If False, skip it. claims: Additional custom claims to merge into the token. expiration_time: If exp is True, how many seconds from now until expiration. Returns: A RS256-signed JWT string. Example usage: token = generate_token( domain="example.us.auth0.com", user_id="user123", audience="my-api", issuer=False, iat=False, exp=False, claims={"scope": "read:stuff"} ) """ token_claims = dict(claims or {}) token_claims.setdefault("sub", user_id) if iat: token_claims["iat"] = int(time.time()) if exp: token_claims["exp"] = int(time.time()) + expiration_time if issuer is not False: token_claims["iss"] = issuer if isinstance(issuer, str) else f"https://{domain}/" if audience: token_claims["aud"] = audience key = JsonWebKey.import_key(PRIVATE_JWK) header = {"alg": "RS256", "kid": PRIVATE_JWK["kid"]} token = jwt.encode(header, token_claims, key) # Ensure we return a string, not bytes return token.decode('utf-8') if isinstance(token, bytes) else token # A private EC P-256 private key for DPoP proof generation (test only) PRIVATE_EC_JWK = { "kty": "EC", "crv": "P-256", "x": "MKBCTNIcKUSDii11ySs3526iDZ8AiTo7Tu6KPAqv7D4", "y": "4Etl6SRW2YiLUrN5vfvVHuhp7x8PxltmWWlbbM4IFyM", "d": "870MB6gfuTJ4HtUnUvYMyJpr5eUZNP4Bk43bVdj3eAE" } async def generate_dpop_proof( access_token: str, http_method: str, http_url: str, jti: Optional[str] = None, iat: bool = True, claims: Optional[dict[str, Any]] = None, header_overrides: Optional[dict[str, Any]] = None, iat_time: Optional[int] = None, include_jti: bool = True ) -> str: """ Generates a real ES256-signed DPoP proof JWT using the EC private key above. Args: access_token: The access token to create proof for (used for ath claim). http_method: The HTTP method (e.g., "GET", "POST") for htm claim. http_url: The HTTP URL for htu claim. jti: The unique identifier for the proof. If omitted, generates random UUID. iat: Whether to set the 'iat' (issued at) claim. If False, skip it. claims: Additional custom claims to merge into the proof. header_overrides: Override header parameters (e.g., for testing invalid headers). iat_time: Fixed time for iat claim (for testing). If None, uses current time. include_jti: Whether to include the 'jti' claim. If False, jti is completely omitted. Returns: An ES256-signed DPoP proof JWT string. Example usage: proof = await generate_dpop_proof( access_token="eyJ...", http_method="GET", http_url="https://api.example.com/resource", iat=False, # Skip iat for testing claims={"custom": "claim"} ) """ proof_claims = dict(claims or {}) if iat: proof_claims["iat"] = iat_time if iat_time is not None else int(time.time()) if include_jti: if jti is not None: proof_claims["jti"] = jti else: proof_claims["jti"] = str(uuid.uuid4()) proof_claims["htm"] = http_method proof_claims["htu"] = normalize_url_for_htu(http_url) proof_claims["ath"] = sha256_base64url(access_token) public_jwk = {k: v for k, v in PRIVATE_EC_JWK.items() if k != "d"} header = { "alg": "ES256", "typ": "dpop+jwt", "jwk": public_jwk } if header_overrides: header.update(header_overrides) key = JsonWebKey.import_key(PRIVATE_EC_JWK) token = jwt.encode(header, proof_claims, key) # Ensure we return a string, not bytes return token.decode('utf-8') if isinstance(token, bytes) else token async def generate_token_with_cnf( domain: str, user_id: str, audience: str, jkt_thumbprint: Optional[str] = None, **kwargs ) -> str: """ Generates an access token with cnf (confirmation) claim for DPoP binding. Extends the existing generate_token() function with DPoP support. Args: domain: The Auth0 domain (used if issuer is not False). user_id: The 'sub' claim in the token. audience: The 'aud' claim in the token. jkt_thumbprint: JWK thumbprint to include in cnf claim. If None, calculates from PRIVATE_EC_JWK. **kwargs: Additional arguments passed to generate_token(). Returns: A RS256-signed JWT string with cnf claim. Example usage: token = await generate_token_with_cnf( domain="auth0.local", user_id="user123", audience="my-api", jkt_thumbprint="custom_thumbprint" ) """ if jkt_thumbprint is None: jkt_thumbprint = calculate_jwk_thumbprint(PRIVATE_EC_JWK) existing_claims = kwargs.get('claims', {}) cnf_claims = dict(existing_claims) cnf_claims["cnf"] = {"jkt": jkt_thumbprint} kwargs['claims'] = cnf_claims return await generate_token( domain=domain, user_id=user_id, audience=audience, **kwargs )