add session route

id.saladeaula.digital/app/routes/authorize.py

@@ -1,49 +1,40 @@
-from http import HTTPStatus
 from http.cookies import SimpleCookie
-from urllib.parse import ParseResult, quote, urlencode, urlunparse
 
+import jwt
 from authlib.oauth2 import OAuth2Error
 from authlib.oauth2.rfc6749 import errors
 from aws_lambda_powertools import Logger
-from aws_lambda_powertools.event_handler import Response
 from aws_lambda_powertools.event_handler.api_gateway import Router
-from jose.exceptions import JWTError
+from aws_lambda_powertools.event_handler.exceptions import BadRequestError
+from layercake.dynamodb import DynamoDBPersistenceLayer, KeyPair
 
-from jose_ import verify_jwt
+from boto3clients import dynamodb_client
+from config import ISSUER, JWT_ALGORITHM, JWT_SECRET, OAUTH2_TABLE
 from oauth2 import server
 
 router = Router()
 logger = Logger(__name__)
+oauth2_layer = DynamoDBPersistenceLayer(OAUTH2_TABLE, dynamodb_client)
 
 
 @router.get('/authorize')
 def authorize():
     current_event = router.current_event
     cookies = _parse_cookies(current_event.get('cookies', []))
-    id_token = cookies.get('id_token')
-    continue_url = _build_continue_url(
-        current_event.path,
-        current_event.query_string_parameters,
-    )
+    session_id = cookies.get('session_id')
 
-    login_url = f'/login?continue={continue_url}'
-
-    try:
-        if not id_token:
-            raise ValueError('Missing id_token')
-
-        user = verify_jwt(id_token)
-    except (ValueError, JWTError):
-        return Response(
-            status_code=HTTPStatus.FOUND,
-            headers={'Location': login_url},
-        )
+    if not session_id:
+        raise BadRequestError('Missing session_id')
 
     try:
+        user_id = verify_session(session_id)
         grant = server.get_consent_grant(
             request=router.current_event,
-            end_user={'id': user['sub']},
+            end_user={'id': user_id},
         )
+    except jwt.exceptions.InvalidTokenError as err:
+        logger.exception(err)
+        raise BadRequestError(str(err))
     except OAuth2Error as err:
         logger.exception(err)
         return dict(err.get_body())
@@ -51,7 +42,7 @@ def authorize():
     try:
         return server.create_authorization_response(
             request=router.current_event,
-            grant_user={'id': user['sub']},
+            grant_user={'id': user_id},
             grant=grant,
         )
     except errors.OAuth2Error as err:
@@ -59,6 +50,29 @@ def authorize():
         return {}
 
 
+def verify_session(session_id: str) -> str:
+    payload = jwt.decode(
+        session_id,
+        JWT_SECRET,
+        algorithms=[JWT_ALGORITHM],
+        issuer=ISSUER,
+        options={
+            'require': ['exp', 'sub', 'iss', 'sid'],
+            'leeway': 60,
+        },
+    )
+
+    oauth2_layer.collection.get_item(
+        KeyPair(
+            pk='SESSION',
+            sk=payload['sid'],
+        ),
+        exc_cls=SessionRevokedError,
+    )
+
+    return payload['sub']
+
+
 def _parse_cookies(cookies: list[str] | None) -> dict[str, str]:
     parsed_cookies = {}
 
@@ -73,9 +87,6 @@ def _parse_cookies(cookies: list[str] | None) -> dict[str, str]:
     return parsed_cookies
 
 
-def _build_continue_url(
-    path: str,
-    query_string_parameters: dict,
-) -> str:
-    query = urlencode(query_string_parameters)
-    return quote(urlunparse(ParseResult('', '', path, '', query, '')), safe='')
+class SessionRevokedError(BadRequestError):
+    def __init__(self, *_):
+        super().__init__('Session revoked')

id.saladeaula.digital/app/routes/login.py

