add jwt

id.saladeaula.digital/app/oauth2.py

@@ -1,7 +1,14 @@
+import time
+from dataclasses import dataclass
+
+from authlib.common.security import generate_token
 from authlib.common.urls import add_params_to_uri
+from authlib.jose import jwt
 from authlib.oauth2 import OAuth2Request, rfc7009, rfc9207
 from authlib.oauth2.rfc6749 import ClientMixin, TokenMixin, grants
+from authlib.oauth2.rfc6750 import BearerTokenGenerator
 from authlib.oauth2.rfc7636 import CodeChallenge
+from authlib.oauth2.rfc9068 import JWTBearerTokenGenerator as JWTBearerTokenGenerator_
 from authlib.oidc.core import OpenIDCode as OpenIDCode_
 from authlib.oidc.core import UserInfo
 from aws_lambda_powertools import Logger
@@ -17,7 +24,7 @@ from layercake.dynamodb import (
 from layercake.funcs import omit, pick
 
 from boto3clients import dynamodb_client
-from config import ISSUER, JWT_ALGORITHM, OAUTH2_TABLE
+from config import ISSUER, OAUTH2_TABLE
 from integrations.apigateway_oauth2.authorization_server import (
     AuthorizationServer,
 )
@@ -25,11 +32,23 @@ from integrations.apigateway_oauth2.tokens import (
     OAuth2AuthorizationCode,
     OAuth2Token,
 )
+from util import read_file_path
 
 logger = Logger(__name__)
 dyn = DynamoDBPersistenceLayer(OAUTH2_TABLE, dynamodb_client)
 
 
+@dataclass(frozen=True)
+class User:
+    id: str
+    name: str
+    email: str
+    email_verified: bool = False
+
+    def get_user_id(self):
+        return self.id
+
+
 class OpenIDCode(OpenIDCode_):
     def exists_nonce(self, nonce: str, request: OAuth2Request) -> bool:
         if not request.payload:
@@ -43,18 +62,19 @@ class OpenIDCode(OpenIDCode_):
 
     def get_jwt_config(self, grant):
         return {
-            'key': 'secret-key',
-            'alg': JWT_ALGORITHM,
+            'key': read_file_path('private.pem'),
+            'alg': 'RS256',
+            'kid': 'test',
             'iss': ISSUER,
-            'exp': 3600,
+            'exp': 3600 * 2,
         }
 
-    def generate_user_info(self, user: dict, scope: str) -> UserInfo:
+    def generate_user_info(self, user: User, scope: str) -> UserInfo:
         return UserInfo(
-            sub=user['id'],
-            name=user['name'],
-            email=user['email'],
-            email_verified=user.get('email_verified', False),
+            sub=user.id,
+            name=user.name,
+            email=user.email,
+            email_verified=user.email_verified,
         ).filter(scope)
 
 
@@ -147,7 +167,7 @@ class AuthorizationCodeGrant(grants.AuthorizationCodeGrant):
     def authenticate_user(
         self,
         authorization_code: OAuth2AuthorizationCode,
-    ) -> dict:
+    ) -> User:
         """Authenticate the user related to this authorization_code."""
         user = dyn.get_item(
             KeyPair(
@@ -155,7 +175,8 @@ class AuthorizationCodeGrant(grants.AuthorizationCodeGrant):
                 sk='0',
             ),
         )
-        return pick(('id', 'name', 'email', 'email_verified'), user)
+
+        return User(**pick(('id', 'name', 'email', 'email_verified'), user))
 
 
 class RefreshTokenNotFoundError(NotFoundError):
@@ -299,6 +320,120 @@ class IssuerParameter(rfc9207.IssuerParameter):
         return ISSUER
 
 
