76 lines
2.0 KiB
Python
76 lines
2.0 KiB
Python
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)
|