codekingpro/portable-devtools

1import functools2from django.http import JsonResponse3from authlib.oauth2 import (4    OAuth2Error,5    ResourceProtector as _ResourceProtector,6)7from authlib.oauth2.rfc6749 import (8    MissingAuthorizationError,9)10from authlib.oauth2.rfc6750 import (11    BearerTokenValidator as _BearerTokenValidator12)13from .requests import DjangoJsonRequest14from .signals import token_authenticated15 16 17class ResourceProtector(_ResourceProtector):18    def acquire_token(self, request, scopes=None, **kwargs):19        """A method to acquire current valid token with the given scope.20 21        :param request: Django HTTP request instance22        :param scopes: a list of scope values23        :return: token object24        """25        req = DjangoJsonRequest(request)26        # backward compatibility27        kwargs['scopes'] = scopes28        for claim in kwargs:29            if isinstance(kwargs[claim], str):30                kwargs[claim] = [kwargs[claim]]31        token = self.validate_request(request=req, **kwargs)32        token_authenticated.send(sender=self.__class__, token=token)33        return token34 35    def __call__(self, scopes=None, optional=False, **kwargs):36        claims = kwargs37        # backward compatibility38        claims['scopes'] = scopes39        def wrapper(f):40            @functools.wraps(f)41            def decorated(request, *args, **kwargs):42                try:43                    token = self.acquire_token(request, **claims)44                    request.oauth_token = token45                except MissingAuthorizationError as error:46                    if optional:47                        request.oauth_token = None48                        return f(request, *args, **kwargs)49                    return return_error_response(error)50                except OAuth2Error as error:51                    return return_error_response(error)52                return f(request, *args, **kwargs)53            return decorated54        return wrapper55 56 57class BearerTokenValidator(_BearerTokenValidator):58    def __init__(self, token_model, realm=None, **extra_attributes):59        self.token_model = token_model60        super().__init__(realm, **extra_attributes)61 62    def authenticate_token(self, token_string):63        try:64            return self.token_model.objects.get(access_token=token_string)65        except self.token_model.DoesNotExist:66            return None67 68 69def return_error_response(error):70    body = dict(error.get_body())71    resp = JsonResponse(body, status=error.status_code)72    headers = error.get_headers()73    for k, v in headers:74        resp[k] = v75    return resp76