diff --git a/awx/main/consumers.py b/awx/main/consumers.py index 223ff84321..bbd155329a 100644 --- a/awx/main/consumers.py +++ b/awx/main/consumers.py @@ -1,20 +1,64 @@ import json +import urlparse from channels import Group from channels.sessions import channel_session +from django.contrib.auth.models import User +from awx.main.models.organization import AuthToken + def discard_groups(message): if 'groups' in message.channel_session: for group in message.channel_session['groups']: Group(group).discard(message.reply_channel) + +def validate_token(token): + try: + auth_token = AuthToken.objects.get(key=token) + if not auth_token.in_valid_tokens: + return None + except AuthToken.DoesNotExist: + return None + return auth_token + + +def user_from_token(auth_token): + try: + return User.objects.get(pk=auth_token.user_id) + except User.DoesNotExist: + return None + +@channel_session +def ws_connect(message): + token = None + qs = urlparse.parse_qs(message['query_string']) + if 'token' in qs: + if len(qs['token']) > 0: + token = qs['token'].pop() + message.channel_session['token'] = token + + @channel_session def ws_disconnect(message): discard_groups(message) + @channel_session def ws_receive(message): + token = message.channel_session.get('token') + + auth_token = validate_token(token) + if auth_token is None: + message.reply_channel.send({"text": json.dumps({"error": "invalid auth token"})}) + return None + + user = user_from_token(auth_token) + if user is None: + message.reply_channel.send({"text": json.dumps({"error": "no valid user"})}) + return None + raw_data = message.content['text'] data = json.loads(raw_data) diff --git a/awx/main/routing.py b/awx/main/routing.py index 67a08ff1bd..0a49f25c6c 100644 --- a/awx/main/routing.py +++ b/awx/main/routing.py @@ -2,6 +2,7 @@ from channels.routing import route channel_routing = [ + route("websocket.connect", "awx.main.consumers.ws_connect", path=r'^/websocket/$'), route("websocket.disconnect", "awx.main.consumers.ws_disconnect", path=r'^/websocket/$'), route("websocket.receive", "awx.main.consumers.ws_receive", path=r'^/websocket/$'), ]