Create auth in websocket

recircle/status/consumers.py

@@ -2,6 +2,7 @@
 from asgiref.sync import async_to_sync
 from channels.generic.websocket import WebsocketConsumer
 from instance.models import Instance
+from rest_framework.authtoken.models import Token
 import logging
 import json
 
@@ -9,22 +10,35 @@ logger = logging.getLogger(__name__)
 
 
 class StatusConsumer(WebsocketConsumer):
+    # Accept connection and wait for login message
     def connect(self):
-        instances = Instance.objects.all()
-        for inst in instances:
-            async_to_sync(self.channel_layer.group_add)(
-                str(inst.id),
-                self.channel_name
-            )
         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:
+                async_to_sync(self.channel_layer.group_add)(
+                    str(inst.id),
+                    self.channel_name
+                )
+        else:
+            self.disconnect(100)
+
+    # Unsubscribe from registered groups
     def disconnect(self, close_code):
-        instances = Instance.objects.all()
-        for inst in instances:
-            async_to_sync(self.channel_layer.group_discard)(
-                str(inst.id),
-                self.channel_name
-            )
+        if self.user:
+            instances = Instance.objects.filter(created_by=self.user)
+            for inst in instances:
+                async_to_sync(self.channel_layer.group_discard)(
+                    str(inst.id),
+                    self.channel_name
+                )
 
     def receive(self, text_message):
         pass