add middleware

http-api/middlewares/tenant_middelware.py

@@ -0,0 +1,75 @@
+from http import HTTPStatus
+
+from aws_lambda_powertools.event_handler.api_gateway import (
+    APIGatewayHttpResolver,
+    Response,
+)
+from aws_lambda_powertools.event_handler.exceptions import BadRequestError, ServiceError
+from aws_lambda_powertools.event_handler.middlewares import (
+    BaseMiddlewareHandler,
+    NextMiddleware,
+)
+from layercake.dynamodb import ComposeKey, DynamoDBCollection, KeyPair
+from pydantic import UUID4, BaseModel
+
+from auth import AuthFlowType
+
+from .authorizer_middleware import User
+
+
+class Tenant(BaseModel):
+    id: UUID4 | str
+    name: str
+
+
+class TenantMiddleware(BaseMiddlewareHandler):
+    def __init__(self, collect: DynamoDBCollection) -> None:
+        self.collect = collect
+
+    def handler(
+        self,
+        app: APIGatewayHttpResolver,
+        next_middleware: NextMiddleware,
+    ) -> Response:
+        context = app.current_event.request_context.authorizer.get_lambda
+        auth_flow_type = context.get('auth_flow_type')
+
+        if auth_flow_type == AuthFlowType.API_AUTH:
+            app.append_context(tenant=Tenant(**context['tenant']))
+
+        if auth_flow_type == AuthFlowType.USER_AUTH:
+            app.append_context(
+                tenant=_tenant(
+                    app.current_event.headers.get('x-tenant'),
+                    app.context.get('user'),  # type: ignore
+                    collect=self.collect,
+                )
+            )
+
+        return next_middleware(app)
+
+
+class ForbiddenError(ServiceError):
+    def __init__(self, msg: str):
+        super().__init__(HTTPStatus.FORBIDDEN, msg)
+
+
+def _tenant(
+    tenant_id: str | None,
+    user: User,
+    /,
+    collect: DynamoDBCollection,
+) -> Tenant:
+    if not tenant_id:
+        raise BadRequestError('Missing tenant')
+
+    collect.get_item(
+        KeyPair(user.id, ComposeKey(tenant_id, prefix='acls')),
+        exception_cls=ForbiddenError,
+    )
+
+    if tenant_id == '*':
+        return Tenant(id=tenant_id, name='default')
+
+    obj = collect.get_item(KeyPair(tenant_id, '0'))
+    return Tenant.parse_obj(obj)