456 lines
15 KiB
Python
456 lines
15 KiB
Python
import time
|
|
from dataclasses import dataclass
|
|
|
|
from authlib.common.security import generate_token
|
|
from authlib.common.urls import add_params_to_uri
|
|
from authlib.jose import jwt
|
|
from authlib.oauth2 import OAuth2Request, rfc7009, rfc9207
|
|
from authlib.oauth2.rfc6749 import ClientMixin, TokenMixin, grants
|
|
from authlib.oauth2.rfc6750 import BearerTokenGenerator
|
|
from authlib.oauth2.rfc7636 import CodeChallenge
|
|
from authlib.oauth2.rfc9068 import JWTBearerTokenGenerator as JWTBearerTokenGenerator_
|
|
from authlib.oidc.core import OpenIDCode as OpenIDCode_
|
|
from authlib.oidc.core import UserInfo
|
|
from aws_lambda_powertools import Logger
|
|
from aws_lambda_powertools.event_handler.api_gateway import Response
|
|
from aws_lambda_powertools.event_handler.exceptions import NotFoundError
|
|
from layercake.dateutils import now, ttl
|
|
from layercake.dynamodb import (
|
|
DynamoDBPersistenceLayer,
|
|
KeyPair,
|
|
SortKey,
|
|
TransactKey,
|
|
)
|
|
from layercake.funcs import omit, pick
|
|
|
|
from boto3clients import dynamodb_client
|
|
from config import ISSUER, OAUTH2_TABLE
|
|
from integrations.apigateway_oauth2.authorization_server import (
|
|
AuthorizationServer,
|
|
)
|
|
from integrations.apigateway_oauth2.tokens import (
|
|
OAuth2AuthorizationCode,
|
|
OAuth2Token,
|
|
)
|
|
from util import read_file_path
|
|
|
|
logger = Logger(__name__)
|
|
dyn = DynamoDBPersistenceLayer(OAUTH2_TABLE, dynamodb_client)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class User:
|
|
id: str
|
|
name: str
|
|
email: str
|
|
email_verified: bool = False
|
|
|
|
def get_user_id(self):
|
|
return self.id
|
|
|
|
|
|
class OpenIDCode(OpenIDCode_):
|
|
def exists_nonce(self, nonce: str, request: OAuth2Request) -> bool:
|
|
if not request.payload:
|
|
raise ValueError('Missing request payload')
|
|
|
|
nonce_ = dyn.get_item(
|
|
KeyPair(pk='OAUTH2#CODE', sk=f'NONCE#{nonce}'),
|
|
)
|
|
|
|
return bool(nonce_)
|
|
|
|
def get_jwt_config(self, grant):
|
|
return {
|
|
'key': read_file_path('private.pem'),
|
|
'alg': 'RS256',
|
|
'kid': 'test',
|
|
'iss': ISSUER,
|
|
'exp': 3600 * 2,
|
|
}
|
|
|
|
def generate_user_info(self, user: User, scope: str) -> UserInfo:
|
|
return UserInfo(
|
|
sub=user.id,
|
|
name=user.name,
|
|
email=user.email,
|
|
email_verified=user.email_verified,
|
|
).filter(scope)
|
|
|
|
|
|
class AuthorizationCodeGrant(grants.AuthorizationCodeGrant):
|
|
TOKEN_ENDPOINT_AUTH_METHODS = [
|
|
'client_secret_basic',
|
|
'client_secret_post',
|
|
'none',
|
|
]
|
|
|
|
def save_authorization_code(
|
|
self,
|
|
code: str,
|
|
request: OAuth2Request,
|
|
) -> None:
|
|
"""Save authorization_code for later use."""
|
|
if not request.payload:
|
|
raise ValueError('Missing request payload')
|
|
|
|
if not request.user:
|
|
raise ValueError('Missing request user')
|
|
|
|
client_id: str = request.payload.client_id
|
|
scope: str = request.payload.scope
|
|
data: dict = request.payload.data
|
|
user_id: str = request.user
|
|
nonce: str | None = data.get('nonce')
|
|
code_challenge: str | None = data.get('code_challenge')
|
|
code_challenge_method: str | None = data.get('code_challenge_method')
|
|
now_ = now()
|
|
ttl_ = ttl(start_dt=now_, minutes=10)
|
|
|
|
with dyn.transact_writer() as transact:
|
|
transact.put(
|
|
item={
|
|
'id': 'OAUTH2#CODE',
|
|
'sk': f'CODE#{code}',
|
|
'redirect_uri': request.payload.redirect_uri,
|
|
'response_type': request.payload.response_type,
|
|
'scope': scope,
|
|
'client_id': client_id,
|
|
'user_id': user_id,
|
|
'nonce': nonce,
|
|
'code_challenge': code_challenge,
|
|
'code_challenge_method': code_challenge_method,
|
|
'created_at': now_,
|
|
'ttl': ttl_,
|
|
},
|
|
)
|
|
|
|
if nonce:
|
|
transact.put(
|
|
item={
|
|
'id': 'OAUTH2#CODE',
|
|
'sk': f'NONCE#{nonce}',
|
|
'client_id': client_id,
|
|
'code': code,
|
|
'created_at': now_,
|
|
'ttl': ttl_,
|
|
},
|
|
)
|
|
|
|
def query_authorization_code(
|
|
self,
|
|
code: str,
|
|
client: ClientMixin,
|
|
) -> OAuth2AuthorizationCode:
|
|
"""Get authorization_code from previously savings."""
|
|
auth_code = dyn.get_item(
|
|
KeyPair(pk='OAUTH2#CODE', sk=f'CODE#{code}'),
|
|
)
|
|
|
|
return OAuth2AuthorizationCode(
|
|
code=code,
|
|
**auth_code,
|
|
)
|
|
|
|
def delete_authorization_code(
|
|
self,
|
|
authorization_code: OAuth2AuthorizationCode,
|
|
) -> None:
|
|
"""Delete authorization code from database or cache."""
|
|
dyn.delete_item(
|
|
KeyPair(
|
|
pk='OAUTH2#CODE',
|
|
sk=f'CODE#{authorization_code.code}',
|
|
),
|
|
)
|
|
|
|
def authenticate_user(
|
|
self,
|
|
authorization_code: OAuth2AuthorizationCode,
|
|
) -> User:
|
|
"""Authenticate the user related to this authorization_code."""
|
|
user = dyn.get_item(
|
|
KeyPair(
|
|
pk=authorization_code.user_id,
|
|
sk='0',
|
|
),
|
|
)
|
|
|
|
return User(**pick(('id', 'name', 'email', 'email_verified'), user))
|
|
|
|
|
|
class RefreshTokenNotFoundError(NotFoundError):
|
|
def __init__(self, *_):
|
|
super().__init__('Refresh token not found')
|
|
|
|
|
|
class RefreshTokenGrant(grants.RefreshTokenGrant):
|
|
TOKEN_ENDPOINT_AUTH_METHODS = [
|
|
'client_secret_basic',
|
|
'client_secret_post',
|
|
'none',
|
|
]
|
|
|
|
# The authorization server MAY issue a new refresh token
|
|
INCLUDE_NEW_REFRESH_TOKEN = True
|
|
|
|
def authenticate_refresh_token(self, refresh_token: str, **kwargs) -> TokenMixin:
|
|
"""Get token information with refresh_token string."""
|
|
token = dyn.collection.get_item(
|
|
KeyPair(
|
|
pk='OAUTH2#TOKEN',
|
|
sk=f'REFRESH_TOKEN#{refresh_token}',
|
|
),
|
|
exc_cls=RefreshTokenNotFoundError,
|
|
)
|
|
|
|
logger.info('Refresh token retrieved', token=token)
|
|
|
|
return OAuth2Token(
|
|
expires_in=int(token['expires_in']),
|
|
issued_at=int(token['issued_at']),
|
|
refresh_token=refresh_token,
|
|
**omit(('expires_in', 'issued_at'), token),
|
|
)
|
|
|
|
def authenticate_user(self, refresh_token: TokenMixin):
|
|
"""Authenticate the user related to this credential."""
|
|
return refresh_token.get_user()
|
|
|
|
def revoke_old_credential(self, refresh_token: TokenMixin) -> None:
|
|
"""The authorization server MAY revoke the old refresh token after
|
|
issuing a new refresh token to the client."""
|
|
|
|
logger.debug('Revoking old refresh token', refresh_token=refresh_token)
|
|
token = getattr(refresh_token, 'refresh_token', None)
|
|
user = refresh_token.get_user()
|
|
|
|
with dyn.transact_writer() as transact:
|
|
transact.delete(
|
|
key=KeyPair(
|
|
pk='OAUTH2#TOKEN',
|
|
sk=f'REFRESH_TOKEN#{token}',
|
|
)
|
|
)
|
|
transact.delete(
|
|
key=KeyPair(
|
|
pk=user.get('id'),
|
|
sk=f'SESSION#REFRESH_TOKEN#{token}',
|
|
)
|
|
)
|
|
|
|
|
|
class RevocationEndpoint(rfc7009.RevocationEndpoint):
|
|
def query_token( # type: ignore
|
|
self,
|
|
token_string: str,
|
|
token_type_hint: str | None = None,
|
|
):
|
|
t = dyn.collection.get_items(
|
|
TransactKey('OAUTH2#TOKEN')
|
|
+ SortKey(sk=f'REFRESH_TOKEN#{token_string}', rename_key='refresh_token')
|
|
+ SortKey(sk=f'ACCESS_TOKEN#{token_string}', rename_key='access_token'),
|
|
flatten_top=False,
|
|
)
|
|
|
|
if not t:
|
|
return None
|
|
|
|
logger.debug('Tokens retrieved', result=t)
|
|
|
|
if not token_type_hint:
|
|
token_type_hint = (
|
|
'refresh_token' if 'refresh_token' in t else 'access_token'
|
|
)
|
|
|
|
token = t[token_type_hint]
|
|
|
|
return OAuth2Token(
|
|
expires_in=int(token['expires_in']),
|
|
issued_at=int(token['issued_at']),
|
|
**{token_type_hint: token_string},
|
|
**omit(('expires_in', 'issued_at', 'refresh_token', 'access_token'), token),
|
|
)
|
|
|
|
def revoke_token(
|
|
self,
|
|
token: OAuth2Token,
|
|
request: OAuth2Request,
|
|
):
|
|
user_id = token.user['id']
|
|
r = dyn.collection.query(KeyPair(pk=user_id, sk='SESSION'))
|
|
|
|
with dyn.transact_writer() as transact:
|
|
# Revoke all sessions, access tokens, and refresh tokens
|
|
for x in r['items']:
|
|
pk, sk = x['id'], x['sk']
|
|
*_, kind, idx = sk.split('#')
|
|
|
|
transact.delete(key=KeyPair(pk, sk))
|
|
transact.delete(
|
|
key=KeyPair(
|
|
pk='SESSION',
|
|
sk=idx,
|
|
)
|
|
if kind == 'SESSION'
|
|
else KeyPair(
|
|
pk='OAUTH2#TOKEN',
|
|
sk=f'{kind}#{idx}',
|
|
)
|
|
)
|
|
|
|
|
|
class IssuerParameter(rfc9207.IssuerParameter):
|
|
def add_issuer_parameter(
|
|
self,
|
|
authorization_server: AuthorizationServer,
|
|
response: Response,
|
|
):
|
|
location = response.headers.get('Location')
|
|
|
|
if self.get_issuer() and location:
|
|
# RFC9207 §2
|
|
# In authorization responses to the client, including error responses,
|
|
# an authorization server supporting this specification MUST indicate
|
|
# its identity by including the iss parameter in the response.
|
|
new_location = add_params_to_uri(location, {'iss': self.get_issuer()}) # type: ignore
|
|
response.headers['Location'] = new_location
|
|
|
|
def get_issuer(self) -> str:
|
|
return ISSUER
|
|
|
|
|
|
class JWTBearerTokenGenerator(JWTBearerTokenGenerator_):
|
|
def get_jwks(self):
|
|
return read_file_path('private.pem')
|
|
|
|
def access_token_generator(self, client, grant_type, user, scope):
|
|
now = int(time.time())
|
|
expires_in = now + self._get_expires_in(client, grant_type)
|
|
|
|
token_data = {
|
|
'iss': self.issuer,
|
|
'exp': expires_in,
|
|
'client_id': client.get_client_id(),
|
|
'iat': now,
|
|
'jti': self.get_jti(client, grant_type, user, scope),
|
|
'scope': scope,
|
|
}
|
|
|
|
# In cases of access tokens obtained through grants where a resource owner is
|
|
# involved, such as the authorization code grant, the value of 'sub' SHOULD
|
|
# correspond to the subject identifier of the resource owner.
|
|
|
|
if user:
|
|
token_data['sub'] = user.get_user_id()
|
|
|
|
# In cases of access tokens obtained through grants where no resource owner is
|
|
# involved, such as the client credentials grant, the value of 'sub' SHOULD
|
|
# correspond to an identifier the authorization server uses to indicate the
|
|
# client application.
|
|
|
|
else:
|
|
token_data['sub'] = client.get_client_id()
|
|
|
|
# If the request includes a 'resource' parameter (as defined in [RFC8707]), the
|
|
# resulting JWT access token 'aud' claim SHOULD have the same value as the
|
|
# 'resource' parameter in the request.
|
|
|
|
# TODO: Implement this with RFC8707
|
|
if False: # pragma: no cover
|
|
...
|
|
|
|
# If the request does not include a 'resource' parameter, the authorization
|
|
# server MUST use a default resource indicator in the 'aud' claim. If a 'scope'
|
|
# parameter is present in the request, the authorization server SHOULD use it to
|
|
# infer the value of the default resource indicator to be used in the 'aud'
|
|
# claim. The mechanism through which scopes are associated with default resource
|
|
# indicator values is outside the scope of this specification.
|
|
|
|
else:
|
|
token_data['aud'] = self.get_audiences(client, user, scope)
|
|
|
|
# If the values in the 'scope' parameter refer to different default resource
|
|
# indicator values, the authorization server SHOULD reject the request with
|
|
# 'invalid_scope' as described in Section 4.1.2.1 of [RFC6749].
|
|
# TODO: Implement this with RFC8707
|
|
|
|
if auth_time := self.get_auth_time(user):
|
|
token_data['auth_time'] = auth_time
|
|
|
|
# The meaning and processing of acr Claim Values is out of scope for this
|
|
# specification.
|
|
|
|
if acr := self.get_acr(user):
|
|
token_data['acr'] = acr
|
|
|
|
# The definition of particular values to be used in the amr Claim is beyond the
|
|
# scope of this specification.
|
|
|
|
if amr := self.get_amr(user):
|
|
token_data['amr'] = amr
|
|
|
|
# Authorization servers MAY return arbitrary attributes not defined in any
|
|
# existing specification, as long as the corresponding claim names are collision
|
|
# resistant or the access tokens are meant to be used only within a private
|
|
# subsystem. Please refer to Sections 4.2 and 4.3 of [RFC7519] for details.
|
|
|
|
token_data.update(self.get_extra_claims(client, grant_type, user, scope))
|
|
|
|
# This specification registers the 'application/at+jwt' media type, which can
|
|
# be used to indicate that the content is a JWT access token. JWT access tokens
|
|
# MUST include this media type in the 'typ' header parameter to explicitly
|
|
# declare that the JWT represents an access token complying with this profile.
|
|
# Per the definition of 'typ' in Section 4.1.9 of [RFC7515], it is RECOMMENDED
|
|
# that the 'application/' prefix be omitted. Therefore, the 'typ' value used
|
|
# SHOULD be 'at+jwt'.
|
|
|
|
header = {'alg': self.alg, 'typ': 'at+jwt', 'kid': 'k1'}
|
|
|
|
access_token = jwt.encode(
|
|
header,
|
|
token_data,
|
|
key=self.get_jwks(),
|
|
check=False,
|
|
)
|
|
return access_token.decode()
|
|
|
|
|
|
GRANT_TYPES_EXPIRES_IN = {
|
|
'refresh_token': 600,
|
|
}
|
|
|
|
|
|
def expires_in(client, grant_type: str) -> int:
|
|
return GRANT_TYPES_EXPIRES_IN.get(
|
|
grant_type, BearerTokenGenerator.DEFAULT_EXPIRES_IN
|
|
)
|
|
|
|
|
|
def create_token_generator(length: int = 42):
|
|
def token_generator(*args, **kwargs):
|
|
return generate_token(length)
|
|
|
|
return token_generator
|
|
|
|
|
|
server = AuthorizationServer(persistence_layer=dyn)
|
|
server.register_grant(
|
|
AuthorizationCodeGrant,
|
|
[
|
|
OpenIDCode(require_nonce=False),
|
|
CodeChallenge(required=True),
|
|
],
|
|
)
|
|
server.register_token_generator(
|
|
'default',
|
|
JWTBearerTokenGenerator(
|
|
issuer=ISSUER,
|
|
refresh_token_generator=create_token_generator(48),
|
|
expires_generator=expires_in,
|
|
),
|
|
)
|
|
server.register_grant(RefreshTokenGrant)
|
|
server.register_endpoint(RevocationEndpoint)
|
|
server.register_extension(IssuerParameter())
|