Create auth in websocket

recircle/status/consumers.py

@@ -2,6 +2,7 @@
 from asgiref.sync import async_to_sync
 from asgiref.sync import async_to_sync
 from channels.generic.websocket import WebsocketConsumer
 from channels.generic.websocket import WebsocketConsumer
 from instance.models import Instance
 from instance.models import Instance
+from rest_framework.authtoken.models import Token
 import logging
 import logging
 import json
 import json
 
 
@@ -9,17 +10,30 @@ logger = logging.getLogger(__name__)
 
 
 
 
 class StatusConsumer(WebsocketConsumer):
 class StatusConsumer(WebsocketConsumer):
+    # Accept connection and wait for login message
     def connect(self):
     def connect(self):
-        instances = Instance.objects.all()
+        self.accept()
+
+    # If login message arrives get the auth token from it
+    # Then subscribe to the channels of the vm's that belongs to the user
+    def login(self, event):
+        user_token = event["auth_token"]
+        token = Token.objects.get(key=user_token)
+        if token:
+            self.user = token.user
+            instances = Instance.objects.filter(created_by=self.user)
             for inst in instances:
             for inst in instances:
                 async_to_sync(self.channel_layer.group_add)(
                 async_to_sync(self.channel_layer.group_add)(
                     str(inst.id),
                     str(inst.id),
                     self.channel_name
                     self.channel_name
                 )
                 )
-        self.accept()
+        else:
+            self.disconnect(100)
 
 
+    # Unsubscribe from registered groups
     def disconnect(self, close_code):
     def disconnect(self, close_code):
-        instances = Instance.objects.all()
+        if self.user:
+            instances = Instance.objects.filter(created_by=self.user)
             for inst in instances:
             for inst in instances:
                 async_to_sync(self.channel_layer.group_discard)(
                 async_to_sync(self.channel_layer.group_discard)(
                     str(inst.id),
                     str(inst.id),