codekingpro/portable-devtools

1from django.http import HttpResponse2from django.utils.module_loading import import_string3from django.conf import settings4from authlib.oauth2 import (5    AuthorizationServer as _AuthorizationServer,6)7from authlib.oauth2.rfc6750 import BearerTokenGenerator8from authlib.common.security import generate_token as _generate_token9from authlib.common.encoding import json_dumps10from .requests import DjangoOAuth2Request, DjangoJsonRequest11from .signals import client_authenticated, token_revoked12 13 14class AuthorizationServer(_AuthorizationServer):15    """Django implementation of :class:`authlib.oauth2.rfc6749.AuthorizationServer`.16    Initialize it with client model and token model::17 18        from authlib.integrations.django_oauth2 import AuthorizationServer19        from your_project.models import OAuth2Client, OAuth2Token20 21        server = AuthorizationServer(OAuth2Client, OAuth2Token)22    """23 24    def __init__(self, client_model, token_model):25        self.config = getattr(settings, 'AUTHLIB_OAUTH2_PROVIDER', {})26        self.client_model = client_model27        self.token_model = token_model28        scopes_supported = self.config.get('scopes_supported')29        super().__init__(scopes_supported=scopes_supported)30        # add default token generator31        self.register_token_generator('default', self.create_bearer_token_generator())32 33    def query_client(self, client_id):34        """Default method for ``AuthorizationServer.query_client``. Developers MAY35        rewrite this function to meet their own needs.36        """37        try:38            return self.client_model.objects.get(client_id=client_id)39        except self.client_model.DoesNotExist:40            return None41 42    def save_token(self, token, request):43        """Default method for ``AuthorizationServer.save_token``. Developers MAY44        rewrite this function to meet their own needs.45        """46        client = request.client47        if request.user:48            user_id = request.user.pk49        else:50            user_id = client.user_id51        item = self.token_model(52            client_id=client.client_id,53            user_id=user_id,54            **token55        )56        item.save()57        return item58 59    def create_oauth2_request(self, request):60        return DjangoOAuth2Request(request)61 62    def create_json_request(self, request):63        return DjangoJsonRequest(request)64 65    def handle_response(self, status_code, payload, headers):66        if isinstance(payload, dict):67            payload = json_dumps(payload)68        resp = HttpResponse(payload, status=status_code)69        for k, v in headers:70            resp[k] = v71        return resp72 73    def send_signal(self, name, *args, **kwargs):74        if name == 'after_authenticate_client':75            client_authenticated.send(sender=self.__class__, *args, **kwargs)76        elif name == 'after_revoke_token':77            token_revoked.send(sender=self.__class__, *args, **kwargs)78 79    def create_bearer_token_generator(self):80        """Default method to create BearerToken generator."""81        conf = self.config.get('access_token_generator', True)82        access_token_generator = create_token_generator(conf, 42)83 84        conf = self.config.get('refresh_token_generator', False)85        refresh_token_generator = create_token_generator(conf, 48)86 87        conf = self.config.get('token_expires_in')88        expires_generator = create_token_expires_in_generator(conf)89 90        return BearerTokenGenerator(91            access_token_generator=access_token_generator,92            refresh_token_generator=refresh_token_generator,93            expires_generator=expires_generator,94        )95 96 97def create_token_generator(token_generator_conf, length=42):98    if callable(token_generator_conf):99        return token_generator_conf100 101    if isinstance(token_generator_conf, str):102        return import_string(token_generator_conf)103    elif token_generator_conf is True:104        def token_generator(*args, **kwargs):105            return _generate_token(length)106        return token_generator107 108 109def create_token_expires_in_generator(expires_in_conf=None):110    data = {}111    data.update(BearerTokenGenerator.GRANT_TYPES_EXPIRES_IN)112    if expires_in_conf:113        data.update(expires_in_conf)114 115    def expires_in(client, grant_type):116        return data.get(grant_type, BearerTokenGenerator.DEFAULT_EXPIRES_IN)117 118    return expires_in119