关于django-channels的token认证实现

向右看齐 2022-07-11 08:46 141阅读 0赞

### django-channels只支持基于session的认证，在实际使用中可能需要基于token或者jwt的认证。 ###

### 参照 [https://www.machinalis.com/blog/jwt-django-channels/][https_www.machinalis.com_blog_jwt-django-channels] 实现基于django token认证。 ###

# ws_authentication.py
    from functools import wraps
    from json import loads, dumps
    
    from django.utils.translation import ugettext_lazy as _
    from django.core.exceptions import ObjectDoesNotExist
    
    from rest_framework import exceptions
    from channels.handler import AsgiRequest
    
    from multi_token.models import Token
    
    
    def _close_reply_channel(message):
        message.reply_channel.send({
       "close": True})
    
    
    def authenticate(token):
        """ Tries to authenticate user based on the supplied token. It also checks the token structure and validity. """
        try:
            auth_token = Token.objects.get(key=token)
        except Token.DoesNotExist:
            msg = _('Invalid token')
            raise exceptions.AuthenticationFailed(msg)
        """Other authenticate operation"""
    
        return auth_token.user, auth_token
    
    
    def ws_auth_request_token(func):
        """ Checks the presence of a "token" request parameter and tries to authenticate the user based on its content. The request url must include token. eg: /v1/channel/1/?token=abcdefghijklmn """
        @wraps(func)
        def inner(message, *args, **kwargs):
            try:
                if "method" not in message.content:
                    message.content['method'] = "FAKE"
                request = AsgiRequest(message)
            except Exception as e:
                raise ValueError("Cannot parse HTTP message - are you sure this is a HTTP consumer? %s" % e)
    
            token = request.GET.get("token", None)
    
            print request.path,request.GET
    
            if token is None:
                _close_reply_channel(message)
                raise ValueError("Missing token request parameter. Closing channel.")
    
            user, token = authenticate(token)
    
            message.token = token
            message.user = user
    
            return func(message, *args, **kwargs)
        return inner
    
    
    def ws_auth_message_token(func):
        """ Checks the presence of a "token" field on the message's text field and tries to authenticate the user based on its content. The text must include "token" field. eg:{'text':{'token': 'abcdefg','otherdata':''}}. """
        @wraps(func)
        def inner(message, *args, **kwargs):
            message_text = message.get('text', None)
            if message_text is None:
                _close_reply_channel(message)
                raise ValueError("Missing text field. Closing channel.")
    
            try:
                message_text_json = loads(message_text)
            except ValueError:
                _close_reply_channel(message)
                raise
    
            token = message_text_json.pop('token', None)
            if token is None:
                _close_reply_channel(message)
                raise ValueError("Missing token field. Closing channel.")
    
            user, token = authenticate(token)
    
            message.token = token
            message.user = user
            message.text = dumps(message_text_json)
    
            return func(message, *args, **kwargs)
        return inner

### 使用： ###

from channels import Group
    from .ws_authentication import ws_auth_request_token, ws_auth_message_token
    
    
    @ws_auth_request_token
    def ws_connect(message, room_id):
        Group("room-%s" % room_id).add(message.reply_channel)
        message.reply_channel.send({
       "accept": True})
    
    
    @ws_auth_message_token
    def ws_receive(message, room_id):
        message.reply_channel.send({
       'text': message['text']})
    
    
    def ws_disconnect(message,room_id):
        Group("room-%s" % room_id).discard(message.reply_channel)

### ws测试： ###

from channels.tests import ChannelTestCase, HttpClient
    
    from rest_framework import status
    from rest_framework.test import APITestCase
    
    
    class WsNotifyTests(ChannelTestCase, APITestCase):
    
        @override_settings(DEBUG=True)
        def get_token(self):
            """Get token"""
            return token
    
        def test_notify(self):
    
            token = self.get_token()
    
            token_str = "token=%s" % token
            client = HttpClient()
            client.send_and_consume(u'websocket.connect',
                                    {
       'path': '/v1/channel/, 'method': 'GET', "query_string": token_str}) self.assertIsNone(client.receive()) Group('restaurant-%s' % restaurant_id).send({'text': 'ok'}, immediately=True) self.assertEqual(client.receive(json=False), 'ok')

[https_www.machinalis.com_blog_jwt-django-channels]: https://www.machinalis.com/blog/jwt-django-channels/