+class JWTBearerTokenGenerator(JWTBearerTokenGenerator_):
+    def get_jwks(self):
+        return read_file_path('private.pem')
+
+    def access_token_generator(self, client, grant_type, user, scope):
+        now = int(time.time())
+        expires_in = now + self._get_expires_in(client, grant_type)
+
+        token_data = {
+            'iss': self.issuer,
+            'exp': expires_in,
+            'client_id': client.get_client_id(),
+            'iat': now,
+            'jti': self.get_jti(client, grant_type, user, scope),
+            'scope': scope,
+        }
+
+        # In cases of access tokens obtained through grants where a resource owner is
+        # involved, such as the authorization code grant, the value of 'sub' SHOULD
+        # correspond to the subject identifier of the resource owner.
+
+        if user:
+            token_data['sub'] = user.get_user_id()
+
+        # In cases of access tokens obtained through grants where no resource owner is
+        # involved, such as the client credentials grant, the value of 'sub' SHOULD
+        # correspond to an identifier the authorization server uses to indicate the
+        # client application.
+
+        else:
+            token_data['sub'] = client.get_client_id()
+
+        # If the request includes a 'resource' parameter (as defined in [RFC8707]), the
+        # resulting JWT access token 'aud' claim SHOULD have the same value as the
+        # 'resource' parameter in the request.
+
+        # TODO: Implement this with RFC8707
+        if False:  # pragma: no cover
+            ...
+
+        # If the request does not include a 'resource' parameter, the authorization
+        # server MUST use a default resource indicator in the 'aud' claim. If a 'scope'
+        # parameter is present in the request, the authorization server SHOULD use it to
+        # infer the value of the default resource indicator to be used in the 'aud'
+        # claim. The mechanism through which scopes are associated with default resource
+        # indicator values is outside the scope of this specification.
+
+        else:
+            token_data['aud'] = self.get_audiences(client, user, scope)
+
+        # If the values in the 'scope' parameter refer to different default resource
+        # indicator values, the authorization server SHOULD reject the request with
+        # 'invalid_scope' as described in Section 4.1.2.1 of [RFC6749].
+        # TODO: Implement this with RFC8707
+
+        if auth_time := self.get_auth_time(user):
+            token_data['auth_time'] = auth_time
+
+        # The meaning and processing of acr Claim Values is out of scope for this
+        # specification.
+
+        if acr := self.get_acr(user):
+            token_data['acr'] = acr
+
+        # The definition of particular values to be used in the amr Claim is beyond the
+        # scope of this specification.
+
+        if amr := self.get_amr(user):
+            token_data['amr'] = amr
+
+        # Authorization servers MAY return arbitrary attributes not defined in any
+        # existing specification, as long as the corresponding claim names are collision
+        # resistant or the access tokens are meant to be used only within a private
+        # subsystem. Please refer to Sections 4.2 and 4.3 of [RFC7519] for details.
+
+        token_data.update(self.get_extra_claims(client, grant_type, user, scope))
+
+        # This specification registers the 'application/at+jwt' media type, which can
+        # be used to indicate that the content is a JWT access token. JWT access tokens
+        # MUST include this media type in the 'typ' header parameter to explicitly
+        # declare that the JWT represents an access token complying with this profile.
+        # Per the definition of 'typ' in Section 4.1.9 of [RFC7515], it is RECOMMENDED
+        # that the 'application/' prefix be omitted. Therefore, the 'typ' value used
+        # SHOULD be 'at+jwt'.
+
+        header = {'alg': self.alg, 'typ': 'at+jwt', 'kid': 'k1'}
+
+        access_token = jwt.encode(
+            header,
+            token_data,
+            key=self.get_jwks(),
+            check=False,
+        )
+        return access_token.decode()
+
+
+GRANT_TYPES_EXPIRES_IN = {
+    'refresh_token': 600,
+}
+
+
+def expires_in(client, grant_type: str) -> int:
+    return GRANT_TYPES_EXPIRES_IN.get(
+        grant_type, BearerTokenGenerator.DEFAULT_EXPIRES_IN
+    )
+
+
+def create_token_generator(length: int = 42):
+    def token_generator(*args, **kwargs):
+        return generate_token(length)
+
+    return token_generator
+
+
 server = AuthorizationServer(persistence_layer=dyn)
 server.register_grant(
     AuthorizationCodeGrant,
@@ -307,6 +442,14 @@ server.register_grant(
         CodeChallenge(required=True),
     ],
 )
+server.register_token_generator(
+    'default',
+    JWTBearerTokenGenerator(
+        issuer=ISSUER,
+        refresh_token_generator=create_token_generator(48),
+        expires_generator=expires_in,
+    ),
+)
 server.register_grant(RefreshTokenGrant)
 server.register_endpoint(RevocationEndpoint)
 server.register_extension(IssuerParameter())