wip

id.saladeaula.digital/app/oauth2.py

@@ -0,0 +1,131 @@
+from authlib.oauth2.rfc6749.grants import (
+    AuthorizationCodeGrant as _AuthorizationCodeGrant,
+)
+from authlib.oidc.core import OpenIDCode as OpenIDCode_
+from authlib.oidc.core import UserInfo
+from aws_lambda_powertools.event_handler.exceptions import NotFoundError
+from layercake.dateutils import now, ttl
+from layercake.dynamodb import DynamoDBPersistenceLayer, KeyPair
+
+from apigateway_oauth2 import (
+    AuthorizationServer,
+    OAuth2Client,
+    OAuth2Token,
+)
+from boto3clients import dynamodb_client
+from config import DYNAMODB_SORT_KEY, OAUTH2_TABLE
+
+oauth2_layer = DynamoDBPersistenceLayer(OAUTH2_TABLE, dynamodb_client)
+
+
+DUMMY_JWT_CONFIG = {
+    'key': 'secret-key',
+    'alg': 'HS256',
+    'iss': 'https://authlib.org',
+    'exp': 3600,
+}
+
+
+def create_save_token_func(persistence_layer: DynamoDBPersistenceLayer):
+    def save_token(token, request) -> OAuth2Token:
+        return OAuth2Token()
+
+    return save_token
+
+
+class ClientNotFoundError(NotFoundError):
+    def __init__(self, *_):
+        super().__init__('Client not found')
+
+
+def create_query_client_func(persistence_layer: DynamoDBPersistenceLayer):
+    def query_client(client_id) -> OAuth2Client:
+        client = persistence_layer.collection.get_item(
+            KeyPair('OAUTH2_CLIENT', f'CLIENT_ID#{client_id}'),
+            exc_cls=ClientNotFoundError,
+        )
+
+        _, client_id = client.get(DYNAMODB_SORT_KEY, '').split('#')
+
+        return OAuth2Client(
+            client_id=client_id,
+            redirect_uris=client['redirect_uris'],
+            response_types=client['response_types'],
+        )
+
+    return query_client
+
+
+def save_authorization_code(code, request):
+    data: dict = request.payload.data  # type: ignore
+    user: dict = request.user  # type: ignore
+    nonce: str | None = data.get('nonce')
+    now_ = now()
+    ttl_ = ttl(start_dt=now_, minutes=15)
+
+    with oauth2_layer.transact_writer() as transact:
+        transact.put(
+            item={
+                'id': f'OAUTH2_CODE#CLIENT_ID#{request.payload.client_id}',
+                'sk': f'CODE#{code}',
+                'redirect_uri': request.payload.redirect_uri,  # type: ignore
+                'scope': request.payload.scope,  # type: ignore
+                'user_id': user['id'],
+                'nonce': nonce,
+                'created_at': now_,
+                'ttl': ttl_,
+            },
+        )
+
+        if nonce:
+            transact.put(
+                item={
+                    'id': f'OAUTH2_CODE#CLIENT_ID#{request.payload.client_id}',
+                    'sk': f'NONCE#{nonce}',
+                    'code': code,
+                    'created_at': now_,
+                    'ttl': ttl_,
+                },
+            )
+
+
+def exists_nonce(nonce, request):
+    nonce_ = oauth2_layer.get_item(
+        KeyPair(
+            f'OAUTH2_CODE#CLIENT_ID#{request.payload.client_id}',
+            f'NONCE#{nonce}',
+        )
+    )
+    return bool(nonce_)
+
+
+class OpenIDCode(OpenIDCode_):
+    def exists_nonce(self, nonce, request):
+        return exists_nonce(nonce, request)
+
+    def get_jwt_config(self, grant):
+        return DUMMY_JWT_CONFIG
+
+    def generate_user_info(self, user, scope):
+        return UserInfo(
+            sub=user.id,
+            name=user.name,
+            email=user.email,
+        ).filter(scope)
+
+
+class AuthorizationCodeGrant(_AuthorizationCodeGrant):
+    TOKEN_ENDPOINT_AUTH_METHODS = [
+        'client_secret_basic',
+        'client_secret_post',
+    ]
+
+    def save_authorization_code(self, code: str, request):
+        return save_authorization_code(code, request)
+
+
+authorization = AuthorizationServer(
+    query_client=create_query_client_func(oauth2_layer),
+    save_token=create_save_token_func(oauth2_layer),
+)
+authorization.register_grant(AuthorizationCodeGrant, [OpenIDCode(require_nonce=False)])