@@ -1,91 +0,0 @@
-from http import HTTPStatus
-from typing import Annotated
-
-from aws_lambda_powertools.event_handler import (
-    Response,
-)
-from aws_lambda_powertools.event_handler.api_gateway import Router
-from aws_lambda_powertools.event_handler.exceptions import ForbiddenError, NotFoundError
-from aws_lambda_powertools.event_handler.openapi.params import Form, Param
-from aws_lambda_powertools.shared.cookies import Cookie
-from jinja2 import Environment, PackageLoader, select_autoescape
-from layercake.dynamodb import DynamoDBPersistenceLayer, KeyPair
-from passlib.hash import pbkdf2_sha256
-
-from boto3clients import dynamodb_client
-from config import OAUTH2_TABLE
-from jose_ import generate_jwt
-
-router = Router()
-oauth2_layer = DynamoDBPersistenceLayer(OAUTH2_TABLE, dynamodb_client)
-templates = Environment(
-    loader=PackageLoader('app'),
-    autoescape=select_autoescape(['html']),
-)
-
-
-@router.get('/login', compress=True)
-def login_form(continue_: Annotated[str, Param(alias='continue')]):
-    template = templates.get_template('login.html')
-    html = template.render(**{'continue': continue_})
-
-    return Response(
-        body=html,
-        status_code=HTTPStatus.OK,
-        content_type='text/html',
-    )
-
-
-@router.post('/login')
-def login(
-    username: Annotated[str, Form()],
-    password: Annotated[str, Form()],
-    continue_: Annotated[str, Form(alias='continue')],
-):
-    user_id, password_hash = _get_user(username)
-
-    if not pbkdf2_sha256.verify(password, password_hash):
-        raise ForbiddenError('Invalid credentials')
-
-    jwt_token = generate_jwt(user_id, username)
-
-    return Response(
-        status_code=HTTPStatus.FOUND,
-        headers={
-            'Location': continue_,
-        },
-        cookies=[
-            Cookie(
-                name='id_token',
-                value=jwt_token,
-                http_only=True,
-                same_site=None,
-            ),
-        ],
-    )
-
-
-def _get_user(username: str) -> tuple[str, str]:
-    r = oauth2_layer.collection.get_item(
-        # Post-migration: uncomment the following line
-        # KeyPair('EMAIL', username),
-        KeyPair('email', username),
-        exc_cls=EmailNotFoundError,
-    )
-
-    password = oauth2_layer.collection.get_item(
-        KeyPair(r['user_id'], 'PASSWORD'),
-        exc_cls=UserNotFoundError,
-    )
-
-    return r['user_id'], password['hash']
-
-
-class EmailNotFoundError(NotFoundError):
-    def __init__(self, *_):
-        super().__init__('Email not found')
-
-
-class UserNotFoundError(NotFoundError):
-    def __init__(self, *_):
-        super().__init__('User not found')

id.saladeaula.digital/app/routes/session.py

@@ -0,0 +1,107 @@
+from http import HTTPStatus
+from typing import Annotated
+from uuid import uuid4
+
+import jwt
+from aws_lambda_powertools.event_handler import (
+    Response,
+)
+from aws_lambda_powertools.event_handler.api_gateway import Router
+from aws_lambda_powertools.event_handler.exceptions import ForbiddenError, NotFoundError
+from aws_lambda_powertools.event_handler.openapi.params import Body
+from aws_lambda_powertools.shared.cookies import Cookie
+from layercake.dateutils import now, ttl
+from layercake.dynamodb import DynamoDBPersistenceLayer, KeyPair, SortKey
+from passlib.hash import pbkdf2_sha256
+
+from boto3clients import dynamodb_client
+from config import ISSUER, JWT_ALGORITHM, JWT_EXP_SECONDS, JWT_SECRET, OAUTH2_TABLE
+
+router = Router()
+oauth2_layer = DynamoDBPersistenceLayer(OAUTH2_TABLE, dynamodb_client)
+
+
+@router.post('/session')
+def session(
+    username: Annotated[str, Body()],
+    password: Annotated[str, Body()],
+):
+    user_id, password_hash = _get_user(username)
+
+    if not pbkdf2_sha256.verify(password, password_hash):
+        raise ForbiddenError('Invalid credentials')
+
+    return Response(
+        status_code=HTTPStatus.FOUND,
+        cookies=[
+            Cookie(
+                name='session_id',
+                value=new_session(user_id),
+                http_only=True,
+                secure=True,
+                same_site=None,
+            )
+        ],
+    )
+
+
+def _get_user(username: str) -> tuple[str, str]:
+    sk = SortKey(username, path_spec='user_id')
+    user = oauth2_layer.collection.get_items(
+        KeyPair(pk='email', sk=sk, rename_key=sk.path_spec)
+        + KeyPair(pk='cpf', sk=sk, rename_key=sk.path_spec),
+        flatten_top=False,
+    )
+
+    if not user:
+        raise UserNotFoundError()
+
+    password = oauth2_layer.collection.get_item(
+        KeyPair(user['user_id'], 'PASSWORD'),
+        exc_cls=UserNotFoundError,
+    )
+
+    return user['user_id'], password['hash']
+
+
+def new_session(sub: str) -> str:
+    now_ = now()
+    sid = str(uuid4())
+    exp = ttl(start_dt=now_, seconds=JWT_EXP_SECONDS)
+    token = jwt.encode(
+        {
+            'sid': sid,
+            'sub': sub,
+            'iss': ISSUER,
+            'iat': int(now_.timestamp()),
+            'exp': exp,
+        },
+        JWT_SECRET,
+        algorithm=JWT_ALGORITHM,
+    )
+
+    with oauth2_layer.transact_writer() as transact:
+        transact.put(
+            item={
+                'id': 'SESSION',
+                'sk': sid,
+                'user_id': sub,
+                'ttl': exp,
+                'created_at': now_,
+            }
+        )
+        transact.put(
+            item={
+                'id': sub,
+                'sk': f'SESSION#{sid}',
+                'ttl': exp,
+                'created_at': now_,
+            }
+        )
+
+    return token
+
+
+class UserNotFoundError(NotFoundError):
+    def __init__(self, *_):
+        super().__init__('User not found')