128 lines
3.7 KiB
Python
128 lines
3.7 KiB
Python
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 boto3clients import dynamodb_client
|
|
from config import DYNAMODB_SORT_KEY, OAUTH2_TABLE
|
|
from integrations.apigateway_oauth2 import (
|
|
AuthorizationServer,
|
|
OAuth2Client,
|
|
OAuth2Token,
|
|
)
|
|
|
|
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
|
|
|
|
|
|
def create_query_client_func(persistence_layer: DynamoDBPersistenceLayer):
|
|
class ClientNotFoundError(NotFoundError):
|
|
def __init__(self, *_):
|
|
super().__init__('Client not found')
|
|
|
|
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_,
|
|
},
|
|
)
|
|
|
|
|
|
class OpenIDCode(OpenIDCode_):
|
|
def exists_nonce(self, nonce, request):
|
|
nonce_ = oauth2_layer.get_item(
|
|
KeyPair(
|
|
f'OAUTH2_CODE#CLIENT_ID#{request.payload.client_id}', # type:ignore
|
|
f'NONCE#{nonce}',
|
|
)
|
|
)
|
|
|
|
return bool(nonce_)
|
|
|
|
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)])
|